[Async CC] Propagate runtime values.

lib/IRGen/GenCall.cpp

@@ -103,6 +103,31 @@ AsyncContextLayout irgen::getAsyncContextLayout(
   auto parameters = substitutedType->getParameters();
   SILFunctionConventions fnConv(substitutedType, IGF.getSILModule());
 
+  // AsyncContext * __ptrauth_swift_async_context_parent Parent;
+  {
+    auto ty = SILType();
+    auto &ti = IGF.IGM.getSwiftContextPtrTypeInfo();
+    valTypes.push_back(ty);
+    typeInfos.push_back(&ti);
+  }
+
+  // TaskContinuationFunction * __ptrauth_swift_async_context_resume
+  //     ResumeParent;
+  {
+    auto ty = SILType();
+    auto &ti = IGF.IGM.getTaskContinuationFunctionPtrTypeInfo();
+    valTypes.push_back(ty);
+    typeInfos.push_back(&ti);
+  }
+
+  // ExecutorRef ResumeParentExecutor;
+  {
+    auto ty = SILType();
+    auto &ti = IGF.IGM.getSwiftExecutorPtrTypeInfo();
+    valTypes.push_back(ty);
+    typeInfos.push_back(&ti);
+  }
+
   //   SwiftError *errorResult;
   auto errorCanType = IGF.IGM.Context.getExceptionType();
   auto errorType = SILType::getPrimitiveObjectType(errorCanType);
@@ -268,11 +293,6 @@ static Alignment getAsyncContextAlignment(IRGenModule &IGM) {
   return IGM.getPointerAlignment();
 }
 
-static llvm::Value *getAsyncTask(IRGenFunction &IGF) {
-  // TODO: Return the appropriate task.
-  return llvm::Constant::getNullValue(IGF.IGM.SwiftTaskPtrTy);
-}
-
 llvm::Value *IRGenFunction::getAsyncTask() {
   assert(isAsync());
   auto *value = CurFn->getArg((unsigned)AsyncFunctionArgumentIndex::Task);
@@ -2165,9 +2185,8 @@ class AsyncCallEmission final : public CallEmission {
   void setArgs(Explosion &llArgs, bool isOutlined,
                WitnessMetadata *witnessMetadata) override {
     Explosion asyncExplosion;
-    asyncExplosion.add(llvm::Constant::getNullValue(IGF.IGM.SwiftTaskPtrTy));
-    asyncExplosion.add(
-        llvm::Constant::getNullValue(IGF.IGM.SwiftExecutorPtrTy));
+    asyncExplosion.add(IGF.getAsyncTask());
+    asyncExplosion.add(IGF.getAsyncExecutor());
     asyncExplosion.add(contextBuffer.getAddress());
     if (getCallee().getRepresentation() ==
         SILFunctionTypeRepresentation::Thick) {
@@ -2176,9 +2195,22 @@ class AsyncCallEmission final : public CallEmission {
     super::setArgs(asyncExplosion, false, witnessMetadata);
     SILFunctionConventions fnConv(getCallee().getSubstFunctionType(),
                                   IGF.getSILModule());
+    auto layout = getAsyncContextLayout();
 
+    // Set caller info into the context.
+    { // caller context
+      Explosion explosion;
+      explosion.add(IGF.getAsyncContext());
+      auto fieldLayout = layout.getParentLayout();
+      saveValue(fieldLayout, explosion, isOutlined);
+    }
+    { // caller executor
+      Explosion explosion;
+      explosion.add(IGF.getAsyncExecutor());
+      auto fieldLayout = layout.getResumeParentExecutorLayout();
+      saveValue(fieldLayout, explosion, isOutlined);
+    }
     // Move all the arguments into the context.
-    auto layout = getAsyncContextLayout();
     for (unsigned index = 0, count = layout.getIndirectReturnCount();
          index < count; ++index) {
       auto fieldLayout = layout.getIndirectReturnLayout(index);
@@ -3354,7 +3386,7 @@ void irgen::emitDeallocYieldManyCoroutineBuffer(IRGenFunction &IGF,
 Address irgen::emitTaskAlloc(IRGenFunction &IGF, llvm::Value *size,
                              Alignment alignment) {
   auto *call = IGF.Builder.CreateCall(IGF.IGM.getTaskAllocFn(),
-                                      {getAsyncTask(IGF), size});
+                                      {IGF.getAsyncTask(), size});
   call->setDoesNotThrow();
   call->setCallingConv(IGF.IGM.SwiftCC);
   call->addAttribute(llvm::AttributeList::FunctionIndex,
@@ -3366,7 +3398,7 @@ Address irgen::emitTaskAlloc(IRGenFunction &IGF, llvm::Value *size,
 void irgen::emitTaskDealloc(IRGenFunction &IGF, Address address,
                             llvm::Value *size) {
   auto *call = IGF.Builder.CreateCall(
-      IGF.IGM.getTaskDeallocFn(), {getAsyncTask(IGF), address.getAddress()});
+      IGF.IGM.getTaskDeallocFn(), {IGF.getAsyncTask(), address.getAddress()});
   call->setDoesNotThrow();
   call->setCallingConv(IGF.IGM.SwiftCC);
   call->addAttribute(llvm::AttributeList::FunctionIndex,

lib/IRGen/GenCall.h

@@ -93,9 +93,15 @@ namespace irgen {
 
   private:
     enum class FixedIndex : unsigned {
-      Error = 0,
+      Parent = 0,
+      ResumeParent = 1,
+      ResumeParentExecutor = 2,
+      Error = 3,
     };
     enum class FixedCount : unsigned {
+      Parent = 1,
+      ResumeParent = 1,
+      ResumeParentExecutor = 1,
       Error = 1,
     };
     IRGenFunction &IGF;
@@ -111,6 +117,13 @@ namespace irgen {
     Optional<TrailingWitnessInfo> trailingWitnessInfo;
     SmallVector<ArgumentInfo, 4> argumentInfos;
 
+    unsigned getParentIndex() { return (unsigned)FixedIndex::Parent; }
+    unsigned getResumeParentIndex() {
+      return (unsigned)FixedIndex::ResumeParent;
+    }
+    unsigned getResumeParentExecutorIndex() {
+      return (unsigned)FixedIndex::ResumeParentExecutor;
+    }
     unsigned getErrorIndex() { return (unsigned)FixedIndex::Error; }
     unsigned getFirstIndirectReturnIndex() {
       return getErrorIndex() + getErrorCount();
@@ -156,6 +169,13 @@ namespace irgen {
     }
 
   public:
+    ElementLayout getParentLayout() { return getElement(getParentIndex()); }
+    ElementLayout getResumeParentLayout() {
+      return getElement(getResumeParentIndex());
+    }
+    ElementLayout getResumeParentExecutorLayout() {
+      return getElement(getResumeParentExecutorIndex());
+    }
     bool canHaveError() { return canHaveValidError; }
     ElementLayout getErrorLayout() { return getElement(getErrorIndex()); }
     unsigned getErrorCount() { return (unsigned)FixedCount::Error; }

lib/IRGen/GenType.cpp

@@ -1453,6 +1453,48 @@ const TypeInfo &TypeConverter::getTypeMetadataPtrTypeInfo() {
   return *TypeMetadataPtrTI;
 }
 
+const TypeInfo &IRGenModule::getSwiftContextPtrTypeInfo() {
+  return Types.getSwiftContextPtrTypeInfo();
+}
+
+const TypeInfo &TypeConverter::getSwiftContextPtrTypeInfo() {
+  if (SwiftContextPtrTI) return *SwiftContextPtrTI;
+  SwiftContextPtrTI = createUnmanagedStorageType(IGM.SwiftContextPtrTy,
+                                                 ReferenceCounting::Unknown,
+                                                 /*isOptional*/false);
+  SwiftContextPtrTI->NextConverted = FirstType;
+  FirstType = SwiftContextPtrTI;
+  return *SwiftContextPtrTI;
+}
+
+const TypeInfo &IRGenModule::getTaskContinuationFunctionPtrTypeInfo() {
+  return Types.getTaskContinuationFunctionPtrTypeInfo();
+}
+
+const TypeInfo &TypeConverter::getTaskContinuationFunctionPtrTypeInfo() {
+  if (TaskContinuationFunctionPtrTI) return *TaskContinuationFunctionPtrTI;
+  TaskContinuationFunctionPtrTI = createUnmanagedStorageType(
+      IGM.TaskContinuationFunctionPtrTy, ReferenceCounting::Unknown,
+      /*isOptional*/ false);
+  TaskContinuationFunctionPtrTI->NextConverted = FirstType;
+  FirstType = TaskContinuationFunctionPtrTI;
+  return *TaskContinuationFunctionPtrTI;
+}
+
+const TypeInfo &IRGenModule::getSwiftExecutorPtrTypeInfo() {
+  return Types.getSwiftExecutorPtrTypeInfo();
+}
+
+const TypeInfo &TypeConverter::getSwiftExecutorPtrTypeInfo() {
+  if (SwiftExecutorPtrTI) return *SwiftExecutorPtrTI;
+  SwiftExecutorPtrTI = createUnmanagedStorageType(IGM.SwiftExecutorPtrTy,
+                                                  ReferenceCounting::Unknown,
+                                                  /*isOptional*/ false);
+  SwiftExecutorPtrTI->NextConverted = FirstType;
+  FirstType = SwiftExecutorPtrTI;
+  return *SwiftExecutorPtrTI;
+}
+
 const LoadableTypeInfo &
 IRGenModule::getReferenceObjectTypeInfo(ReferenceCounting refcounting) {
   switch (refcounting) {

lib/IRGen/GenType.h

@@ -105,6 +105,9 @@ class TypeConverter {
   const LoadableTypeInfo *RawPointerTI = nullptr;
   const LoadableTypeInfo *WitnessTablePtrTI = nullptr;
   const TypeInfo *TypeMetadataPtrTI = nullptr;
+  const TypeInfo *SwiftContextPtrTI = nullptr;
+  const TypeInfo *TaskContinuationFunctionPtrTI = nullptr;
+  const TypeInfo *SwiftExecutorPtrTI = nullptr;
   const TypeInfo *ObjCClassPtrTI = nullptr;
   const LoadableTypeInfo *EmptyTI = nullptr;
   const LoadableTypeInfo *IntegerLiteralTI = nullptr;
@@ -181,6 +184,9 @@ class TypeConverter {
   const LoadableTypeInfo &getBridgeObjectTypeInfo();
   const LoadableTypeInfo &getRawPointerTypeInfo();
   const TypeInfo &getTypeMetadataPtrTypeInfo();
+  const TypeInfo &getSwiftContextPtrTypeInfo();
+  const TypeInfo &getTaskContinuationFunctionPtrTypeInfo();
+  const TypeInfo &getSwiftExecutorPtrTypeInfo();
   const TypeInfo &getObjCClassPtrTypeInfo();
   const LoadableTypeInfo &getWitnessTablePtrTypeInfo();
   const LoadableTypeInfo &getEmptyTypeInfo();

lib/IRGen/IRGenModule.cpp

@@ -594,6 +594,14 @@ IRGenModule::IRGenModule(IRGenerator &irgen,
   SwiftTaskPtrTy = SwiftTaskTy->getPointerTo(DefaultAS);
   SwiftExecutorPtrTy = SwiftExecutorTy->getPointerTo(DefaultAS);
 
+  // using TaskContinuationFunction =
+  //   SWIFT_CC(swift)
+  //   void (AsyncTask *, ExecutorRef, AsyncContext *);
+  TaskContinuationFunctionTy = llvm::FunctionType::get(
+      VoidTy, {SwiftTaskPtrTy, SwiftExecutorPtrTy, SwiftContextPtrTy},
+      /*isVarArg*/ false);
+  TaskContinuationFunctionPtrTy = TaskContinuationFunctionTy->getPointerTo();
+
   DifferentiabilityWitnessTy = createStructType(
       *this, "swift.differentiability_witness", {Int8PtrTy, Int8PtrTy});
 }

lib/IRGen/IRGenModule.h

@@ -729,6 +729,8 @@ class IRGenModule {
   llvm::PointerType *SwiftContextPtrTy;
   llvm::PointerType *SwiftTaskPtrTy;
   llvm::PointerType *SwiftExecutorPtrTy;
+  llvm::FunctionType *TaskContinuationFunctionTy;
+  llvm::PointerType *TaskContinuationFunctionPtrTy;
 
   llvm::StructType *DifferentiabilityWitnessTy; // { i8*, i8* }
 
@@ -891,6 +893,9 @@ class IRGenModule {
   const TypeInfo &getTypeInfo(SILType T);
   const TypeInfo &getWitnessTablePtrTypeInfo();
   const TypeInfo &getTypeMetadataPtrTypeInfo();
+  const TypeInfo &getSwiftContextPtrTypeInfo();
+  const TypeInfo &getTaskContinuationFunctionPtrTypeInfo();
+  const TypeInfo &getSwiftExecutorPtrTypeInfo();
   const TypeInfo &getObjCClassPtrTypeInfo();
   const LoadableTypeInfo &getOpaqueStorageTypeInfo(Size size, Alignment align);
   const LoadableTypeInfo &