IRGen: Fix lowering of Builtin.createAsyncTask and Builtin.createAsyncTaskFuture

Thick async functions store their async context size in the closure
context. Only if the closure context is nil can we assume the
partial_apply_forwarder function to be the address of an async function
pointer struct value.

Author: aschwaighofer

include/swift/Runtime/RuntimeFunctions.def

@@ -1498,23 +1498,24 @@ FUNCTION(TaskCancel,
          ATTRS(NoUnwind, ArgMemOnly))
 
 // AsyncTaskAndContext swift_task_create(
-//     size_t flags, AsyncTask *task, AsyncFunctionPointer *function);
+//     size_t flags, AsyncTask *task, TaskContinuationFunction* function,
+//     size_t contextSize);
 FUNCTION(TaskCreateFunc,
-         swift_task_create, SwiftCC,
+         swift_task_create_f, SwiftCC,
          ConcurrencyAvailability,
          RETURNS(AsyncTaskAndContextTy),
-         ARGS(SizeTy, SwiftTaskPtrTy, AsyncFunctionPointerPtrTy),
+         ARGS(SizeTy, SwiftTaskPtrTy, TaskContinuationFunctionPtrTy, SizeTy),
          ATTRS(NoUnwind, ArgMemOnly))
 
-// AsyncTaskAndContext swift_task_create_future(
+// AsyncTaskAndContext swift_task_create_future_f(
 //     size_t flags, AsyncTask *task, const Metadata *futureResultType,
-//     AsyncFunctionPointer *function);
+//     TaskContinuationFunction *function, size_t contextSize);
 FUNCTION(TaskCreateFutureFunc,
-         swift_task_create_future, SwiftCC,
+         swift_task_create_future_f, SwiftCC,
          ConcurrencyAvailability,
          RETURNS(AsyncTaskAndContextTy),
          ARGS(SizeTy, SwiftTaskPtrTy, TypeMetadataPtrTy,
-              AsyncFunctionPointerPtrTy),
+              TaskContinuationFunctionPtrTy, SizeTy),
          ATTRS(NoUnwind, ArgMemOnly))
 
 #undef RETURNS

lib/IRGen/GenCall.cpp

@@ -3600,8 +3600,6 @@ llvm::Value *irgen::emitTaskCreate(
     SubstitutionMap subs) {
   parentTask = IGF.Builder.CreateBitOrPointerCast(
       parentTask, IGF.IGM.SwiftTaskPtrTy);
-  taskFunction = IGF.Builder.CreateBitOrPointerCast(
-      taskFunction, IGF.IGM.AsyncFunctionPointerPtrTy);
 
   // Determine the size of the async context for the closure.
   ASTContext &ctx = IGF.IGM.IRGen.SIL.getASTContext();
@@ -3625,29 +3623,36 @@ llvm::Value *irgen::emitTaskCreate(
 
   // Call the function.
   llvm::CallInst *result;
+  llvm::Value *theSize, *theFunction;
+  std::tie(theFunction, theSize) =
+      getAsyncFunctionAndSize(IGF, SILFunctionTypeRepresentation::Thick,
+                              FunctionPointer::forExplosionValue(
+                                  IGF, taskFunction, taskFunctionCanSILType),
+                              localContextInfo);
+  theFunction = IGF.Builder.CreateBitOrPointerCast(
+      theFunction, IGF.IGM.TaskContinuationFunctionPtrTy);
+  theSize = IGF.Builder.CreateZExtOrBitCast(theSize, IGF.IGM.SizeTy);
   if (futureResultType) {
     result = IGF.Builder.CreateCall(
       IGF.IGM.getTaskCreateFutureFuncFn(),
-      { flags, parentTask, futureResultType, taskFunction });
+      { flags, parentTask, futureResultType, theFunction, theSize });
   } else {
-    result = IGF.Builder.CreateCall(
-      IGF.IGM.getTaskCreateFuncFn(),
-      { flags, parentTask, taskFunction });
+    result = IGF.Builder.CreateCall(IGF.IGM.getTaskCreateFuncFn(),
+                                    {flags, parentTask, theFunction, theSize});
   }
   result->setDoesNotThrow();
   result->setCallingConv(IGF.IGM.SwiftCC);
 
   // Write the local context information into the initial context for the task.
-  if (layout.hasLocalContext()) {
-    // Dig out the initial context returned from task creation.
-    auto initialContext = IGF.Builder.CreateExtractValue(result, { 1 });
-    Address initialContextAddr = layout.emitCastTo(IGF, initialContext);
-
-    auto localContextLayout = layout.getLocalContextLayout();
-    auto localContextAddr = localContextLayout.project(
-        IGF, initialContextAddr, llvm::None);
-    IGF.Builder.CreateStore(localContextInfo, localContextAddr);
-  }
+  assert(layout.hasLocalContext());
+  // Dig out the initial context returned from task creation.
+  auto initialContext = IGF.Builder.CreateExtractValue(result, {1});
+  Address initialContextAddr = layout.emitCastTo(IGF, initialContext);
+
+  auto localContextLayout = layout.getLocalContextLayout();
+  auto localContextAddr =
+      localContextLayout.project(IGF, initialContextAddr, llvm::None);
+  IGF.Builder.CreateStore(localContextInfo, localContextAddr);
 
   return result;
 }

