IRGen/runtime: change the code generation for dynamically replaceable functions

Instead of a thunk insert the dispatch into the original function.
If the original function should be executed the prolog just jumps to the "real" code in the function. Otherwise the replacement function is called.
There is one little complication here: when the replacement function calls the original function, the original function should not dispatch to the replacement again.
To pass this information, we use a flag in thread local storage.
The setting and reading of the flag is done in two new runtime functions.

rdar://problem/51043781

Author: eeckstein

include/swift/Runtime/Exclusivity.h

@@ -40,6 +40,20 @@ SWIFT_RUNTIME_EXPORT
 void swift_beginAccess(void *pointer, ValueBuffer *buffer,
                        ExclusivityFlags flags, void *pc);
 
+/// Loads the replacement function pointer from \p ReplFnPtr and returns the
+/// replacement function if it should be called.
+/// Returns null if the original function (which is passed in \p CurrFn) should
+/// be called.
+SWIFT_RUNTIME_EXPORT
+char *swift_getFunctionReplacement(char **ReplFnPtr, char *CurrFn);
+
+/// Returns the original function of a replaced function, which is loaded from
+/// \p OrigFnPtr.
+/// This function is called from a replacement function to call the original
+/// function.
+SWIFT_RUNTIME_EXPORT
+char *swift_getOrigOfReplaceable(char **OrigFnPtr);
+
 /// Stop dynamically tracking an access.
 SWIFT_RUNTIME_EXPORT
 void swift_endAccess(ValueBuffer *buffer);

include/swift/Runtime/RuntimeFunctions.def

@@ -1225,6 +1225,16 @@ FUNCTION(EndAccess, swift_endAccess, C_CC, AlwaysAvailable,
          ARGS(getFixedBufferTy()->getPointerTo()),
          ATTRS(NoUnwind))
 
