Merge pull request #30076 from aschwaighofer/silgen_fix_dynamic_replacement_before_original

SILGen: Fix dynamic replacement before original function

Author: aschwaighofer

include/swift/SIL/SILFunctionBuilder.h

@@ -80,10 +80,16 @@ class SILFunctionBuilder {
       SubclassScope subclassScope = SubclassScope::NotApplicable);
 
   /// Return the declaration of a function, or create it if it doesn't exist.
-  SILFunction *
-  getOrCreateFunction(SILLocation loc, SILDeclRef constant,
-                      ForDefinition_t forDefinition,
-                      ProfileCounter entryCount = ProfileCounter());
+  /// Creating a function that references another function (a dynamic
+  /// replacement) requires getting/creating a SIL function for this reference.
+  /// In this case the function will call the \p getOrCreateDeclaration
+  /// callback.
+  SILFunction *getOrCreateFunction(
+      SILLocation loc, SILDeclRef constant, ForDefinition_t forDefinition,
+      llvm::function_ref<SILFunction *(SILLocation loc, SILDeclRef constant)>
+          getOrCreateDeclaration = [](SILLocation loc, SILDeclRef constant)
+          -> SILFunction * { return nullptr; },
+      ProfileCounter entryCount = ProfileCounter());
 
   /// Create a function declaration.
   ///
@@ -104,8 +110,11 @@ class SILFunctionBuilder {
                  SILFunction *InsertBefore = nullptr,
                  const SILDebugScope *DebugScope = nullptr);
 
-  void addFunctionAttributes(SILFunction *F, DeclAttributes &Attrs,
-                             SILModule &M, SILDeclRef constant = SILDeclRef());
+  void addFunctionAttributes(
+      SILFunction *F, DeclAttributes &Attrs, SILModule &M,
+      llvm::function_ref<SILFunction *(SILLocation loc, SILDeclRef constant)>
+          getOrCreateDeclaration,
+      SILDeclRef constant = SILDeclRef());
 
   /// We do not expose this to everyone, instead we allow for our users to opt
   /// into this if they need to. Please do not do this in general! We only want

lib/SIL/SILFunctionBuilder.cpp

@@ -37,10 +37,11 @@ SILFunction *SILFunctionBuilder::getOrCreateFunction(
   return fn;
 }
 
-void SILFunctionBuilder::addFunctionAttributes(SILFunction *F,
-                                               DeclAttributes &Attrs,
-                                               SILModule &M,
-                                               SILDeclRef constant) {
+void SILFunctionBuilder::addFunctionAttributes(
+    SILFunction *F, DeclAttributes &Attrs, SILModule &M,
+    llvm::function_ref<SILFunction *(SILLocation loc, SILDeclRef constant)>
+        getOrCreateDeclaration,
+    SILDeclRef constant) {
 
   for (auto *A : Attrs.getAttributes<SemanticsAttr>())
     F->addSemanticsAttr(cast<SemanticsAttr>(A)->Value);
@@ -92,8 +93,7 @@ void SILFunctionBuilder::addFunctionAttributes(SILFunction *F,
     return;
 
   SILDeclRef declRef(replacedDecl, constant.kind, false);
-  auto *replacedFunc =
-      getOrCreateFunction(replacedDecl, declRef, NotForDefinition);
+  auto *replacedFunc = getOrCreateDeclaration(replacedDecl, declRef);
 
   assert(replacedFunc->getLoweredFunctionType() ==
              F->getLoweredFunctionType() ||
@@ -102,10 +102,11 @@ void SILFunctionBuilder::addFunctionAttributes(SILFunction *F,
   F->setDynamicallyReplacedFunction(replacedFunc);
 }
 
-SILFunction *
-SILFunctionBuilder::getOrCreateFunction(SILLocation loc, SILDeclRef constant,
-                                        ForDefinition_t forDefinition,
-                                        ProfileCounter entryCount) {
+SILFunction *SILFunctionBuilder::getOrCreateFunction(
+    SILLocation loc, SILDeclRef constant, ForDefinition_t forDefinition,
+    llvm::function_ref<SILFunction *(SILLocation loc, SILDeclRef constant)>
+        getOrCreateDeclaration,
+    ProfileCounter entryCount) {
   auto nameTmp = constant.mangle();
   auto constantType = mod.Types.getConstantFunctionType(
       TypeExpansionContext::minimal(), constant);
@@ -169,9 +170,11 @@ SILFunctionBuilder::getOrCreateFunction(SILLocation loc, SILDeclRef constant,
     if (auto *accessor = dyn_cast<AccessorDecl>(decl)) {
       auto *storage = accessor->getStorage();
       // Add attributes for e.g. computed properties.
-      addFunctionAttributes(F, storage->getAttrs(), mod);
+      addFunctionAttributes(F, storage->getAttrs(), mod,
+                            getOrCreateDeclaration);
     }
-    addFunctionAttributes(F, decl->getAttrs(), mod, constant);
+    addFunctionAttributes(F, decl->getAttrs(), mod, getOrCreateDeclaration,
+                          constant);
   }
 
   return F;

lib/SILGen/SILGen.cpp

@@ -612,9 +612,13 @@ SILFunction *SILGenModule::getFunction(SILDeclRef constant,
 
   // Note: Do not provide any SILLocation. You can set it afterwards.
   SILGenFunctionBuilder builder(*this);
-  auto *F = builder.getOrCreateFunction(constant.hasDecl() ? constant.getDecl()
-                                                           : (Decl *)nullptr,
-                                        constant, forDefinition);
+  auto &IGM = *this;
+  auto *F = builder.getOrCreateFunction(
+      constant.hasDecl() ? constant.getDecl() : (Decl *)nullptr, constant,
+      forDefinition,
+      [&IGM](SILLocation loc, SILDeclRef constant) -> SILFunction * {
+        return IGM.getFunction(constant, NotForDefinition);
+      });
   setUpForProfiling(constant, F, forDefinition);
 
   assert(F && "SILFunction should have been defined");

test/SILGen/dynamically_replaceable.swift

@@ -418,3 +418,13 @@ struct WrapperWithInitialValue<T> {
 public struct SomeStruct {
   @WrapperWithInitialValue var t = false
 }
+
+// Make sure that declaring the replacement before the original does not assert.
+@_dynamicReplacement(for: orig)
+func opaqueReplacement() -> Int {
+  return 2
+}
+
+dynamic func orig() -> Int {
+  return 1
+}