Partial apply for coroutines

asavonic · asl · commit 3f5278ac4521 · 2024-08-06T10:30:10.000-07:00
The patch adds lowering of partial_apply instructions for coroutines. This pattern seems to trigger a lot of type mismatch errors in IRGen, because coroutine functions are not substituted in the same way as regular functions (see the patch 07f03bd "Use pattern substitutions to consistently abstract yields" for more details). The odd type conversions in the patch are related to this issue, and these should be checked carefully. Perhaps it is better to enable substitutions for coroutine functions instead (at least for some cases). Other than that, lowering of partial_apply for coroutines is straightforward: we generate another coroutine that captures arguments passed to the partial_apply instructions. It calls the original coroutine for yields (first return) and yields the resulting values. Then it calls the original function's continuation for return or unwind, and forwards them to the caller as well. After IRGen, LLVM's Coroutine pass transforms the generated coroutine (along with all other coroutines) and eliminates llvm.coro.* intrinsics. LIT tests check LLVM IR after this transformation.
diff --git a/lib/IRGen/GenDecl.cpp b/lib/IRGen/GenDecl.cpp
@@ -6163,6 +6163,7 @@ IRGenModule::getAddrOfContinuationPrototype(CanSILFunctionType fnType) {
   llvm::Function *&entry = GlobalFuncs[entity];
   if (entry) return entry;
 
+  GenericContextScope scope(*this, fnType->getInvocationGenericSignature());
   auto signature = Signature::forCoroutineContinuation(*this, fnType);
   LinkInfo link = LinkInfo::get(*this, entity, NotForDefinition);
   entry = createFunction(*this, link, signature);
diff --git a/lib/IRGen/GenFunc.cpp b/lib/IRGen/GenFunc.cpp
@@ -91,6 +91,7 @@
 #include "llvm/Support/Debug.h"
 
 #include "BitPatternBuilder.h"
+#include "CallEmission.h"
 #include "Callee.h"
 #include "ConstantBuilder.h"
 #include "EnumPayload.h"
@@ -109,11 +110,11 @@
 #include "HeapTypeInfo.h"
 #include "IRGenDebugInfo.h"
 #include "IRGenFunction.h"
+#include "IRGenMangler.h"
 #include "IRGenModule.h"
 #include "IndirectTypeInfo.h"
 #include "ScalarPairTypeInfo.h"
 #include "Signature.h"
-#include "IRGenMangler.h"
 
 using namespace swift;
 using namespace irgen;
@@ -1119,11 +1120,11 @@ class PartialApplicationForwarderEmission {
   virtual void addDynamicFunctionContext(Explosion &explosion) = 0;
   virtual void addDynamicFunctionPointer(Explosion &explosion) = 0;
 
-  void addSelf(Explosion &explosion) { addArgument(explosion); }
-  void addWitnessSelfMetadata(llvm::Value *value) {
+  virtual void addSelf(Explosion &explosion) { addArgument(explosion); }
+  virtual void addWitnessSelfMetadata(llvm::Value *value) {
     addArgument(value);
   }
-  void addWitnessSelfWitnessTable(llvm::Value *value) {
+  virtual void addWitnessSelfWitnessTable(llvm::Value *value) {
     addArgument(value);
   }
   virtual void forwardErrorResult() = 0;
@@ -1136,6 +1137,14 @@ class PartialApplicationForwarderEmission {
   virtual void end(){};
   virtual ~PartialApplicationForwarderEmission() {}
 };
+
+static Size getYieldOnceCoroutineBufferSize(IRGenModule &IGM) {
+  return NumWords_YieldOnceBuffer * IGM.getPointerSize();
+}
+static Alignment getYieldOnceCoroutineBufferAlignment(IRGenModule &IGM) {
+  return IGM.getPointerAlignment();
+}
+
 class SyncPartialApplicationForwarderEmission
     : public PartialApplicationForwarderEmission {
   using super = PartialApplicationForwarderEmission;
@@ -1422,6 +1431,239 @@ class AsyncPartialApplicationForwarderEmission
     super::end();
   }
 };
+
+class CoroPartialApplicationForwarderEmission
+    : public PartialApplicationForwarderEmission {
+  using super = PartialApplicationForwarderEmission;
+
+private:
+  llvm::Value *Self;
+  llvm::Value *FirstData;
+  llvm::Value *SecondData;
+  WitnessMetadata Witness;
+
+public:
+  CoroPartialApplicationForwarderEmission(
+      IRGenModule &IGM, IRGenFunction &subIGF, llvm::Function *fwd,
+      const std::optional<FunctionPointer> &staticFnPtr, bool calleeHasContext,
+      const Signature &origSig, CanSILFunctionType origType,
+      CanSILFunctionType substType, CanSILFunctionType outType,
+      SubstitutionMap subs, HeapLayout const *layout,
+      ArrayRef<ParameterConvention> conventions)
+      : PartialApplicationForwarderEmission(
+            IGM, subIGF, fwd, staticFnPtr, calleeHasContext, origSig, origType,
+            substType, outType, subs, layout, conventions),
+        Self(nullptr), FirstData(nullptr), SecondData(nullptr) {}
+
+  void begin() override {
+    auto prototype = subIGF.IGM.getOpaquePtr(
+        subIGF.IGM.getAddrOfContinuationPrototype(origType));
+
+    // Use malloc and free as our allocator.
+    auto allocFn = subIGF.IGM.getOpaquePtr(subIGF.IGM.getMallocFn());
+    auto deallocFn = subIGF.IGM.getOpaquePtr(subIGF.IGM.getFreeFn());
+
+    // Call the right 'llvm.coro.id.retcon' variant.
+    llvm::Value *buffer = origParams.claimNext();
+    llvm::Value *id = subIGF.Builder.CreateIntrinsicCall(
+        llvm::Intrinsic::coro_id_retcon_once,
+        {llvm::ConstantInt::get(
+             subIGF.IGM.Int32Ty,
+             getYieldOnceCoroutineBufferSize(subIGF.IGM).getValue()),
+         llvm::ConstantInt::get(
+             subIGF.IGM.Int32Ty,
+             getYieldOnceCoroutineBufferAlignment(subIGF.IGM).getValue()),
+         buffer, prototype, allocFn, deallocFn});
+
+    // Call 'llvm.coro.begin', just for consistency with the normal pattern.
+    // This serves as a handle that we can pass around to other intrinsics.
+    auto hdl = subIGF.Builder.CreateIntrinsicCall(
+        llvm::Intrinsic::coro_begin,
+        {id, llvm::ConstantPointerNull::get(subIGF.IGM.Int8PtrTy)});
+
+    // Set the coroutine handle; this also flags that is a coroutine so that
+    // e.g. dynamic allocas use the right code generation.
+    subIGF.setCoroutineHandle(hdl);
+
+    auto *pt = subIGF.Builder.IRBuilderBase::CreateAlloca(
+        subIGF.IGM.Int1Ty,
+        /*array size*/ nullptr, "earliest insert point");
+    subIGF.setEarliestInsertionPoint(pt);
+  }
+
+  void gatherArgumentsFromApply() override {
+    super::gatherArgumentsFromApply(false);
+  }
+  llvm::Value *getDynamicFunctionPointer() override {
+    llvm::Value *Ret = SecondData;
+    SecondData = nullptr;
+    return Ret;
+  }
+  llvm::Value *getDynamicFunctionContext() override {
+    llvm::Value *Ret = FirstData;
+    FirstData = nullptr;
+    return Ret;
+  }
+  void addDynamicFunctionContext(Explosion &explosion) override {
+    assert(!Self && "context value overrides 'self'");
+    FirstData = explosion.claimNext();
+  }
+  void addDynamicFunctionPointer(Explosion &explosion) override {
+    SecondData = explosion.claimNext();
+  }
+  void addSelf(Explosion &explosion) override {
+    assert(!FirstData && "'self' overrides another context value");
+    if (!hasSelfContextParameter(origType)) {
+      // witness methods can be declared on types that are not classes. Pass
+      // such "self" argument as a plain argument.
+      addArgument(explosion);
+      return;
+    }
+    Self = explosion.claimNext();
+    FirstData = Self;
+  }
+
+  void addWitnessSelfMetadata(llvm::Value *value) override {
+    Witness.SelfMetadata = value;
+  }
+
+  void addWitnessSelfWitnessTable(llvm::Value *value) override {
+    Witness.SelfWitnessTable = value;
+  }
+
+  void forwardErrorResult() override {
+    bool isTypedError = origConv.isTypedError();
+    SILType origErrorTy =
+        origConv.getSILErrorType(subIGF.IGM.getMaximalTypeExpansionContext());
+    auto errorAlignment =
+        isTypedError ? subIGF.IGM.getPointerAlignment()
+                     : cast<FixedTypeInfo>(subIGF.getTypeInfo(origErrorTy))
+                           .getFixedAlignment();
+    auto errorStorageType =
+        isTypedError ? IGM.Int8PtrTy
+                     : cast<FixedTypeInfo>(subIGF.getTypeInfo(origErrorTy))
+                           .getStorageType();
+    llvm::Value *errorResultPtr = origParams.claimNext();
+    subIGF.setCallerErrorResultSlot(
+        Address(errorResultPtr, errorStorageType, errorAlignment));
+  }
+
+  Explosion callCoroutine(FunctionPointer &fnPtr) {
+    Callee callee({origType, substType, subs}, fnPtr, FirstData, SecondData);
+
+    std::unique_ptr<CallEmission> emitSuspend =
+        getCallEmission(subIGF, Self, std::move(callee));
+
+    emitSuspend->begin();
+    emitSuspend->setArgs(args, /*isOutlined=*/false, &Witness);
+    Explosion yieldedValues;
+    emitSuspend->emitToExplosion(yieldedValues, /*isOutlined=*/false);
+    emitSuspend->end();
+    emitSuspend->claimTemporaries().destroyAll(subIGF);
+
+    if (origConv.getSILResultType(subIGF.IGM.getMaximalTypeExpansionContext())
+            .hasTypeParameter()) {
+
+      ArrayRef<llvm::Value *> yieldValues = yieldedValues.claimAll();
+      ArrayRef<llvm::Type *> retTypes =
+          cast<llvm::StructType>(fwd->getReturnType())->elements();
+      Explosion yieldCoerced;
+      assert(yieldValues.size() == retTypes.size() &&
+             "mismatch between return types of the wrapper and the callee");
+      for (unsigned i = 0; i < yieldValues.size(); ++i) {
+        llvm::Value *v = yieldValues[i];
+        if (v->getType() != retTypes[i]) {
+          v = subIGF.coerceValue(v, retTypes[i], subIGF.IGM.DataLayout);
+        }
+        yieldCoerced.add(v);
+      }
+      return yieldCoerced;
+    }
+
+    return yieldedValues;
+  }
+
+  llvm::CallInst *createCall(FunctionPointer &fnPtr) override {
+    /// Call the wrapped coroutine
+    ///
+    Address calleeBuf = emitAllocYieldOnceCoroutineBuffer(subIGF);
+    llvm::Value *calleeHandle = calleeBuf.getAddress();
+    args.insert(0, calleeHandle);
+    Explosion yieldedValues = callCoroutine(fnPtr);
+
+    /// Get the continuation function pointer
+    ///
+    PointerAuthInfo newAuthInfo =
+        fnPtr.getAuthInfo().getCorrespondingCodeAuthInfo();
+    FunctionPointer contFn = FunctionPointer::createSigned(
+        FunctionPointer::Kind::Function, yieldedValues.claimNext(), newAuthInfo,
+        Signature::forCoroutineContinuation(subIGF.IGM, origType));
+
+    /// Forward the remaining yields of the wrapped coroutine
+    ///
+    llvm::Value *condUnwind = emitYield(subIGF, substType, yieldedValues);
+
+    llvm::BasicBlock *unwindBB = subIGF.createBasicBlock("unwind");
+    llvm::BasicBlock *resumeBB = subIGF.createBasicBlock("resume");
+    llvm::BasicBlock *cleanupBB = subIGF.createBasicBlock("cleanup");
+    subIGF.CurFn->insert(subIGF.CurFn->end(), unwindBB);
+    subIGF.CurFn->insert(subIGF.CurFn->end(), resumeBB);
+    subIGF.CurFn->insert(subIGF.CurFn->end(), cleanupBB);
+    subIGF.Builder.CreateCondBr(condUnwind, unwindBB, resumeBB);
+
+    /// Call for the results
+    ///
+    subIGF.Builder.SetInsertPoint(resumeBB);
+
+    auto isResume = llvm::ConstantInt::get(IGM.Int1Ty, /*isAbort*/ false);
+    auto *call = subIGF.Builder.CreateCall(contFn, {calleeHandle, isResume});
+
+    /// Emit coro_end for results and forward them
+    ///
+    llvm::Type *callTy = call->getType();
+    llvm::Value *noneToken =
+        llvm::ConstantTokenNone::get(subIGF.Builder.getContext());
+    llvm::Value *resultToken = nullptr;
+    if (callTy->isVoidTy()) {
+      resultToken = noneToken;
+    } else if (llvm::StructType *sty = dyn_cast<llvm::StructType>(callTy)) {
+      Explosion splitCall;
+      subIGF.emitAllExtractValues(call, sty, splitCall);
+      resultToken = subIGF.Builder.CreateIntrinsicCall(
+          llvm::Intrinsic::coro_end_results, splitCall.claimAll());
+    } else {
+      resultToken = subIGF.Builder.CreateIntrinsicCall(
+          llvm::Intrinsic::coro_end_results, call);
+    }
+
+    llvm::Value *fwdHandle = subIGF.getCoroutineHandle();
+    subIGF.Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_end,
+                                       {fwdHandle, isResume, resultToken});
+    subIGF.Builder.CreateBr(cleanupBB);
+
+    /// Emit coro_end for unwind
+    ///
+    subIGF.Builder.SetInsertPoint(unwindBB);
+    auto isUnwind = llvm::ConstantInt::get(IGM.Int1Ty, /*isAbort*/ true);
+    subIGF.Builder.CreateCall(contFn, {calleeHandle, isUnwind});
+    subIGF.Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_end,
+                                       {fwdHandle, isUnwind, noneToken});
+    subIGF.Builder.CreateBr(cleanupBB);
+
+    subIGF.Builder.SetInsertPoint(cleanupBB);
+    emitDeallocYieldOnceCoroutineBuffer(subIGF, calleeBuf);
+    llvm::Instruction *cleanupPt = subIGF.Builder.CreateUnreachable();
+    subIGF.Builder.SetInsertPoint(cleanupPt);
+
+    return nullptr;
+  }
+
+  void createReturn(llvm::CallInst *call) override {
+    // Do nothing, yield/return/unwind blocks are already created in createCall.
+  }
+  void end() override { super::end(); }
+};
+
 std::unique_ptr<PartialApplicationForwarderEmission>
 getPartialApplicationForwarderEmission(
     IRGenModule &IGM, IRGenFunction &subIGF, llvm::Function *fwd,
@@ -1434,6 +1676,11 @@ getPartialApplicationForwarderEmission(
     return std::make_unique<AsyncPartialApplicationForwarderEmission>(
         IGM, subIGF, fwd, staticFnPtr, calleeHasContext, origSig, origType,
         substType, outType, subs, layout, conventions);
+  } else if (origType->isCoroutine()) {
+    return std::make_unique<CoroPartialApplicationForwarderEmission>(
+        IGM, subIGF, fwd, staticFnPtr, calleeHasContext, origSig, origType,
+        substType, outType, subs, layout, conventions);
+
   } else {
     return std::make_unique<SyncPartialApplicationForwarderEmission>(
         IGM, subIGF, fwd, staticFnPtr, calleeHasContext, origSig, origType,
@@ -1611,8 +1858,6 @@ static llvm::Value *emitPartialApplicationForwarder(
           ti.isSingleSwiftRetainablePointer(ResilienceExpansion::Maximal))
         ref = subIGF.coerceValue(rawData, ti.getStorageType(),
                                  subIGF.IGM.DataLayout);
-      else
-        ref = subIGF.Builder.CreateBitCast(rawData, ti.getStorageType());
       param.add(ref);
       bindPolymorphicParameter(subIGF, origType, substType, param, paramI);
       (void)param.claimAll();
@@ -1687,7 +1932,7 @@ static llvm::Value *emitPartialApplicationForwarder(
     auto argIndex = emission->getCurrentArgumentIndex();
     if (haveContextArgument)
       argIndex += polyArgs.size();
-    if (origType->isAsync())
+    if (origType->isAsync() || origType->isCoroutine())
       argIndex += 1;
 
     llvm::Type *expectedArgTy = origSig.getType()->getParamType(argIndex);
@@ -1712,10 +1957,15 @@ static llvm::Value *emitPartialApplicationForwarder(
     } else {
       argValue = subIGF.Builder.CreateBitCast(rawData, expectedArgTy);
     }
-    emission->addArgument(argValue);
+    if (haveContextArgument) {
+      Explosion e;
+      e.add(argValue);
+      emission->addDynamicFunctionContext(e);
+    } else
+      emission->addArgument(argValue);
 
-  // If there's a data pointer required, grab it and load out the
-  // extra, previously-curried parameters.
+    // If there's a data pointer required, grab it and load out the
+    // extra, previously-curried parameters.
   } else {
     unsigned origParamI = outType->getParameters().size();
     unsigned extraFieldIndex = 0;
@@ -1975,8 +2225,8 @@ static llvm::Value *emitPartialApplicationForwarder(
 
   llvm::CallInst *call = emission->createCall(fnPtr);
 
-  if (!origType->isAsync() && addressesToDeallocate.empty() && !needsAllocas &&
-      (!consumesContext || !dependsOnContextLifetime))
+  if (!origType->isAsync() && !origType->isCoroutine() && addressesToDeallocate.empty() &&
+      !needsAllocas &&  (!consumesContext || !dependsOnContextLifetime))
     call->setTailCall();
 
   // Deallocate everything we allocated above.
@@ -2030,7 +2280,7 @@ std::optional<StackAddress> irgen::emitFunctionPartialApplication(
   bool considerParameterSources = true;
   for (auto param : params) {
     SILType argType = IGF.IGM.silConv.getSILType(
-        param, origType, IGF.IGM.getMaximalTypeExpansionContext());
+        param, substType, IGF.IGM.getMaximalTypeExpansionContext());
     auto argLoweringTy = getArgumentLoweringType(argType.getASTType(), param,
                                                  outType->isNoEscape());
     auto &ti = IGF.getTypeInfoForLowered(argLoweringTy);
@@ -2043,7 +2293,7 @@ std::optional<StackAddress> irgen::emitFunctionPartialApplication(
 
   auto addParam = [&](SILParameterInfo param) {
     SILType argType = IGF.IGM.silConv.getSILType(
-        param, origType, IGF.IGM.getMaximalTypeExpansionContext());
+        param, substType, IGF.IGM.getMaximalTypeExpansionContext());
 
     auto argLoweringTy = getArgumentLoweringType(argType.getASTType(), param,
                                                  outType->isNoEscape());
diff --git a/lib/IRGen/IRGenSIL.cpp b/lib/IRGen/IRGenSIL.cpp
@@ -6867,8 +6867,12 @@ void IRGenSILFunction::visitConvertFunctionInst(swift::ConvertFunctionInst *i) {
       fnType->getRepresentation() != SILFunctionType::Representation::Block) {
     auto *fn = temp.claimNext();
     Explosion res;
-    auto sig = IGM.getSignature(fnType);
-    res.add(Builder.CreateBitCast(fn, sig.getType()->getPointerTo()));
+    auto &fnTI = IGM.getTypeInfoForLowered(fnType);
+    auto &fnNative = fnTI.nativeReturnValueSchema(IGM);
+    llvm::Value *newFn =
+        Builder.CreateBitCast(fn, fnNative.getExpandedType(IGM));
+    extractScalarResults(*this, newFn->getType(), newFn, res);
+
     setLoweredExplosion(i, res);
     return;
   }
diff --git a/test/IRGen/partial_apply_coro.sil b/test/IRGen/partial_apply_coro.sil