+FUNCTION(GetOrigOfReplaceable, swift_getOrigOfReplaceable, C_CC, AlwaysAvailable,
+         RETURNS(FunctionPtrTy),
+         ARGS(FunctionPtrTy->getPointerTo()),
+         ATTRS(NoUnwind))
+
+FUNCTION(GetReplacement, swift_getFunctionReplacement, C_CC, AlwaysAvailable,
+         RETURNS(FunctionPtrTy),
+         ARGS(FunctionPtrTy->getPointerTo(), FunctionPtrTy),
+         ATTRS(NoUnwind))
+
 FUNCTION(InstantiateObjCClass, swift_instantiateObjCClass,
          C_CC, AlwaysAvailable,
          RETURNS(VoidTy),

lib/IRGen/GenDecl.cpp

@@ -2175,6 +2175,71 @@ static llvm::GlobalVariable *createGlobalForDynamicReplacementFunctionKey(
   return key;
 }
 
+/// Creates the prolog for a dynamically replaceable function.
+/// It checks if the replaced function or the original function should be called
+/// (by calling the swift_getFunctionReplacement runtime function). In case of
+/// the original function, it just jumps to the "real" code of the function,
+/// otherwise it tail calls the replacement.
+void IRGenModule::createReplaceableProlog(IRGenFunction &IGF, SILFunction *f) {
+  LinkEntity varEntity =
+      LinkEntity::forDynamicallyReplaceableFunctionVariable(f);
+  LinkEntity keyEntity =
+      LinkEntity::forDynamicallyReplaceableFunctionKey(f);
+  Signature signature = getSignature(f->getLoweredFunctionType());
+
+  // Create and initialize the first link entry for the chain of replacements.
+  // The first implementation is initialized with 'implFn'.
+  auto linkEntry = getChainEntryForDynamicReplacement(*this, varEntity, IGF.CurFn);
+
+  // Create the key data structure. This is used from other modules to refer to
+  // the chain of replacements.
+  createGlobalForDynamicReplacementFunctionKey(*this, keyEntity, linkEntry);
+
+  llvm::Constant *indices[] = {llvm::ConstantInt::get(Int32Ty, 0),
+                               llvm::ConstantInt::get(Int32Ty, 0)};
+
+  auto *fnPtrAddr =
+      llvm::ConstantExpr::getInBoundsGetElementPtr(nullptr, linkEntry, indices);
+
+  auto *ReplAddr =
+    llvm::ConstantExpr::getPointerBitCastOrAddrSpaceCast(fnPtrAddr,
+      FunctionPtrTy->getPointerTo());
+
+  auto *FnAddr = llvm::ConstantExpr::getPointerBitCastOrAddrSpaceCast(
+      IGF.CurFn, FunctionPtrTy);
+
+  // Call swift_getFunctionReplacement to check which function to call.
+  auto *ReplFn = IGF.Builder.CreateCall(getGetReplacementFn(),
+                                        { ReplAddr, FnAddr });
+  ReplFn->setDoesNotThrow();
+  auto *hasReplFn = IGF.Builder.CreateICmpEQ(ReplFn,
+          llvm::ConstantExpr::getNullValue(ReplFn->getType()));
+
+  auto *replacedBB = IGF.createBasicBlock("forward_to_replaced");
+  auto *origEntryBB = IGF.createBasicBlock("original_entry");
+  IGF.Builder.CreateCondBr(hasReplFn, origEntryBB, replacedBB);
+
+  IGF.Builder.emitBlock(replacedBB);
+
+  // Call the replacement function.
+  SmallVector<llvm::Value *, 16> forwardedArgs;
+  for (auto &arg : IGF.CurFn->args())
+    forwardedArgs.push_back(&arg);
+  auto *fnType = signature.getType()->getPointerTo();
+  auto *realReplFn = IGF.Builder.CreateBitCast(ReplFn, fnType);
+
+  auto *Res = IGF.Builder.CreateCall(FunctionPointer(realReplFn, signature),
+                                     forwardedArgs);
+  Res->setTailCall();
+  if (IGF.CurFn->getReturnType()->isVoidTy())
+    IGF.Builder.CreateRetVoid();
+  else
+    IGF.Builder.CreateRet(Res);
+
+  IGF.Builder.emitBlock(origEntryBB);
+
+}
+
 /// Emit the thunk that dispatches to the dynamically replaceable function.
 static void emitDynamicallyReplaceableThunk(IRGenModule &IGM,
                                             LinkEntity varEntity,
@@ -2303,11 +2368,17 @@ void IRGenModule::emitDynamicReplacementOriginalFunctionThunk(SILFunction *f) {
   llvm::Constant *indices[] = {llvm::ConstantInt::get(Int32Ty, 0),
                                llvm::ConstantInt::get(Int32Ty, 0)};
 
-  auto *fnPtr = IGF.Builder.CreateLoad(
+  auto *fnPtrAddr =
+    llvm::ConstantExpr::getPointerBitCastOrAddrSpaceCast(
       llvm::ConstantExpr::getInBoundsGetElementPtr(nullptr, linkEntry, indices),
-      getPointerAlignment());
+      FunctionPtrTy->getPointerTo());
+
+  auto *OrigFn = IGF.Builder.CreateCall(getGetOrigOfReplaceableFn(),
+                                        { fnPtrAddr });
+  OrigFn->setDoesNotThrow();
+
   auto *typeFnPtr =
-      IGF.Builder.CreateBitOrPointerCast(fnPtr, implFn->getType());
+      IGF.Builder.CreateBitOrPointerCast(OrigFn, implFn->getType());
 
   SmallVector<llvm::Value *, 16> forwardedArgs;
   for (auto &arg : implFn->args())
@@ -2338,25 +2409,6 @@ llvm::Function *IRGenModule::getAddrOfSILFunction(
   if (fn) {
     if (forDefinition) {
       updateLinkageForDefinition(*this, fn, entity);
-      if (isDynamicallyReplaceableImplementation) {
-        // Create the dynamically replacement thunk.
-        LinkEntity implEntity = LinkEntity::forSILFunction(f, true);
-        auto existingImpl = Module.getFunction(implEntity.mangleAsString());
-        assert(!existingImpl);
-        (void) existingImpl;
-        Signature signature = getSignature(f->getLoweredFunctionType());
-        addLLVMFunctionAttributes(f, signature);
-        LinkInfo implLink = LinkInfo::get(*this, implEntity, forDefinition);
-        auto implFn = createFunction(*this, implLink, signature, fn,
-                                     f->getOptimizationMode());
-        LinkEntity varEntity =
-            LinkEntity::forDynamicallyReplaceableFunctionVariable(f);
-        LinkEntity keyEntity =
-            LinkEntity::forDynamicallyReplaceableFunctionKey(f);
-        emitDynamicallyReplaceableThunk(*this, varEntity, keyEntity, fn, implFn,
-                                        signature);
-        return implFn;
-      }
     }
     return fn;
   }
@@ -2431,20 +2483,6 @@ llvm::Function *IRGenModule::getAddrOfSILFunction(
   if (hasOrderNumber) {
     EmittedFunctionsByOrder.insert(orderNumber, fn);
   }
-
-  if (isDynamicallyReplaceableImplementation && forDefinition) {
-    LinkEntity implEntity = LinkEntity::forSILFunction(f, true);
-    LinkInfo implLink = LinkInfo::get(*this, implEntity, forDefinition);
-    auto implFn = createFunction(*this, implLink, signature, insertBefore,
-                                 f->getOptimizationMode());
-
-    LinkEntity varEntity =
-        LinkEntity::forDynamicallyReplaceableFunctionVariable(f);
-    LinkEntity keyEntity = LinkEntity::forDynamicallyReplaceableFunctionKey(f);
-    emitDynamicallyReplaceableThunk(*this, varEntity, keyEntity, fn, implFn,
-                                    signature);
-    return implFn;
-  }
   return fn;
 }
 

lib/IRGen/IRGenModule.h

@@ -1437,6 +1437,8 @@ private:                            \
   llvm::Function *getAddrOfOpaqueTypeDescriptorAccessFunction(
       OpaqueTypeDecl *decl, ForDefinition_t forDefinition, bool implementation);
 
+  void createReplaceableProlog(IRGenFunction &IGF, SILFunction *f);
+
   void emitOpaqueTypeDescriptorAccessor(OpaqueTypeDecl *);
 
 private:

lib/IRGen/IRGenSIL.cpp

@@ -1244,6 +1244,10 @@ IRGenSILFunction::IRGenSILFunction(IRGenModule &IGM, SILFunction *f)
   if (f->getDynamicallyReplacedFunction()) {
     IGM.emitDynamicReplacementOriginalFunctionThunk(f);
   }
+
+  if (f->isDynamicallyReplaceable()) {
+    IGM.createReplaceableProlog(*this, f);
+  }
 }
 
 IRGenSILFunction::~IRGenSILFunction() {
@@ -1646,7 +1650,7 @@ void IRGenSILFunction::emitSILFunction() {
     IGM.DebugInfo->emitFunction(*CurSILFn, CurFn);
 
   // Map the entry bb.
-  LoweredBBs[&*CurSILFn->begin()] = LoweredBB(&*CurFn->begin(), {});
+  LoweredBBs[&*CurSILFn->begin()] = LoweredBB(&CurFn->back(), {});
   // Create LLVM basic blocks for the other bbs.
   for (auto bi = std::next(CurSILFn->begin()), be = CurSILFn->end(); bi != be;
        ++bi) {

lib/TBDGen/TBDGen.cpp

@@ -196,8 +196,6 @@ void TBDGenVisitor::visitAbstractFunctionDecl(AbstractFunctionDecl *AFD) {
     bool useAllocator = shouldUseAllocatorMangling(AFD);
     addSymbol(LinkEntity::forDynamicallyReplaceableFunctionVariable(
         AFD, useAllocator));
-    addSymbol(
-        LinkEntity::forDynamicallyReplaceableFunctionImpl(AFD, useAllocator));
     addSymbol(
         LinkEntity::forDynamicallyReplaceableFunctionKey(AFD, useAllocator));
 

stdlib/public/runtime/Exclusivity.cpp

@@ -157,8 +157,16 @@ static_assert(sizeof(Access) <= sizeof(ValueBuffer) &&
               "Access doesn't fit in a value buffer!");
 
 /// A set of accesses that we're tracking.  Just a singly-linked list.
+/// TODO: rename this class to something like SwiftTLSContext, because it also
+/// contains fields not related to access tracking.
 class AccessSet {
   Access *Head = nullptr;
+
+  // Not related to access tracking: The "implicit" boolean parameter which is
+  // passed to a dynamically replaceable function.
+  // If true, the original function should be executed instead of the
+  // replacement function.
+  bool CallOriginalOfReplacedFunction = false;
 public:
   constexpr AccessSet() {}
 
@@ -221,6 +229,21 @@ class AccessSet {
     }
   }
 #endif
+
+  /// Called immediately before a replacement function calls its original
+  /// function.
+  void aboutToCallOriginalOfReplacedFunction() {
+    CallOriginalOfReplacedFunction = true;
+  }
+
+  /// Checked in the prolog of a replaceable function. Returns true if the
+  /// original function should be called instead of the replacement function.
+  /// Also clears the CallOriginalOfReplacedFunction flag.
+  bool shouldCallOriginalOfReplacedFunction() {
+    bool callOrig = CallOriginalOfReplacedFunction;
+    CallOriginalOfReplacedFunction = false;
+    return callOrig;
+  }
 };
 
 } // end anonymous namespace
@@ -327,6 +350,22 @@ void swift::swift_endAccess(ValueBuffer *buffer) {
   getAccessSet().remove(access);
 }
 
+char *swift::swift_getFunctionReplacement(char **ReplFnPtr, char *CurrFn) {
+  char *ReplFn = *ReplFnPtr;
+  if (ReplFn == CurrFn)
+    return nullptr;
+  if (getAccessSet().shouldCallOriginalOfReplacedFunction()) {
+    return nullptr;
+  }
+  return ReplFn;
+}
+
+char *swift::swift_getOrigOfReplaceable(char **OrigFnPtr) {
+  char *OrigFn = *OrigFnPtr;
+  getAccessSet().aboutToCallOriginalOfReplacedFunction();
+  return OrigFn;
+}
+
 #ifndef NDEBUG
 
 // Dump the accesses that are currently being tracked by the runtime.

test/IRGen/dynamic_replaceable.sil

@@ -2,7 +2,7 @@
 
 // REQUIRES: objc_interop
 
-// CHECK: @test_dynamically_replaceableTX = global %swift.dyn_repl_link_entry { i8* bitcast (void ()* @test_dynamically_replaceableTI to i8*), %swift.dyn_repl_link_entry* null }
+// CHECK: @test_dynamically_replaceableTX = global %swift.dyn_repl_link_entry { i8* bitcast (void ()* @test_dynamically_replaceable to i8*), %swift.dyn_repl_link_entry* null }
 // CHECK: @test_dynamically_replaceableTx = constant %swift.dyn_repl_key { i32{{.*}}%swift.dyn_repl_link_entry* @test_dynamically_replaceableTX{{.*}}, i32 0 }, section "__TEXT,__const"
 // CHECK: @test_replacementTX = global %swift.dyn_repl_link_entry zeroinitializer
 
@@ -23,17 +23,17 @@
 
 // CHECK-LABEL: define swiftcc void @test_dynamically_replaceable()
 // CHECK-NEXT: entry:
-// CHECK-NEXT:   [[FUN_PTR:%.*]] = load i8*, i8** getelementptr inbounds (%swift.dyn_repl_link_entry, %swift.dyn_repl_link_entry* @test_dynamically_replaceableTX, i32 0, i32 0)
+// CHECK-NEXT:   [[FUN_PTR:%.*]] =  call i8* @swift_getFunctionReplacement(i8** getelementptr inbounds (%swift.dyn_repl_link_entry, %swift.dyn_repl_link_entry* @test_dynamically_replaceableTX, i32 0, i32 0), i8* bitcast (void ()* @test_dynamically_replaceable to i8*))
+// CHECK-NEXT:   [[CMP:%.*]] = icmp eq i8* [[FUN_PTR]], null
+// CHECK-NEXT:   br i1 [[CMP]], label %[[ORIGBB:[a-z_]*]], label %[[FWBB:[a-z_]*]]
+// CHECK:      [[FWBB]]:
 // CHECK-NEXT:   [[TYPED_PTR:%.*]] = bitcast i8* [[FUN_PTR]] to void ()*
 // CHECK-NEXT:   tail call swiftcc void [[TYPED_PTR]]()
 // CHECK-NEXT:   ret void
+// CHECK:      [[ORIGBB]]:
+// CHECK-NEXT:   ret void
 // CHECK-NEXT: }
 
-// CHECK-LABEL: define swiftcc void @test_dynamically_replaceableTI()
-// CHECK: entry:
-// CHECK:   ret void
-// CHECK: }
-
 sil [dynamically_replacable] @test_dynamically_replaceable : $@convention(thin) () -> () {
 bb0:
   %0 = tuple ()
@@ -49,7 +49,7 @@ bb0:
 // The thunk that implement the prev_dynamic_function_ref lookup.
 // CHECK-LABEL: define swiftcc void @test_replacementTI()
 // CHECK: entry:
-// CHECK:   [[FUN_PTR:%.*]] = load i8*, i8** getelementptr inbounds (%swift.dyn_repl_link_entry, %swift.dyn_repl_link_entry* @test_replacementTX, i32 0, i32 0)
+// CHECK:   [[FUN_PTR:%.*]] = call i8* @swift_getOrigOfReplaceable(i8** getelementptr inbounds (%swift.dyn_repl_link_entry, %swift.dyn_repl_link_entry* @test_replacementTX, i32 0, i32 0))
 // CHECK:   [[TYPED_PTR:%.*]] = bitcast i8* [[FUN_PTR]] to void ()*
 // CHECK:   call swiftcc void [[TYPED_PTR]]()
 // CHECK:   ret void