Merge pull request swiftlang#34525 from jckarter/foreign-async-silgen

jckarter · web-flow · commit 0ca29b504f84 · 2020-11-12T06:55:04.000-08:00
SILGen: Caller-side codegen for invoking foreign async functions
diff --git a/docs/ABI/Mangling.rst b/docs/ABI/Mangling.rst
@@ -220,7 +220,8 @@ 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 ::= impl-function-type type 'TZ'     // objc-to-swift-async completion handler block implementation (predefined by runtime)
   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,15 @@ class ASTMangler : public Mangler {
                                              Type SelfType,
                                              ModuleDecl *Module);
 
+  /// Mangle a completion handler block implementation function, used for importing ObjC
+  /// APIs as async.
+  ///
+  /// - If `predefined` is true, this mangles the symbol name of the completion handler
+  /// predefined in the Swift runtime for the given type signature.
+  std::string mangleObjCAsyncCompletionHandlerImpl(CanSILFunctionType BlockType,
+                                                   CanType ResultType,
+                                                   bool predefined);
+  
   /// Mangle the derivative function (JVP/VJP) for the given:
   /// - Mangled original function name.
   /// - Derivative function kind.
diff --git a/include/swift/Demangling/DemangleNodes.def b/include/swift/Demangling/DemangleNodes.def
@@ -160,6 +160,7 @@ NODE(NominalTypeDescriptor)
 NODE(NonObjCAttribute)
 NODE(Number)
 NODE(ObjCAsyncCompletionHandlerImpl)
+NODE(PredefinedObjCAsyncCompletionHandlerImpl)
 NODE(ObjCAttribute)
 NODE(ObjCBlock)
 NODE(EscapingObjCBlock)
diff --git a/lib/AST/ASTMangler.cpp b/lib/AST/ASTMangler.cpp
@@ -394,6 +394,17 @@ std::string ASTMangler::mangleReabstractionThunkHelper(
   return finalize();
 }
 
