[Async CC] Add coroutine fields to context.

The following fields are now available when the function is a coroutine:
- TaskContinuationFunction * __ptrauth(...) yieldToCaller?
- TaskContinuationFunction * __ptrauth(...) resumeFromYield?
- TaskContinuationFunction * __ptrauth(...) abortFromYield?
- ExecutorRef calleeActorDuringYield?
- YieldTypes yieldValues...

These fields have yet to be filled in.

The following field are now available when the function is NOT a
coroutine (whereas previously they were always available):
- ResultTypes directResults...

Author: nate-chandler

lib/IRGen/GenCall.cpp

@@ -97,6 +97,8 @@ AsyncContextLayout irgen::getAsyncContextLayout(
   SmallVector<const TypeInfo *, 4> typeInfos;
   SmallVector<SILType, 4> valTypes;
   SmallVector<AsyncContextLayout::ArgumentInfo, 4> paramInfos;
+  bool isCoroutine = originalType->isCoroutine();
+  SmallVector<SILYieldInfo, 4> yieldInfos;
   SmallVector<SILResultInfo, 4> indirectReturnInfos;
   SmallVector<SILResultInfo, 4> directReturnInfos;
 
@@ -141,6 +143,11 @@ AsyncContextLayout irgen::getAsyncContextLayout(
     typeInfos.push_back(&ti);
   }
 
+  if (isCoroutine) {
+    // SwiftPartialFunction * __ptrauth(...) yieldToCaller?;
+    addTaskContinuationFunction();
+  }
+
   //   SwiftError *errorResult;
   auto errorCanType = IGF.IGM.Context.getExceptionType();
   auto errorType = SILType::getPrimitiveObjectType(errorCanType);
@@ -160,15 +167,33 @@ AsyncContextLayout irgen::getAsyncContextLayout(
     indirectReturnInfos.push_back(indirectResult);
   }
 
-  //     ResultTypes directResults...;
-  auto directResults = fnConv.getDirectSILResults();
-  for (auto result : directResults) {
-    auto ty =
-        fnConv.getSILType(result, IGF.IGM.getMaximalTypeExpansionContext());
-    auto &ti = IGF.getTypeInfoForLowered(ty.getASTType());
-    valTypes.push_back(ty);
-    typeInfos.push_back(&ti);
-    directReturnInfos.push_back(result);
+  // union {
+  if (isCoroutine) {
+    // SwiftPartialFunction * __ptrauth(...) resumeFromYield?
+    addTaskContinuationFunction();
+    // SwiftPartialFunction * __ptrauth(...) abortFromYield?
+    addTaskContinuationFunction();
+    // SwiftActor * __ptrauth(...) calleeActorDuringYield?
+    addExecutor();
+    // YieldTypes yieldValues...
+    for (auto yield : fnConv.getYields()) {
+      auto ty =
+          fnConv.getSILType(yield, IGF.IGM.getMaximalTypeExpansionContext());
+      auto &ti = IGF.getTypeInfoForLowered(ty.getASTType());
+      valTypes.push_back(ty);
+      typeInfos.push_back(&ti);
+      yieldInfos.push_back(yield);
+    }
+  } else {
+    //     ResultTypes directResults...;
+    for (auto result : fnConv.getDirectSILResults()) {
+      auto ty =
+          fnConv.getSILType(result, IGF.IGM.getMaximalTypeExpansionContext());
+      auto &ti = IGF.getTypeInfoForLowered(ty.getASTType());
+      valTypes.push_back(ty);
+      typeInfos.push_back(&ti);
+      directReturnInfos.push_back(result);
+    }
   }
 
   //   SelfType self?;
@@ -257,7 +282,8 @@ AsyncContextLayout irgen::getAsyncContextLayout(
       IGF.IGM, LayoutStrategy::Optimal, valTypes, typeInfos, IGF, originalType,
       substitutedType, substitutionMap, std::move(bindings),
       trailingWitnessInfo, errorType, canHaveValidError, paramInfos,
-      indirectReturnInfos, directReturnInfos, localContextInfo);
+      isCoroutine, yieldInfos, indirectReturnInfos, directReturnInfos,
+      localContextInfo);
 }
 
 AsyncContextLayout::AsyncContextLayout(
@@ -267,14 +293,16 @@ AsyncContextLayout::AsyncContextLayout(
     SubstitutionMap substitutionMap, NecessaryBindings &&bindings,
     Optional<TrailingWitnessInfo> trailingWitnessInfo, SILType errorType,
     bool canHaveValidError, ArrayRef<ArgumentInfo> argumentInfos,
+    bool isCoroutine, ArrayRef<SILYieldInfo> yieldInfos,
     ArrayRef<SILResultInfo> indirectReturnInfos,
     ArrayRef<SILResultInfo> directReturnInfos,
     Optional<AsyncContextLayout::ArgumentInfo> localContextInfo)
     : StructLayout(IGM, /*decl=*/nullptr, LayoutKind::NonHeapObject, strategy,
                    fieldTypeInfos, /*typeToFill*/ nullptr),
       IGF(IGF), originalType(originalType), substitutedType(substitutedType),
       substitutionMap(substitutionMap), errorType(errorType),
-      canHaveValidError(canHaveValidError),
+      canHaveValidError(canHaveValidError), isCoroutine(isCoroutine),
+      yieldInfos(yieldInfos.begin(), yieldInfos.end()),
       directReturnInfos(directReturnInfos.begin(), directReturnInfos.end()),
       indirectReturnInfos(indirectReturnInfos.begin(),
                           indirectReturnInfos.end()),
@@ -284,6 +312,11 @@ AsyncContextLayout::AsyncContextLayout(
 #ifndef NDEBUG
   assert(fieldTypeInfos.size() == fieldTypes.size() &&
          "type infos don't match types");
+  if (isCoroutine) {
+    assert(directReturnInfos.empty());
+  } else {
+    assert(yieldInfos.empty());
+  }
   if (!bindings.empty()) {
     assert(fieldTypeInfos.size() >= 2 && "no field for bindings");
     auto fixedBindingsField =

lib/IRGen/GenCall.h

@@ -97,7 +97,7 @@ namespace irgen {
       ResumeParent = 1,
       ResumeParentExecutor = 2,
       Flags = 3,
-      Error = 4,
+      YieldToParent = 4,
     };
     enum class FixedCount : unsigned {
       Parent = 1,
@@ -111,6 +111,8 @@ namespace irgen {
     SubstitutionMap substitutionMap;
     SILType errorType;
     bool canHaveValidError;
+    bool isCoroutine;
+    SmallVector<SILYieldInfo, 4> yieldInfos;
     SmallVector<SILResultInfo, 4> directReturnInfos;
     SmallVector<SILResultInfo, 4> indirectReturnInfos;
     Optional<ArgumentInfo> localContextInfo;
@@ -126,20 +128,56 @@ namespace irgen {
       return (unsigned)FixedIndex::ResumeParentExecutor;
     }
     unsigned getFlagsIndex() { return (unsigned)FixedIndex::Flags; }
-    unsigned getErrorIndex() { return (unsigned)FixedIndex::Error; }
+    unsigned getYieldToParentIndex() {
+      assert(isCoroutine);
+      return (unsigned)FixedIndex::YieldToParent;
+    }
+    unsigned getErrorIndex() {
+      return (isCoroutine ? getYieldToParentIndex() : getFlagsIndex()) + 1;
+    }
     unsigned getFirstIndirectReturnIndex() {
       return getErrorIndex() + getErrorCount();
     }
     unsigned getIndexAfterIndirectReturns() {
       return getFirstIndirectReturnIndex() + getIndirectReturnCount();
     }
     unsigned getFirstDirectReturnIndex() {
+      assert(!isCoroutine);
       return getIndexAfterIndirectReturns();
     }
     unsigned getIndexAfterDirectReturns() {
+      assert(!isCoroutine);
       return getFirstDirectReturnIndex() + getDirectReturnCount();
     }
-    unsigned getFirstArgumentIndex() { return getIndexAfterDirectReturns(); }
+    unsigned getResumeFromYieldIndex() {
+      assert(isCoroutine);
+      return getIndexAfterIndirectReturns();
+    }
+    unsigned getAbortFromYieldIndex() {
+      assert(isCoroutine);
+      return getResumeFromYieldIndex() + 1;
+      ;
+    }
+    unsigned getCalleeExecutorDuringYieldIndex() {
+      assert(isCoroutine);
+      return getAbortFromYieldIndex() + 1;
+    }
+    unsigned getFirstYieldIndex() {
+      assert(isCoroutine);
+      return getCalleeExecutorDuringYieldIndex() + 1;
+    }
+    unsigned getIndexAfterYields() {
+      assert(isCoroutine);
+      return getFirstYieldIndex() + getYieldCount();
+    }
+    unsigned getIndexAfterUnion() {
+      if (isCoroutine) {
+        return getIndexAfterYields();
+      } else {
+        return getIndexAfterDirectReturns();
+      }
+    }
+    unsigned getFirstArgumentIndex() { return getIndexAfterUnion(); }
     unsigned getIndexAfterArguments() {
       return getFirstArgumentIndex() + getArgumentCount();
     }
@@ -230,10 +268,18 @@ namespace irgen {
       return getElement(getSelfWitnessTableIndex());
     }
 
-    unsigned getDirectReturnCount() { return directReturnInfos.size(); }
+    unsigned getDirectReturnCount() {
+      assert(!isCoroutine);
+      return directReturnInfos.size();
+    }
     ElementLayout getDirectReturnLayout(unsigned index) {
+      assert(!isCoroutine);
       return getElement(getFirstDirectReturnIndex() + index);
     }
+    unsigned getYieldCount() {
+      assert(isCoroutine);
+      return yieldInfos.size();
+    }
 
     AsyncContextLayout(
         IRGenModule &IGM, LayoutStrategy strategy, ArrayRef<SILType> fieldTypes,
@@ -242,6 +288,7 @@ namespace irgen {
         SubstitutionMap substitutionMap, NecessaryBindings &&bindings,
         Optional<TrailingWitnessInfo> trailingWitnessInfo, SILType errorType,
         bool canHaveValidError, ArrayRef<ArgumentInfo> argumentInfos,
+        bool isCoroutine, ArrayRef<SILYieldInfo> yieldInfos,
         ArrayRef<SILResultInfo> indirectReturnInfos,
         ArrayRef<SILResultInfo> directReturnInfos,
         Optional<ArgumentInfo> localContextInfo);