Change _wait(_throwing) ABIs to reduce code size

Changes the task, taskGroup, asyncLet wait funtion call ABIs.

To reduce code size pass the context parameters and resumption function
as arguments to the wait function.

This means that the suspend point does not need to store parent context
and resumption to the suspend point's context.

```
  void swift_task_future_wait_throwing(
    OpaqueValue * result,
    SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
    AsyncTask *task,
    ThrowingTaskFutureWaitContinuationFunction *resume,
    AsyncContext *callContext);
```

The runtime passes the caller context to the resume entry point saving
the load of the parent context in the resumption function.

This patch adds a `Metadata *` field to `GroupImpl`. The await entry
pointer no longer pass the metadata pointer and there is a path through
the runtime where the task future is no longer available.

Author: aschwaighofer

include/swift/ABI/Executor.h

@@ -128,6 +128,11 @@ using TaskContinuationFunction =
   SWIFT_CC(swiftasync)
   void (SWIFT_ASYNC_CONTEXT AsyncContext *);
 
+using ThrowingTaskFutureWaitContinuationFunction =
+  SWIFT_CC(swiftasync)
+  void (SWIFT_ASYNC_CONTEXT AsyncContext *, SWIFT_CONTEXT void *);
+
+
 template <class AsyncSignature>
 class AsyncFunctionPointer;
 template <class AsyncSignature>

include/swift/Runtime/Concurrency.h

@@ -125,12 +125,6 @@ swift_task_escalate(AsyncTask *task, JobPriority newPriority);
 // TODO: "async let wait" and "async let destroy" would be expressed
 //       similar to like TaskFutureWait;
 
-/// This matches the ABI of a closure `<T>(Builtin.NativeObject) async -> T`
-using TaskFutureWaitSignature =
-    SWIFT_CC(swiftasync)
-    void(OpaqueValue *,
-         SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *, Metadata *);
-
 /// Wait for a non-throwing future task to complete.
 ///
 /// This can be called from any thread. Its Swift signature is
@@ -141,12 +135,9 @@ using TaskFutureWaitSignature =
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
 void swift_task_future_wait(OpaqueValue *,
-         SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *, Metadata *);
-
-using TaskFutureWaitThrowingSignature =
-    SWIFT_CC(swiftasync)
-    void(OpaqueValue *,
-         SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *, Metadata *);
+         SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *,
+         TaskContinuationFunction *,
+         AsyncContext *);
 
 /// Wait for a potentially-throwing future task to complete.
 ///
@@ -157,15 +148,12 @@ using TaskFutureWaitThrowingSignature =
 ///    async throws -> Success
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
-void swift_task_future_wait_throwing(OpaqueValue *,
-                                     SWIFT_ASYNC_CONTEXT AsyncContext *,
-                                     AsyncTask *, Metadata *);
-
-using TaskGroupFutureWaitThrowingSignature =
-SWIFT_CC(swiftasync)
-  void(OpaqueValue *,
-       SWIFT_ASYNC_CONTEXT AsyncContext *, AsyncTask *, TaskGroup *,
-       const Metadata *successType);
+void swift_task_future_wait_throwing(
+  OpaqueValue *,
+  SWIFT_ASYNC_CONTEXT AsyncContext *,
+  AsyncTask *,
+  ThrowingTaskFutureWaitContinuationFunction *,
+  AsyncContext *);
 
 /// Wait for a readyQueue of a Channel to become non empty.
 ///
@@ -180,8 +168,9 @@ SWIFT_CC(swiftasync)
 SWIFT_EXPORT_FROM(swift_Concurrency)
 SWIFT_CC(swiftasync)
 void swift_taskGroup_wait_next_throwing(
-    OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *rawContext,
-    TaskGroup *group, const Metadata *successType);
+    OpaqueValue *resultPointer, SWIFT_ASYNC_CONTEXT AsyncContext *callerContext,
+    TaskGroup *group, ThrowingTaskFutureWaitContinuationFunction *resumeFn,
+    AsyncContext *callContext);
 
 /// Initialize a `TaskGroup` in the passed `group` memory location.
 /// The caller is responsible for retaining and managing the group's lifecycle.
@@ -192,7 +181,7 @@ void swift_taskGroup_wait_next_throwing(
 /// func swift_taskGroup_initialize(group: Builtin.RawPointer)
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-void swift_taskGroup_initialize(TaskGroup *group);
+void swift_taskGroup_initialize(TaskGroup *group, const Metadata *T);
 
 /// Attach a child task to the parent task's task group record.
 ///
@@ -310,7 +299,8 @@ using AsyncLetWaitSignature =
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
 void swift_asyncLet_wait(OpaqueValue *,
                          SWIFT_ASYNC_CONTEXT AsyncContext *,
-                         AsyncLet *, Metadata *);
+                         AsyncLet *, TaskContinuationFunction *,
+                         AsyncContext *);
 
 /// Wait for a potentially-throwing async-let to complete.
 ///