test/IRGen/async/builtins.sil

@@ -34,7 +34,7 @@ sil hidden [ossa] @launch_task : $@convention(method) @async (Int, Optional<Buil
 bb0(%0 : $Int, %1: @unowned $Optional<Builtin.NativeObject>, %2: @guaranteed $@async @callee_guaranteed () -> (@error Error)):
   %3 = begin_borrow %1 : $Optional<Builtin.NativeObject>
   // CHECK: call %swift.refcounted* @swift_retain(%swift.refcounted* returned [[FN_CONTEXT:%.*]])
-  // CHECK: [[NEW_TASK_AND_CONTEXT:%.*]] = call swiftcc %swift.async_task_and_context @swift_task_create(
+  // CHECK: [[NEW_TASK_AND_CONTEXT:%.*]] = call swiftcc %swift.async_task_and_context @swift_task_create_f(
   // CHECK-NEXT: [[NEW_CONTEXT_RAW:%.*]] = extractvalue %swift.async_task_and_context [[NEW_TASK_AND_CONTEXT]], 1
   // CHECK-NEXT: [[NEW_CONTEXT:%.*]] = bitcast %swift.context* [[NEW_CONTEXT_RAW]] to
   // CHECK-NEXT: [[CONTEXT_INFO_LOC:%.*]] = getelementptr inbounds <{{.*}}>* [[NEW_CONTEXT]]
@@ -53,7 +53,7 @@ bb0(%0 : $Int, %1: @unowned $Optional<Builtin.NativeObject>, %2: @guaranteed $@a
   // CHECK: call %swift.refcounted* @swift_retain(%swift.refcounted* returned [[FN_CONTEXT:%.*]])
   %9 = metatype $@thick T.Type
   %10 = init_existential_metatype %9 : $@thick T.Type, $@thick Any.Type
-  // CHECK: [[NEW_TASK_AND_CONTEXT:%.*]] = call swiftcc %swift.async_task_and_context @swift_task_create_future(
+  // CHECK: [[NEW_TASK_AND_CONTEXT:%.*]] = call swiftcc %swift.async_task_and_context @swift_task_create_future_f(
   // CHECK-NEXT: [[NEW_CONTEXT_RAW:%.*]] = extractvalue %swift.async_task_and_context [[NEW_TASK_AND_CONTEXT]], 1
   // CHECK-NEXT: [[NEW_CONTEXT:%.*]] = bitcast %swift.context* [[NEW_CONTEXT_RAW]] to
   // CHECK-NEXT: [[CONTEXT_INFO_LOC:%.*]] = getelementptr inbounds <{{.*}}>* [[NEW_CONTEXT]]