SILGen: Caller-side codegen for invoking foreign async functions

jckarter · jckarter · commit e7ec8c35af3c · 2020-11-03T08:28:30.000-08:00
Immediately before invoking the ObjC API, get the current continuation, capture it into a block to
pass as the completion handler, and then await the continuation, whose resume/error successors
serve as the semantic return/throw result of the call. This should complete the caller-side part
of SILGen; the completion handler block implementation is however still only a stub.
diff --git a/docs/ABI/Mangling.rst b/docs/ABI/Mangling.rst
@@ -219,7 +219,7 @@ types where the metadata itself has unknown layout.)
   global ::= entity                      // some identifiable thing
   global ::= from-type to-type generic-signature? 'TR'  // reabstraction thunk
   global ::= from-type to-type generic-signature? 'TR'  // reabstraction thunk
-  global ::= impl-function-type 'Tz'     // objc-to-swift-async completion handler block implementation
+  global ::= impl-function-type type 'Tz'     // objc-to-swift-async completion handler block implementation
   global ::= from-type to-type self-type generic-signature? 'Ty'  // reabstraction thunk with dynamic 'Self' capture
   global ::= from-type to-type generic-signature? 'Tr'  // obsolete mangling for reabstraction thunk
   global ::= entity generic-signature? type type* 'TK' // key path getter
diff --git a/include/swift/AST/ASTMangler.h b/include/swift/AST/ASTMangler.h
@@ -155,6 +155,9 @@ class ASTMangler : public Mangler {
                                              Type SelfType,
                                              ModuleDecl *Module);
 
+  std::string mangleObjCAsyncCompletionHandlerImpl(CanSILFunctionType BlockType,
+                                                   CanType ResultType);
+  
   /// Mangle the derivative function (JVP/VJP) for the given:
   /// - Mangled original function name.
   /// - Derivative function kind.
diff --git a/lib/AST/ASTMangler.cpp b/lib/AST/ASTMangler.cpp
@@ -394,6 +394,16 @@ std::string ASTMangler::mangleReabstractionThunkHelper(
   return finalize();
 }
 
