IRGen: Define and use an earliest insertion point

This is to deal with the fact that LLVM's coroutine can't handle instructions
with side-effects well that are inserted before the coro.begin.

rdar://81113950

Author: aschwaighofer

lib/IRGen/GenCall.cpp

@@ -174,7 +174,12 @@ Alignment IRGenModule::getAsyncContextAlignment() const {
 void IRGenFunction::setupAsync(unsigned asyncContextIndex) {
   llvm::Value *c = CurFn->getArg(asyncContextIndex);
   asyncContextLocation = createAlloca(c->getType(), IGM.getPointerAlignment());
-  Builder.CreateStore(c, asyncContextLocation);
+
+  IRBuilder builder(IGM.getLLVMContext(), IGM.DebugInfo != nullptr);
+  // Insert the stores after the coro.begin.
+  builder.SetInsertPoint(getEarliestInsertionPoint()->getParent(),
+                         getEarliestInsertionPoint()->getIterator());
+  builder.CreateStore(c, asyncContextLocation);
 }
 
 llvm::Value *IRGenFunction::getAsyncTask() {
@@ -3771,6 +3776,11 @@ emitRetconCoroutineEntry(IRGenFunction &IGF, CanSILFunctionType fnType,
   // Set the coroutine handle; this also flags that is a coroutine so that
   // e.g. dynamic allocas use the right code generation.
   IGF.setCoroutineHandle(hdl);
+
+  auto *pt = IGF.Builder.IRBuilderBase::CreateAlloca(IGF.IGM.Int1Ty,
+                                                     /*array size*/ nullptr,
+                                                     "earliest insert point");
+  IGF.setEarliestInsertionPoint(pt);
 }
 
 void irgen::emitAsyncFunctionEntry(IRGenFunction &IGF,
@@ -3796,6 +3806,11 @@ void irgen::emitAsyncFunctionEntry(IRGenFunction &IGF,
   // Set the coroutine handle; this also flags that is a coroutine so that
   // e.g. dynamic allocas use the right code generation.
   IGF.setCoroutineHandle(hdl);
+  auto *pt = IGF.Builder.IRBuilderBase::CreateAlloca(IGF.IGM.Int1Ty,
+                                                     /*array size*/ nullptr,
+                                                     "earliest insert point");
+  IGF.setEarliestInsertionPoint(pt);
+  IGF.setupAsync(asyncContextIndex);
 }
 
 void irgen::emitYieldOnceCoroutineEntry(
@@ -4054,6 +4069,11 @@ Address IRGenFunction::createErrorResultSlot(SILType errorType, bool isAsync) {
   auto addr = createAlloca(errorTI.getStorageType(),
                            errorTI.getFixedAlignment(), "swifterror");
 
+  if (!isAsync) {
+    builder.SetInsertPoint(getEarliestInsertionPoint()->getParent(),
+                           getEarliestInsertionPoint()->getIterator());
+  }
+
   // Only add the swifterror attribute on ABIs that pass it in a register.
   // We create a shadow stack location of the swifterror parameter for the
   // debugger on platforms that pass swifterror by reference and so we can't
@@ -4131,6 +4151,7 @@ void IRGenFunction::emitPrologue() {
   AllocaIP = Builder.IRBuilderBase::CreateAlloca(IGM.Int1Ty,
                                                  /*array size*/ nullptr,
                                                  "alloca point");
+  EarliestIP = AllocaIP;
 }
 
 /// Emit a branch to the return block and set the insert point there.
@@ -4170,6 +4191,8 @@ bool IRGenFunction::emitBranchToReturnBB() {
 
 /// Emit the epilogue for the function.
 void IRGenFunction::emitEpilogue() {
+  if (EarliestIP != AllocaIP)
+    EarliestIP->eraseFromParent();
   // Destroy the alloca insertion point.
   AllocaIP->eraseFromParent();
 }

lib/IRGen/GenFunc.cpp

@@ -1275,8 +1275,6 @@ static llvm::Value *emitPartialApplicationForwarder(IRGenModule &IGM,
         FunctionPointer::Kind(
             FunctionPointer::BasicKind::AsyncFunctionPointer)));
 
-    subIGF.setupAsync(asyncContextIdx);
-
     //auto *calleeAFP = staticFnPtr->getDirectPointer();
     LinkEntity entity = LinkEntity::forPartialApplyForwarder(fwd);
     assert(!asyncFunctionPtr &&

lib/IRGen/GenThunk.cpp

@@ -255,7 +255,6 @@ void IRGenThunk::emit() {
     auto asyncContextIdx = Signature::forAsyncEntry(
                                IGF.IGM, origTy, /*useSpecialConvention*/ false)
                                .getAsyncContextIndex();
-    IGF.setupAsync(asyncContextIdx);
 
     auto entity = LinkEntity::forDispatchThunk(declRef);
     emitAsyncFunctionEntry(IGF, *asyncLayout, entity, asyncContextIdx);

lib/IRGen/IRGenFunction.h

@@ -364,8 +364,20 @@ class IRGenFunction {
 private:
   llvm::Instruction *AllocaIP;
   const SILDebugScope *DbgScope;
+  /// The insertion point where we should but instructions we would normally put
+  /// at the beginning of the function. LLVM's coroutine lowering really does
+  /// not like it if we put instructions with side-effectrs before the
+  /// coro.begin.
+  llvm::Instruction *EarliestIP;
 
-//--- Reference-counting methods -----------------------------------------------
+public:
+  void setEarliestInsertionPoint(llvm::Instruction *inst) { EarliestIP = inst; }
+  /// Returns the first insertion point before which we should insert
+  /// instructions which have side-effects.
+  llvm::Instruction *getEarliestInsertionPoint() const { return EarliestIP; }
+
+  //--- Reference-counting methods
+  //-----------------------------------------------
 public:
   // Returns the default atomicity of the module.
   Atomicity getDefaultAtomicity();

lib/IRGen/IRGenSIL.cpp

@@ -687,7 +687,8 @@ class IRGenSILFunction :
       return;
 
     llvm::IRBuilder<> ZeroInitBuilder(AI->getNextNode());
-
+    ZeroInitBuilder.SetInsertPoint(getEarliestInsertionPoint()->getParent(),
+                                   getEarliestInsertionPoint()->getIterator());
     // No debug location is how LLVM marks prologue instructions.
     ZeroInitBuilder.SetCurrentDebugLocation(nullptr);
     ZeroInitBuilder.CreateMemSet(
@@ -1731,12 +1732,6 @@ IRGenSILFunction::IRGenSILFunction(IRGenModule &IGM, SILFunction *f)
   if (f->isDynamicallyReplaceable() && !f->isAsync()) {
     IGM.createReplaceableProlog(*this, f);
   }
-
-  if (f->getLoweredFunctionType()->isAsync()) {
-    setupAsync(Signature::forAsyncEntry(IGM, f->getLoweredFunctionType(),
-                                        /*useSpecialConvention*/ false)
-                   .getAsyncContextIndex());
-  }
 }
 
 IRGenSILFunction::~IRGenSILFunction() {
@@ -1935,11 +1930,6 @@ static void emitEntryPointArgumentsNativeCC(IRGenSILFunction &IGF,
                                    witnessMetadata);
   }
 
-  // Bind the error result by popping it off the parameter list.
-  if (funcTy->hasErrorResult() && !funcTy->isAsync()) {
-    IGF.setCallerErrorResultSlot(emission->getCallerErrorResultArgument());
-  }
-
   // The coroutine context should be the first parameter.
   switch (funcTy->getCoroutineKind()) {
   case SILCoroutineKind::None:
@@ -1965,6 +1955,11 @@ static void emitEntryPointArgumentsNativeCC(IRGenSILFunction &IGF,
     }
   }
 
+  // Bind the error result by popping it off the parameter list.
+  if (funcTy->hasErrorResult() && !funcTy->isAsync()) {
+    IGF.setCallerErrorResultSlot(emission->getCallerErrorResultArgument());
+  }
+
   SILFunctionConventions conv(funcTy, IGF.getSILModule());
 
   // The 'self' argument might be in the context position, which is

test/IRGen/dynamic_replaceable_coroutine.swift

@@ -0,0 +1,31 @@
+// RUN: %target-swift-frontend -module-name A -enable-implicit-dynamic -emit-ir %s | %FileCheck %s
+
+
+extension Int {
+  public struct Thing {
+    var _int: Int
+    init(_ int: Int) {
+      self._int = int
+    }
+  }
+
+  public var thing: Thing {
+    get { Thing(self) }
+    // Make sure the initialization of `thing` is after the dynamic replacement
+    // check. Coro splitting does not like memsets before the coro.begin.
+
+    // CHECK: define{{.*}} swiftcc { i8*, %TSi1AE5ThingV* } @"$sSi1AE5thingSiAAE5ThingVvM"
+    // CHECK: call i8* @swift_getFunctionReplacement
+    // CHECK: br
+    // CHECK: original_entry:
+    // CHECK: [[FRAMEPTR:%.*]] = bitcast i8* %0 to
+    // CHECK: [[THING:%.*]] = getelementptr inbounds {{.*}}* [[FRAMEPTR]], i32 0
+    // CHECK: [[THING2:%.*]] = bitcast %TSi1AE5ThingV* [[THING]] to i8*
+    // CHECK: call void @llvm.memset{{.*}}(i8* {{.*}} [[THING2]]
+    // CHECK: ret
+    _modify {
+      var thing = Thing(self)
+      yield &thing
+    }
+  }
+}