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

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

@@ -75,6 +75,10 @@
 using namespace swift;
 using namespace irgen;
 
+llvm::cl::opt<bool> UseBasicDynamicReplacement(
+    "basic-dynamic-replacement", llvm::cl::init(false),
+    llvm::cl::desc("Basic implementation of dynamic replacement"));
+
 namespace {
 
 /// Add methods, properties, and protocol conformances from a JITed extension
@@ -2175,6 +2179,79 @@ 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);
+
+  llvm::Value *ReplFn = nullptr, *rhs = nullptr;
+
+  if (UseBasicDynamicReplacement) {
+    ReplFn = IGF.Builder.CreateLoad(fnPtrAddr, getPointerAlignment());
+    rhs = FnAddr;
+  } else {
+    // Call swift_getFunctionReplacement to check which function to call.
+    auto *callRTFunc = IGF.Builder.CreateCall(getGetReplacementFn(),
+                                             { ReplAddr, FnAddr });
+    callRTFunc->setDoesNotThrow();
+    ReplFn = callRTFunc;
+    rhs = llvm::ConstantExpr::getNullValue(ReplFn->getType());
+  }
+  auto *hasReplFn = IGF.Builder.CreateICmpEQ(ReplFn, rhs);
+
+  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,
@@ -2279,6 +2356,9 @@ void IRGenModule::emitOpaqueTypeDescriptorAccessor(OpaqueTypeDecl *opaque) {
 void IRGenModule::emitDynamicReplacementOriginalFunctionThunk(SILFunction *f) {
   assert(f->getDynamicallyReplacedFunction());
 
+  if (UseBasicDynamicReplacement)
+    return;
+
   auto entity = LinkEntity::forSILFunction(f, true);
 
   Signature signature = getSignature(f->getLoweredFunctionType());
@@ -2303,11 +2383,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 +2424,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 +2498,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/GenDecl.h

@@ -56,4 +56,6 @@ namespace irgen {
 }
 }
 
+extern llvm::cl::opt<bool> UseBasicDynamicReplacement;
+
 #endif

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

@@ -60,6 +60,7 @@
 #include "GenCast.h"
 #include "GenClass.h"
 #include "GenConstant.h"
+#include "GenDecl.h"
 #include "GenEnum.h"
 #include "GenExistential.h"
 #include "GenFunc.h"
@@ -1244,6 +1245,10 @@ IRGenSILFunction::IRGenSILFunction(IRGenModule &IGM, SILFunction *f)
   if (f->getDynamicallyReplacedFunction()) {
     IGM.emitDynamicReplacementOriginalFunctionThunk(f);
   }
+
+  if (f->isDynamicallyReplaceable()) {
+    IGM.createReplaceableProlog(*this, f);
+  }
 }
 
 IRGenSILFunction::~IRGenSILFunction() {
@@ -1646,7 +1651,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) {
@@ -1870,6 +1875,11 @@ void IRGenSILFunction::visitDynamicFunctionRefInst(DynamicFunctionRefInst *i) {
 
 void IRGenSILFunction::visitPreviousDynamicFunctionRefInst(
     PreviousDynamicFunctionRefInst *i) {
+  if (UseBasicDynamicReplacement) {
+    IGM.unimplemented(i->getLoc().getSourceLoc(),
+      ": calling the original implementation of a dynamic function is not "
+      "supported with -Xllvm -basic-dynamic-replacement");
+  }
   visitFunctionRefBaseInst(i);
 }
 

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));