+std::string ASTMangler::mangleObjCAsyncCompletionHandlerImpl(
+                                                   CanSILFunctionType BlockType,
+                                                   CanType ResultType) {
+  beginMangling();
+  appendType(BlockType);
+  appendType(ResultType);
+  appendOperator("Tz");
+  return finalize();
+}
+
 std::string ASTMangler::mangleAutoDiffDerivativeFunctionHelper(
     StringRef name, AutoDiffDerivativeFunctionKind kind,
     AutoDiffConfig config) {
diff --git a/lib/Demangling/Demangler.cpp b/lib/Demangling/Demangler.cpp
@@ -2264,8 +2264,10 @@ NodePointer Demangler::popProtocolConformance() {
       return createWithChild(Node::Kind::CoroutineContinuationPrototype, type);
     }
     case 'z': {
-      NodePointer implType = popNode(Node::Kind::ImplFunctionType);
-      return createWithChild(Node::Kind::ObjCAsyncCompletionHandlerImpl, implType);
+      NodePointer resultType = popNode(Node::Kind::Type);
+      NodePointer implType = popNode(Node::Kind::Type);
+      return createWithChildren(Node::Kind::ObjCAsyncCompletionHandlerImpl,
+                                implType, resultType);
     }
     case 'V': {
       NodePointer Base = popNode(isEntity);
diff --git a/lib/Demangling/NodePrinter.cpp b/lib/Demangling/NodePrinter.cpp
@@ -2536,6 +2536,8 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
   case Node::Kind::ObjCAsyncCompletionHandlerImpl:
     Printer << "@objc completion handler block implementation for ";
     print(Node->getChild(0));
+    Printer << " with result type ";
+    print(Node->getChild(1));
     return nullptr;
   }
   printer_unreachable("bad node kind!");
diff --git a/lib/SILGen/ResultPlan.cpp b/lib/SILGen/ResultPlan.cpp
@@ -453,25 +453,115 @@ class TupleInitializationResultPlan final : public ResultPlan {
 
 class ForeignAsyncInitializationPlan final : public ResultPlan {
   SILLocation loc;
+  CalleeTypeInfo calleeTypeInfo;
+  SILType opaqueResumeType;
+  SILValue resumeBuf;
+  SILValue continuation;
+  
 public:
-  ForeignAsyncInitializationPlan(SILLocation loc) : loc(loc) {}
+  ForeignAsyncInitializationPlan(SILGenFunction &SGF, SILLocation loc,
+                                 const CalleeTypeInfo &calleeTypeInfo)
+    : loc(loc), calleeTypeInfo(calleeTypeInfo)
+  {
+    // Allocate space to receive the resume value when the continuation is
+    // resumed.
+    opaqueResumeType = SGF.getLoweredType(AbstractionPattern::getOpaque(),
+                                          calleeTypeInfo.substResultType);
+    resumeBuf = SGF.emitTemporaryAllocation(loc, opaqueResumeType);
+  }
   
   void
   gatherIndirectResultAddrs(SILGenFunction &SGF, SILLocation loc,
                             SmallVectorImpl<SILValue> &outList) const override {
-    // TODO: Move values from the continuation result buffer to the individual
-    // out argument buffers, unless we were able to emit the resume buffer
-    // in-place.
+    // A foreign async function shouldn't have any indirect results.
+  }
+  
+  ManagedValue
+  emitForeignAsyncCompletionHandler(SILGenFunction &SGF, SILLocation loc)
+  override {
+    // Get the current continuation for the task.
+    auto continuationDecl = calleeTypeInfo.foreign.async->completionHandlerErrorParamIndex()
+      ? SGF.getASTContext().getUnsafeThrowingContinuationDecl()
+      : SGF.getASTContext().getUnsafeContinuationDecl();
+    
+    auto continuationTy = BoundGenericType::get(continuationDecl, Type(),
+                                                calleeTypeInfo.substResultType)
+      ->getCanonicalType();
+    
+    
+    continuation = SGF.B.createGetAsyncContinuationAddr(loc, resumeBuf,
+                               SILType::getPrimitiveObjectType(continuationTy));
+    
+    // Stash it in a buffer for a block object.
+    auto blockStorageTy = SILType::getPrimitiveAddressType(SILBlockStorageType::get(continuationTy));
+    auto blockStorage = SGF.emitTemporaryAllocation(loc, blockStorageTy);
+    auto continuationAddr = SGF.B.createProjectBlockStorage(loc, blockStorage);
+    SGF.B.createStore(loc, continuation, continuationAddr,
+                      StoreOwnershipQualifier::Trivial);
+
+    // Get the block invocation function for the given completion block type.
+    auto completionHandlerIndex = calleeTypeInfo.foreign.async
+      ->completionHandlerParamIndex();
+    auto implTy = cast<SILFunctionType>(calleeTypeInfo.substFnType
+      ->getParameters()[completionHandlerIndex]
+      .getInterfaceType());
+    SILFunction *impl = SGF.SGM
+      .getOrCreateForeignAsyncCompletionHandlerImplFunction(implTy,
+                                                            continuationTy);
+    auto implRef = SGF.B.createFunctionRef(loc, impl);
+    
+    // Initialize the block object for the completion handler.
+    auto block = SGF.B.createInitBlockStorageHeader(loc, blockStorage, implRef,
+                                  SILType::getPrimitiveObjectType(implTy), {});
+    // We don't need to manage the block because it's still on the stack. We
+    // know we won't escape it locally so the callee can be responsible for
+    // _Block_copy-ing it.
+    return ManagedValue::forUnmanaged(block);
   }
   
   RValue finish(SILGenFunction &SGF, SILLocation loc, CanType substType,
                 ArrayRef<ManagedValue> &directResults) override {
     // There should be no direct results from the call.
     assert(directResults.empty());
     
-    // TODO: Get the actual result values from the awaited continuation.
-    // For now, produce an undef RValue.
-    return SGF.emitUndefRValue(loc, substType);
+    // Await the continuation we handed off to the completion handler.
+    SILBasicBlock *resumeBlock = SGF.createBasicBlock();
+    SILBasicBlock *errorBlock = nullptr;
+    auto errorParamIndex = calleeTypeInfo.foreign.async->completionHandlerErrorParamIndex();
+    if (errorParamIndex) {
+      errorBlock = SGF.createBasicBlock(FunctionSection::Postmatter);
+    }
+    
+    SGF.B.createAwaitAsyncContinuation(loc, continuation, resumeBlock, errorBlock);
+    
+    // Propagate an error if we have one.
+    if (errorBlock) {
+      SGF.B.emitBlock(errorBlock);
+      
+      Scope errorScope(SGF, loc);
+      
+      auto errorTy = SGF.getASTContext().getErrorDecl()->getDeclaredType()
+        ->getCanonicalType();
+      auto errorVal
+        = SGF.B.createOwnedPhiArgument(SILType::getPrimitiveObjectType(errorTy));
+      
+      SGF.emitThrow(loc, errorVal, true);
+    }
+    
+    SGF.B.emitBlock(resumeBlock);
+    
+    // The incoming value is the maximally-abstracted result type of the
+    // continuation. Move it out of the resume buffer and reabstract it if
+    // necessary.
+    auto resumeResult = SGF.emitLoad(loc, resumeBuf,
+      calleeTypeInfo.origResultType
+         ? *calleeTypeInfo.origResultType
+         : AbstractionPattern(calleeTypeInfo.substResultType),
+                 calleeTypeInfo.substResultType,
+                 SGF.getTypeLowering(calleeTypeInfo.substResultType),
+                 SGFContext(), IsTake);
+    
+    return RValue(SGF, loc, calleeTypeInfo.substResultType, resumeResult);
   }
 };
 
@@ -572,8 +662,7 @@ ResultPlanPtr ResultPlanBuilder::buildTopLevelResult(Initialization *init,
     // Create a result plan that gets the result schema from the completion
     // handler callback's arguments.
     // completion handler.
-    return ResultPlanPtr(new ForeignAsyncInitializationPlan(loc));
-    
+    return ResultPlanPtr(new ForeignAsyncInitializationPlan(SGF, loc, calleeTypeInfo));
   } else if (auto foreignError = calleeTypeInfo.foreign.error) {
     // Handle the foreign error first.
     //
diff --git a/lib/SILGen/ResultPlan.h b/lib/SILGen/ResultPlan.h
@@ -49,6 +49,11 @@ class ResultPlan {
   emitForeignErrorArgument(SILGenFunction &SGF, SILLocation loc) {
     return None;
   }
+  
+  virtual ManagedValue
+  emitForeignAsyncCompletionHandler(SILGenFunction &SGF, SILLocation loc) {
+    return {};
+  }
 };
 
 using ResultPlanPtr = std::unique_ptr<ResultPlan>;
diff --git a/lib/SILGen/SILGen.h b/lib/SILGen/SILGen.h
@@ -163,6 +163,13 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
                                            CanSILFunctionType fromType,
                                            CanSILFunctionType toType,
                                            CanType dynamicSelfType);
+  
+  /// Get or create the declaration of a completion handler block
+  /// implementation function for an ObjC API that was imported
+  /// as `async` in Swift.
+  SILFunction *getOrCreateForeignAsyncCompletionHandlerImplFunction(
+                                           CanSILFunctionType blockType,
+                                           CanType continuationTy);
 
   /// Determine whether the given class has any instance variables that
   /// need to be destroyed.
diff --git a/lib/SILGen/SILGenApply.cpp b/lib/SILGen/SILGenApply.cpp
@@ -3276,14 +3276,15 @@ class ArgEmitter {
     if (Foreign.async
         && Foreign.async->completionHandlerParamIndex() == Args.size()) {
       SILParameterInfo param = claimNextParameters(1).front();
-
-      // TODO: Get or create the completion handler block implementation
-      // function for the given argument type, then create a block containing
-      // the current continuation. (This probably needs to be deferred to right
-      // before the actual call, since evaluating other arguments to the call
-      // may suspend the task)
-      auto argTy = SILType::getPrimitiveObjectType(param.getInterfaceType());
-      Args.push_back(ManagedValue::forUnmanaged(SILUndef::get(argTy, SGF.F)));
+      (void)param;
+
+      // Leave a placeholder in the position. We'll fill this in with a block
+      // capturing the current continuation right before we invoke the
+      // function.
+      // (We can't do this immediately, because evaluating other arguments
+      // may require suspending the async task, which is not allowed while its
+      // continuation is active.)
+      Args.push_back(ManagedValue::forInContext());
     } else if (Foreign.error
                && Foreign.error->getErrorParameterIndex() == Args.size()) {
       SILParameterInfo param = claimNextParameters(1).front();
@@ -4327,14 +4328,21 @@ RValue SILGenFunction::emitApply(ResultPlanPtr &&resultPlan,
 
   // If there's a foreign error or async parameter, fill it in.
   ManagedValue errorTemp;
-  if (calleeTypeInfo.foreign.async) {
-    // TODO: prepare the callback continuation and block here.
-    
-  } else if (calleeTypeInfo.foreign.error) {
+  if (auto foreignAsync = calleeTypeInfo.foreign.async) {
+    unsigned completionIndex = foreignAsync->completionHandlerParamIndex();
+
+    // Ram the emitted error into the argument list, over the placeholder
+    // we left during the first pass.
+    auto &completionArgSlot = const_cast<ManagedValue &>(args[completionIndex]);
+
+    completionArgSlot = resultPlan->emitForeignAsyncCompletionHandler(*this, loc);
+
+  } else if (auto foreignError = calleeTypeInfo.foreign.error) {
     unsigned errorParamIndex =
-        calleeTypeInfo.foreign.error->getErrorParameterIndex();
+        foreignError->getErrorParameterIndex();
 
-    // This is pretty evil.
+    // Ram the emitted error into the argument list, over the placeholder
+    // we left during the first pass.
     auto &errorArgSlot = const_cast<ManagedValue &>(args[errorParamIndex]);
 
     std::tie(errorTemp, errorArgSlot) =
@@ -4442,9 +4450,6 @@ RValue SILGenFunction::emitApply(ResultPlanPtr &&resultPlan,
                           *foreignError);
   }
   
-  // TODO(async): If there's a foreign async convention, await the continuation
-  // to get the result from the completion callback.
-
   auto directResultsArray = makeArrayRef(directResults);
   RValue result =
     resultPlan->finish(*this, loc, substResultType, directResultsArray);
diff --git a/lib/SILGen/SILGenThunk.cpp b/lib/SILGen/SILGenThunk.cpp
@@ -138,6 +138,60 @@ SILGenFunction::emitGlobalFunctionRef(SILLocation loc, SILDeclRef constant,
     return B.createFunctionRefFor(loc, f);
 }
 
+SILFunction *
+SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
+                                         CanSILFunctionType blockType,
+                                         CanType continuationTy) {
+  // Extract the result type from the continuation type.
+  auto resumeType = cast<BoundGenericType>(continuationTy).getGenericArgs()[0];
+  
+  // Build up the implementation function type, which matches the
+  // block signature with an added block storage argument that points at the
+  // block buffer. The block storage holds the continuation we feed the
+  // result values into.
+  SmallVector<SILParameterInfo, 4> implArgs;
+  auto blockStorageTy = SILBlockStorageType::get(continuationTy);
+  implArgs.push_back(SILParameterInfo(blockStorageTy,
+                                ParameterConvention::Indirect_InoutAliasable));
+  
+  std::copy(blockType->getParameters().begin(),
+            blockType->getParameters().end(),
+            std::back_inserter(implArgs));
+  
+  auto implTy = SILFunctionType::get(GenericSignature(),
+         blockType->getExtInfo()
+           .withRepresentation(SILFunctionTypeRepresentation::CFunctionPointer),
+         SILCoroutineKind::None,
+         ParameterConvention::Direct_Unowned,
+         implArgs, {}, blockType->getResults(),
+         None,
+         SubstitutionMap(), SubstitutionMap(), getASTContext());
+
+  auto loc = RegularLocation::getAutoGeneratedLocation();
+
+  Mangle::ASTMangler Mangler;
+  auto name = Mangler.mangleObjCAsyncCompletionHandlerImpl(blockType,
+                                                           resumeType);
+  
+  SILGenFunctionBuilder builder(*this);
+  auto F = builder.getOrCreateSharedFunction(loc, name, implTy,
+                                           IsBare, IsTransparent, IsSerializable,
+                                           ProfileCounter(),
+                                           IsThunk,
+                                           IsNotDynamic);
+  
+  if (F->empty()) {
+    // TODO: Emit the implementation.
+    SILGenFunction SGF(*this, *F, SwiftModule);
+    SmallVector<ManagedValue, 4> params;
+    SGF.collectThunkParams(loc, params);
+
+    SGF.B.createUnreachable(loc);
+  }
+  
+  return F;
+}
+
 SILFunction *SILGenModule::
 getOrCreateReabstractionThunk(CanSILFunctionType thunkType,
                               CanSILFunctionType fromType,
diff --git a/test/SILGen/objc_async.swift b/test/SILGen/objc_async.swift
@@ -6,9 +6,39 @@ import ObjCConcurrency
 
 // CHECK-LABEL: sil {{.*}}@${{.*}}14testSlowServer
 func testSlowServer(slowServer: SlowServer) async throws {
-  // CHECK: objc_method {{.*}} $@convention(objc_method) (NSString, @convention(block) (Int) -> (), SlowServer) -> ()
+  // CHECK: [[RESUME_BUF:%.*]] = alloc_stack $Int
+  // CHECK: [[METHOD:%.*]] = objc_method {{.*}} $@convention(objc_method) (NSString, @convention(block) (Int) -> (), SlowServer) -> ()
+  // CHECK: [[CONT:%.*]] = get_async_continuation_addr $Int, [[RESUME_BUF]]
+  // CHECK: [[BLOCK_STORAGE:%.*]] = alloc_stack $@block_storage UnsafeContinuation<Int>
+  // CHECK: [[CONT_SLOT:%.*]] = project_block_storage [[BLOCK_STORAGE]]
+  // CHECK: store [[CONT]] to [[CONT_SLOT]]
+  // CHECK: [[BLOCK_IMPL:%.*]] = %20 = function_ref @{{.*}} : $@convention(c) (@inout_aliasable @block_storage UnsafeContinuation<Int>, Int) -> ()
+  // CHECK: [[BLOCK:%.*]] = init_block_storage_header [[BLOCK_STORAGE]] {{.*}}, invoke [[BLOCK_IMPL]]
+  // CHECK: apply [[METHOD]]({{.*}}, [[BLOCK]], %0)
+  // CHECK: await_async_continuation [[CONT]] {{.*}}, resume [[RESUME:bb[0-9]+]]
+  // CHECK: [[RESUME]]:
+  // CHECK: [[RESULT:%.*]] = load [trivial] [[RESUME_BUF]]
+  // CHECK: dealloc_stack [[RESUME_BUF]]
   let _: Int = await slowServer.doSomethingSlow("mail")
-  // CHECK: objc_method {{.*}} $@convention(objc_method) (@convention(block) (Optional<NSString>, Optional<NSError>) -> (), SlowServer) -> ()
+
+  // CHECK: [[RESUME_BUF:%.*]] = alloc_stack $Optional<String>
+  // CHECK: [[METHOD:%.*]] = objc_method {{.*}} $@convention(objc_method) (@convention(block) (Optional<NSString>, Optional<NSError>) -> (), SlowServer) -> ()
+  // CHECK: [[CONT:%.*]] = get_async_continuation_addr $Optional<String>, [[RESUME_BUF]]
+  // CHECK: [[BLOCK_STORAGE:%.*]] = alloc_stack $@block_storage UnsafeContinuation<Optional<String>>
+  // CHECK: [[CONT_SLOT:%.*]] = project_block_storage [[BLOCK_STORAGE]]
+  // CHECK: store [[CONT]] to [[CONT_SLOT]]
+  // CHECK: [[BLOCK_IMPL:%.*]] = %20 = function_ref @{{.*}} : $@convention(c) (@inout_aliasable @block_storage UnsafeContinuation<Optional<String>>, Optional<NSString>, Optional<NSError>) -> ()
+  // CHECK: [[BLOCK:%.*]] = init_block_storage_header [[BLOCK_STORAGE]] {{.*}}, invoke [[BLOCK_IMPL]]
+  // CHECK: apply [[METHOD]]({{.*}}, [[BLOCK]], %0)
+  // CHECK: await_async_continuation [[CONT]] {{.*}}, resume [[RESUME:bb[0-9]+]], error [[ERROR:bb[0-9]+]]
+  // CHECK: [[RESUME]]:
+  // CHECK: [[RESULT:%.*]] = load [take] [[RESUME_BUF]]
+  // CHECK: destroy_value [[RESULT]]
+  // CHECK: dealloc_stack [[RESUME_BUF]]
+  // CHECK: [[ERROR]]([[ERROR_VALUE:%.*]] : @owned $Error):
+  // CHECK: dealloc_stack [[RESUME_BUF]]
+  // CHECK: throw [[ERROR_VALUE]]
+
   let _: String? = try await slowServer.findAnswer()
 
   // CHECK: objc_method {{.*}} $@convention(objc_method) (NSString, @convention(block) () -> (), SlowServer) -> ()

Original file line number	Diff line number	Diff line change
`@@ -2264,8 +2264,10 @@ NodePointer Demangler::popProtocolConformance() {`
`2264`	`2264`	`return createWithChild(Node::Kind::CoroutineContinuationPrototype, type);`
`2265`	`2265`	`}`
`2266`	`2266`	`case 'z': {`
`2267`		`- NodePointer implType = popNode(Node::Kind::ImplFunctionType);`
`2268`		`- return createWithChild(Node::Kind::ObjCAsyncCompletionHandlerImpl, implType);`
	`2267`	`+ NodePointer resultType = popNode(Node::Kind::Type);`
	`2268`	`+ NodePointer implType = popNode(Node::Kind::Type);`
	`2269`	`+ return createWithChildren(Node::Kind::ObjCAsyncCompletionHandlerImpl,`
	`2270`	`+ implType, resultType);`
`2269`	`2271`	`}`
`2270`	`2272`	`case 'V': {`
`2271`	`2273`	`NodePointer Base = popNode(isEntity);`
Original file line number	Diff line number	Diff line change
`@@ -2536,6 +2536,8 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {`
`2536`	`2536`	`case Node::Kind::ObjCAsyncCompletionHandlerImpl:`
`2537`	`2537`	`Printer << "@objc completion handler block implementation for ";`
`2538`	`2538`	`print(Node->getChild(0));`
	`2539`	`+ Printer << " with result type ";`
	`2540`	`+ print(Node->getChild(1));`
`2539`	`2541`	`return nullptr;`
`2540`	`2542`	`}`
`2541`	`2543`	`printer_unreachable("bad node kind!");`