@@ -324,7 +314,9 @@ void swift_asyncLet_wait(OpaqueValue *,
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swiftasync)
 void swift_asyncLet_wait_throwing(OpaqueValue *,
                                   SWIFT_ASYNC_CONTEXT AsyncContext *,
-                                  AsyncLet *, Metadata *);
+                                  AsyncLet *,
+                                  ThrowingTaskFutureWaitContinuationFunction *,
+                                  AsyncContext *);
 
 /// Its Swift signature is
 ///

include/swift/Runtime/RuntimeFunctions.def

@@ -1744,7 +1744,7 @@ FUNCTION(TaskGroupInitialize,
          swift_taskGroup_initialize, SwiftCC,
          ConcurrencyAvailability,
          RETURNS(VoidTy),
-         ARGS(Int8PtrTy),
+         ARGS(Int8PtrTy, TypeMetadataPtrTy),
          ATTRS(NoUnwind))
 
 // void swift_taskGroup_destroy(TaskGroup *group);

lib/AST/Builtins.cpp

@@ -1437,7 +1437,7 @@ static ValueDecl *getCreateAsyncTaskFuture(ASTContext &ctx, Identifier id) {
 
 static ValueDecl *getCreateAsyncTaskGroupFuture(ASTContext &ctx, Identifier id) {
   BuiltinFunctionBuilder builder(ctx);
-  auto genericParam = makeGenericParam().build(builder);
+  auto genericParam = makeGenericParam().build(builder); // <T>
   builder.addParameter(makeConcrete(ctx.getIntType())); // flags
   builder.addParameter(
       makeConcrete(OptionalType::get(ctx.TheRawPointerType))); // group
@@ -1520,7 +1520,8 @@ static ValueDecl *getEndAsyncLet(ASTContext &ctx, Identifier id) {
 
 static ValueDecl *getCreateTaskGroup(ASTContext &ctx, Identifier id) {
   return getBuiltinFunction(ctx, id, _thin,
-                            _parameters(),
+                            _generics(_unrestricted),
+                            _parameters(_metatype(_typeparam(0))),
                             _rawPointer);
 }
 

lib/IRGen/CallEmission.h

@@ -119,6 +119,9 @@ class CallEmission {
 
     return result;
   }
+
+  virtual llvm::Value *getResumeFunctionPointer() = 0;
+  virtual llvm::Value *getAsyncContext() = 0;
 };
 
 std::unique_ptr<CallEmission>

lib/IRGen/Callee.h

@@ -209,24 +209,16 @@ namespace irgen {
       ///
       /// This is a micro-optimization we apply to certain special functions
       /// that we know don't need generics.
-      bool suppressGenerics() const {
+      bool useSpecialConvention() const {
         if (!isSpecial()) return false;
 
         switch (getSpecialKind()) {
-        case SpecialKind::TaskFutureWait:
+        case SpecialKind::AsyncLetWaitThrowing:
         case SpecialKind::TaskFutureWaitThrowing:
+        case SpecialKind::TaskFutureWait:
         case SpecialKind::AsyncLetWait:
-        case SpecialKind::AsyncLetWaitThrowing:
         case SpecialKind::TaskGroupWaitNext:
-          // FIXME: I have disabled this optimization, if we bring it back we
-          // need to debug why it currently does not work (call emission
-          // computes an undef return pointer) and change the runtime entries to
-          // remove the extra type parameter.
-          //
-          // We suppress generics from these as a code-size optimization
-          // because the runtime can recover the success type from the
-          // future.
-          return false;
+          return true;
         }
         llvm_unreachable("covered switch");
       }
@@ -372,9 +364,7 @@ namespace irgen {
       return !kind.isAsyncFunctionPointer();
     }
 
-    bool suppressGenerics() const {
-      return kind.suppressGenerics();
-    }
+    bool useSpecialConvention() const { return kind.useSpecialConvention(); }
   };
 
   class Callee {
@@ -438,9 +428,7 @@ namespace irgen {
       return Fn.getSignature();
     }
 
-    bool suppressGenerics() const {
-      return Fn.suppressGenerics();
-    }
+    bool useSpecialConvention() const { return Fn.useSpecialConvention(); }
 
     /// If this callee has a value for the Swift context slot, return
     /// it; otherwise return non-null.
@@ -458,6 +446,7 @@ namespace irgen {
     llvm::Value *getObjCMethodSelector() const;
   };
 
+  FunctionPointer::Kind classifyFunctionPointerKind(SILFunction *fn);
 } // end namespace irgen
 } // end namespace swift
 

lib/IRGen/GenBuiltin.cpp

@@ -250,7 +250,9 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
   }
 
   if (Builtin.ID == BuiltinValueKind::CreateTaskGroup) {
-    out.add(emitCreateTaskGroup(IGF));
+    // Claim metadata pointer.
+    (void)args.claimAll();
+    out.add(emitCreateTaskGroup(IGF, substitutions));
     return;
   }