+std::string ASTMangler::mangleObjCAsyncCompletionHandlerImpl(
+                                                   CanSILFunctionType BlockType,
+                                                   CanType ResultType,
+                                                   bool predefined) {
+  beginMangling();
+  appendType(BlockType);
+  appendType(ResultType);
+  appendOperator(predefined ? "TZ" : "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
@@ -2266,9 +2266,14 @@ NodePointer Demangler::popProtocolConformance() {
       NodePointer type = popNode(Node::Kind::Type);
       return createWithChild(Node::Kind::CoroutineContinuationPrototype, type);
     }
-    case 'z': {
-      NodePointer implType = popNode(Node::Kind::ImplFunctionType);
-      return createWithChild(Node::Kind::ObjCAsyncCompletionHandlerImpl, implType);
+    case 'z':
+    case 'Z': {
+      NodePointer resultType = popNode(Node::Kind::Type);
+      NodePointer implType = popNode(Node::Kind::Type);
+      return createWithChildren(c == 'z'
+                                  ? Node::Kind::ObjCAsyncCompletionHandlerImpl
+                                  : Node::Kind::PredefinedObjCAsyncCompletionHandlerImpl,
+                                implType, resultType);
     }
     case 'V': {
       NodePointer Base = popNode(isEntity);
diff --git a/lib/Demangling/NodePrinter.cpp b/lib/Demangling/NodePrinter.cpp
@@ -455,6 +455,7 @@ class NodePrinter {
     case Node::Kind::PartialApplyForwarder:
     case Node::Kind::PartialApplyObjCForwarder:
     case Node::Kind::PostfixOperator:
+    case Node::Kind::PredefinedObjCAsyncCompletionHandlerImpl:
     case Node::Kind::PrefixOperator:
     case Node::Kind::PrivateDeclName:
     case Node::Kind::PropertyDescriptor:
@@ -2534,9 +2535,14 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
       Printer << ')';
     }
     return nullptr;
+  case Node::Kind::PredefinedObjCAsyncCompletionHandlerImpl:
+    Printer << "predefined ";
+    LLVM_FALLTHROUGH;
   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;
   case Node::Kind::CanonicalPrespecializedGenericTypeCachingOnceToken:
     Printer << "flag for loading of canonical specialized generic type "
diff --git a/lib/Demangling/OldRemangler.cpp b/lib/Demangling/OldRemangler.cpp
@@ -2166,6 +2166,9 @@ void Remangler::mangleGlobalVariableOnceFunction(Node *node) {
 void Remangler::mangleGlobalVariableOnceDeclList(Node *node) {
   unreachable("unsupported");
 }
+void Remangler::manglePredefinedObjCAsyncCompletionHandlerImpl(Node *node) {
+  unreachable("unsupported");
+}
 void Remangler::mangleObjCAsyncCompletionHandlerImpl(Node *node) {
   unreachable("unsupported");
 }
diff --git a/lib/Demangling/Remangler.cpp b/lib/Demangling/Remangler.cpp
@@ -805,6 +805,11 @@ void Remangler::mangleCoroutineContinuationPrototype(Node *node) {
   Buffer << "TC";
 }
 
+void Remangler::manglePredefinedObjCAsyncCompletionHandlerImpl(Node *node) {
+  mangleChildNodes(node);
+  Buffer << "TZ";
+}
+
 void Remangler::mangleObjCAsyncCompletionHandlerImpl(Node *node) {
   mangleChildNodes(node);
   Buffer << "Tz";
diff --git a/lib/SILGen/ResultPlan.cpp b/lib/SILGen/ResultPlan.cpp
@@ -453,25 +453,116 @@ 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,
+                                                *calleeTypeInfo.foreign.async);
+    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 +663,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.cpp b/lib/SILGen/SILGen.cpp
@@ -327,6 +327,53 @@ SILGenModule::getConformanceToBridgedStoredNSError(SILLocation loc, Type type) {
   return SwiftModule->lookupConformance(type, proto);
 }
 
+static FuncDecl *
+lookUpResumeContinuationIntrinsic(ASTContext &C,
+                                  Optional<FuncDecl*> &cache,
+                                  StringRef name) {
+  if (cache)
+    return *cache;
+  
+  auto *module = C.getLoadedModule(C.Id_Concurrency);
+  if (!module) {
+    cache = nullptr;
+    return nullptr;
+  }
+  
+  SmallVector<ValueDecl *, 1> decls;
+  module->lookupQualified(module,
+                     DeclNameRef(C.getIdentifier(name)),
+                     NL_QualifiedDefault | NL_IncludeUsableFromInline,
+                     decls);
+
+  if (decls.size() != 1) {
+    cache = nullptr;
+    return nullptr;
+  }
+  auto func = dyn_cast<FuncDecl>(decls[0]);
+  cache = func;
+  return func;
+}
+
+FuncDecl *
+SILGenModule::getResumeUnsafeContinuation() {
+  return lookUpResumeContinuationIntrinsic(getASTContext(),
+                                           ResumeUnsafeContinuation,
+                                           "_resumeUnsafeContinuation");
+}
+FuncDecl *
+SILGenModule::getResumeUnsafeThrowingContinuation() {
+  return lookUpResumeContinuationIntrinsic(getASTContext(),
+                                           ResumeUnsafeThrowingContinuation,
+                                           "_resumeUnsafeThrowingContinuation");
+}
+FuncDecl *
+SILGenModule::getResumeUnsafeThrowingContinuationWithError() {
+  return lookUpResumeContinuationIntrinsic(getASTContext(),
+                                 ResumeUnsafeThrowingContinuationWithError,
+                                 "_resumeUnsafeThrowingContinuationWithError");
+}
+
 ProtocolConformance *SILGenModule::getNSErrorConformanceToError() {
   if (NSErrorConformanceToError)
     return *NSErrorConformanceToError;
diff --git a/lib/SILGen/SILGen.h b/lib/SILGen/SILGen.h
@@ -27,6 +27,7 @@
 
 namespace swift {
   class SILBasicBlock;
+  class ForeignAsyncConvention;
 
 namespace Lowering {
   class TypeConverter;
@@ -118,6 +119,10 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
 
   Optional<ProtocolConformance *> NSErrorConformanceToError;
 
+  Optional<FuncDecl*> ResumeUnsafeContinuation;
+  Optional<FuncDecl*> ResumeUnsafeThrowingContinuation;
+  Optional<FuncDecl*> ResumeUnsafeThrowingContinuationWithError;
+
 public:
   SILGenModule(SILModule &M, ModuleDecl *SM);
 
@@ -163,6 +168,14 @@ 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,
+                                           ForeignAsyncConvention convention);
 
   /// Determine whether the given class has any instance variables that
   /// need to be destroyed.
@@ -460,6 +473,13 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// Retrieve the conformance of NSError to the Error protocol.
   ProtocolConformance *getNSErrorConformanceToError();
 
+  /// Retrieve the _Concurrency._resumeUnsafeContinuation intrinsic.
+  FuncDecl *getResumeUnsafeContinuation();
+  /// Retrieve the _Concurrency._resumeUnsafeThrowingContinuation intrinsic.
+  FuncDecl *getResumeUnsafeThrowingContinuation();
+  /// Retrieve the _Concurrency._resumeUnsafeThrowingContinuationWithError intrinsic.
+  FuncDecl *getResumeUnsafeThrowingContinuationWithError();
+
   SILFunction *getKeyPathProjectionCoroutine(bool isReadAccess,
                                              KeyPathTypeKind typeKind);
 
diff --git a/lib/SILGen/SILGenApply.cpp b/lib/SILGen/SILGenApply.cpp
diff --git a/lib/SILGen/SILGenThunk.cpp b/lib/SILGen/SILGenThunk.cpp
diff --git a/stdlib/public/Concurrency/PartialAsyncTask.swift b/stdlib/public/Concurrency/PartialAsyncTask.swift
diff --git a/test/SILGen/objc_async.swift b/test/SILGen/objc_async.swift

Original file line number	Diff line number	Diff line change
`@@ -2166,6 +2166,9 @@ void Remangler::mangleGlobalVariableOnceFunction(Node *node) {`
`2166`	`2166`	`void Remangler::mangleGlobalVariableOnceDeclList(Node *node) {`
`2167`	`2167`	`unreachable("unsupported");`
`2168`	`2168`	`}`
	`2169`	`+void Remangler::manglePredefinedObjCAsyncCompletionHandlerImpl(Node *node) {`
	`2170`	`+ unreachable("unsupported");`
	`2171`	`+}`
`2169`	`2172`	`void Remangler::mangleObjCAsyncCompletionHandlerImpl(Node *node) {`
`2170`	`2173`	`unreachable("unsupported");`
`2171`	`2174`	`}`