SILGen: Generate bodies for completion handler block impls

jckarter · jckarter · commit 368dc0f4016c · 2020-11-10T16:36:50.000-08:00
diff --git a/lib/SILGen/ResultPlan.cpp b/lib/SILGen/ResultPlan.cpp
@@ -507,7 +507,8 @@ class ForeignAsyncInitializationPlan final : public ResultPlan {
       .getInterfaceType());
     SILFunction *impl = SGF.SGM
       .getOrCreateForeignAsyncCompletionHandlerImplFunction(implTy,
-                                                            continuationTy);
+                                                continuationTy,
+                                                *calleeTypeInfo.foreign.async);
     auto implRef = SGF.B.createFunctionRef(loc, impl);
     
     // Initialize the block object for the completion handler.
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);
 
@@ -169,7 +174,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// as `async` in Swift.
   SILFunction *getOrCreateForeignAsyncCompletionHandlerImplFunction(
                                            CanSILFunctionType blockType,
-                                           CanType continuationTy);
+                                           CanType continuationTy,
+                                           ForeignAsyncConvention convention);
 
   /// Determine whether the given class has any instance variables that
   /// need to be destroyed.
@@ -467,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/SILGenThunk.cpp b/lib/SILGen/SILGenThunk.cpp
@@ -27,6 +27,7 @@
 #include "swift/AST/ASTMangler.h"
 #include "swift/AST/DiagnosticsSIL.h"
 #include "swift/AST/FileUnit.h"
+#include "swift/AST/ForeignAsyncConvention.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/SIL/PrettyStackTrace.h"
 #include "swift/SIL/SILArgument.h"
@@ -141,7 +142,8 @@ SILGenFunction::emitGlobalFunctionRef(SILLocation loc, SILDeclRef constant,
 SILFunction *
 SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
                                          CanSILFunctionType blockType,
-                                         CanType continuationTy) {
+                                         CanType continuationTy,
+                                         ForeignAsyncConvention convention) {
   // Extract the result type from the continuation type.
   auto resumeType = cast<BoundGenericType>(continuationTy).getGenericArgs()[0];
   
@@ -184,10 +186,125 @@ SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
   if (F->empty()) {
     // TODO: Emit the implementation.
     SILGenFunction SGF(*this, *F, SwiftModule);
-    SmallVector<ManagedValue, 4> params;
-    SGF.collectThunkParams(loc, params);
+    {
+      Scope scope(SGF, loc);
+      SmallVector<ManagedValue, 4> params;
+      SGF.collectThunkParams(loc, params);
 
-    SGF.B.createUnreachable(loc);
+      // Get the continuation out of the block object.
+      auto blockStorage = params[0].getValue();
+      auto continuationAddr = SGF.B.createProjectBlockStorage(loc, blockStorage);
+      auto continuationVal = SGF.B.createLoad(loc, continuationAddr,
+                                           LoadOwnershipQualifier::Trivial);
+      auto continuation = ManagedValue::forUnmanaged(continuationVal);
+      
+      // Check for an error if the convention includes one.
+      auto errorIndex = convention.completionHandlerErrorParamIndex();
+    
+      FuncDecl *resumeIntrinsic, *errorIntrinsic;
+
+      SILBasicBlock *returnBB = nullptr;
+      if (errorIndex) {
+        resumeIntrinsic = getResumeUnsafeThrowingContinuation();
+        errorIntrinsic = getResumeUnsafeThrowingContinuationWithError();
+        
+        auto errorArgument = params[*errorIndex + 1];
+        auto someErrorBB = SGF.createBasicBlock(FunctionSection::Postmatter);
+        auto noneErrorBB = SGF.createBasicBlock();
+        returnBB = SGF.createBasicBlockAfter(noneErrorBB);
+
+        auto &C = SGF.getASTContext();
+        std::pair<EnumElementDecl *, SILBasicBlock *> switchErrorBBs[] = {
+          {C.getOptionalSomeDecl(), someErrorBB},
+          {C.getOptionalNoneDecl(), noneErrorBB}
+        };
+        
+        SGF.B.createSwitchEnum(loc, errorArgument.borrow(SGF, loc).getValue(),
+                               /*default*/ nullptr,
+                               switchErrorBBs);
+        
+        SGF.B.emitBlock(someErrorBB);
+        
+        auto matchedErrorTy = errorArgument.getType().getOptionalObjectType();
+        auto matchedError = SGF.B
+          .createGuaranteedTransformingTerminatorArgument(matchedErrorTy);
+        
+        // Resume the continuation as throwing the given error, bridged to a
+        // native Swift error.
+        auto nativeError = SGF.emitBridgedToNativeError(loc, matchedError);
+        Type replacementTypes[] = {resumeType};
+        auto subs = SubstitutionMap::get(errorIntrinsic->getGenericSignature(),
+                                         replacementTypes,
+                                         ArrayRef<ProtocolConformanceRef>{});
+        SGF.emitApplyOfLibraryIntrinsic(loc, errorIntrinsic, subs,
+                                        {continuation, nativeError},
+                                        SGFContext());
+        
+        SGF.B.createBranch(loc, returnBB);
+        SGF.B.emitBlock(noneErrorBB);
+      } else {
+        resumeIntrinsic = getResumeUnsafeContinuation();
+      }
+            
+      auto loweredResumeTy = SGF.getLoweredType(AbstractionPattern::getOpaque(),
+                                                resumeType);
+      
+      // Prepare the argument for the resume intrinsic, using the non-error
+      // arguments to the callback.
+      {
+        Scope resumeScope(SGF, loc);
+        unsigned errorIndexBoundary = errorIndex ? *errorIndex : ~0u;
+        auto resumeArgBuf = SGF.emitTemporaryAllocation(loc,
+                                              loweredResumeTy.getAddressType());
+
+        auto prepareArgument = [&](SILValue destBuf, ManagedValue arg) {
+          // Convert the ObjC argument to the bridged Swift representation we
+          // want.
+          ManagedValue bridgedArg = SGF.emitBridgedToNativeValue(loc,
+                                       arg,
+                                       arg.getType().getASTType(),
+                                       // FIXME: pass down formal type
+                                       destBuf->getType().getASTType(),
+                                       destBuf->getType().getObjectType());
+          bridgedArg.forwardInto(SGF, loc, destBuf);
+        };
+
+        if (auto resumeTuple = dyn_cast<TupleType>(resumeType)) {
+          assert(params.size() == resumeTuple->getNumElements()
+                                   + 1 + (bool)errorIndex);
+          for (auto i : indices(resumeTuple.getElementTypes())) {
+            auto resumeEltBuf = SGF.B.createTupleElementAddr(loc,
+                                                             resumeArgBuf, i);
+            auto arg = params[1 + i + (i >= errorIndexBoundary)];
+            prepareArgument(resumeEltBuf, arg);
+          }
+        } else {
+          assert(params.size() == 2 + (bool)errorIndex);
+          prepareArgument(resumeArgBuf, params[1 + (errorIndexBoundary == 0)]);
+        }
+        
+        
+        // Resume the continuation with the composed bridged result.
+        ManagedValue resumeArg = SGF.emitManagedBufferWithCleanup(resumeArgBuf);
+        Type replacementTypes[] = {resumeType};
+        auto subs = SubstitutionMap::get(resumeIntrinsic->getGenericSignature(),
+                                         replacementTypes,
+                                         ArrayRef<ProtocolConformanceRef>{});
+        SGF.emitApplyOfLibraryIntrinsic(loc, resumeIntrinsic, subs,
+                                        {continuation, resumeArg},
+                                        SGFContext());
+      }
+      
+      // Now we've resumed the continuation one way or another. Return from the
+      // completion callback.
+      if (returnBB) {
+        SGF.B.createBranch(loc, returnBB);
+        SGF.B.emitBlock(returnBB);
+      }
+    }
+
+    SGF.B.createReturn(loc,
+                       SILUndef::get(SGF.SGM.Types.getEmptyTupleType(), SGF.F));
   }
   
   return F;
diff --git a/stdlib/public/Concurrency/PartialAsyncTask.swift b/stdlib/public/Concurrency/PartialAsyncTask.swift
@@ -24,15 +24,46 @@ public struct PartialAsyncTask {
 public struct UnsafeContinuation<T> {
   private var context: UnsafeRawPointer
 
-  public func resume(_: T) { }
+  public func resume(_: __owned T) { }
 }
 
 @frozen
 public struct UnsafeThrowingContinuation<T> {
   private var context: UnsafeRawPointer
 
-  public func resume(_: T) { }
-  public func fail(_: Error) { }
+  public func resume(_: __owned T) { }
+  public func fail(_: __owned Error) { }
 }
 
+#if _runtime(_ObjC)
 
+// Intrinsics used by SILGen to resume or fail continuations
+// for
+@_alwaysEmitIntoClient
+@usableFromInline
+internal func _resumeUnsafeContinuation<T>(
+  _ continuation: UnsafeContinuation<T>,
+  _ value: __owned T
+) {
+  continuation.resume(value)
+}
+
+@_alwaysEmitIntoClient
+@usableFromInline
+internal func _resumeUnsafeThrowingContinuation<T>(
+  _ continuation: UnsafeThrowingContinuation<T>,
+  _ value: __owned T
+) {
+  continuation.resume(value)
+}
+
+@_alwaysEmitIntoClient
+@usableFromInline
+internal func _resumeUnsafeThrowingContinuationWithError<T>(
+  _ continuation: UnsafeThrowingContinuation<T>,
+  _ error: __owned Error
+) {
+  continuation.fail(error)
+}
+
+#endif
diff --git a/test/SILGen/objc_async.swift b/test/SILGen/objc_async.swift
@@ -12,7 +12,7 @@ func testSlowServer(slowServer: SlowServer) async throws {
   // CHECK: [[BLOCK_STORAGE:%.*]] = alloc_stack $@block_storage UnsafeContinuation<Int>
   // CHECK: [[CONT_SLOT:%.*]] = project_block_storage [[BLOCK_STORAGE]]
   // CHECK: store [[CONT]] to [trivial] [[CONT_SLOT]]
-  // CHECK: [[BLOCK_IMPL:%.*]] = function_ref @{{.*}} : $@convention(c) (@inout_aliasable @block_storage UnsafeContinuation<Int>, Int) -> ()
+  // CHECK: [[BLOCK_IMPL:%.*]] = function_ref @[[INT_COMPLETION_BLOCK:.*]] : $@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]+]]
@@ -27,7 +27,7 @@ func testSlowServer(slowServer: SlowServer) async throws {
   // CHECK: [[BLOCK_STORAGE:%.*]] = alloc_stack $@block_storage UnsafeThrowingContinuation<String>
   // CHECK: [[CONT_SLOT:%.*]] = project_block_storage [[BLOCK_STORAGE]]
   // CHECK: store [[CONT]] to [trivial] [[CONT_SLOT]]
-  // CHECK: [[BLOCK_IMPL:%.*]] = function_ref @{{.*}} : $@convention(c) (@inout_aliasable @block_storage UnsafeThrowingContinuation<String>, Optional<NSString>, Optional<NSError>) -> ()
+  // CHECK: [[BLOCK_IMPL:%.*]] = function_ref @[[STRING_COMPLETION_THROW_BLOCK:.*]] : $@convention(c) (@inout_aliasable @block_storage UnsafeThrowingContinuation<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]+]]
@@ -38,10 +38,50 @@ func testSlowServer(slowServer: SlowServer) async throws {
   let _: String = try await slowServer.findAnswer()
 
   // CHECK: objc_method {{.*}} $@convention(objc_method) (NSString, @convention(block) () -> (), SlowServer) -> ()
+  // CHECK: [[BLOCK_IMPL:%.*]] = function_ref @[[VOID_COMPLETION_BLOCK:.*]] : $@convention(c) (@inout_aliasable @block_storage UnsafeContinuation<()>) -> ()
   await slowServer.serverRestart("somewhere")
 
   // CHECK: [[ERROR]]([[ERROR_VALUE:%.*]] : @owned $Error):
   // CHECK: dealloc_stack [[RESUME_BUF]]
   // CHECK: throw [[ERROR_VALUE]]
 
 }
+
+// CHECK: sil{{.*}}@[[INT_COMPLETION_BLOCK]]
+// CHECK:   [[CONT_ADDR:%.*]] = project_block_storage %0
+// CHECK:   [[CONT:%.*]] = load [trivial] [[CONT_ADDR]]
+// CHECK:   [[RESULT_BUF:%.*]] = alloc_stack $Int
+// CHECK:   store %1 to [trivial] [[RESULT_BUF]]
+// CHECK:   [[RESUME:%.*]] = function_ref @{{.*}}resumeUnsafeContinuation
+// CHECK:   apply [[RESUME]]<Int>([[CONT]], [[RESULT_BUF]])
+
+// CHECK: sil{{.*}}@[[STRING_COMPLETION_THROW_BLOCK]]
+// CHECK:   [[RESUME_IN:%.*]] = copy_value %1
+// CHECK:   [[ERROR_IN:%.*]] = copy_value %2
+// CHECK:   [[CONT_ADDR:%.*]] = project_block_storage %0
+// CHECK:   [[CONT:%.*]] = load [trivial] [[CONT_ADDR]]
+// CHECK:   [[ERROR_IN_B:%.*]] = begin_borrow [[ERROR_IN]]
+// CHECK:   switch_enum [[ERROR_IN_B]] : {{.*}}, case #Optional.some!enumelt: [[ERROR_BB:bb[0-9]+]], case #Optional.none!enumelt: [[RESUME_BB:bb[0-9]+]]
+// CHECK: [[RESUME_BB]]:
+// CHECK:   [[RESULT_BUF:%.*]] = alloc_stack $String
+// CHECK:   [[BRIDGE:%.*]] = function_ref @{{.*}}unconditionallyBridgeFromObjectiveC
+// CHECK:   [[BRIDGED_RESULT:%.*]] = apply [[BRIDGE]]([[RESUME_IN]]
+// CHECK:   store [[BRIDGED_RESULT]] to [init] [[RESULT_BUF]]
+// CHECK:   [[RESUME:%.*]] = function_ref @{{.*}}resumeUnsafeThrowingContinuation
+// CHECK:   apply [[RESUME]]<String>([[CONT]], [[RESULT_BUF]])
+// CHECK:   br [[END_BB:bb[0-9]+]]
+// CHECK: [[END_BB]]:
+// CHECK:   return
+// CHECK: [[ERROR_BB]]([[ERROR_IN_UNWRAPPED:%.*]] : @guaranteed $NSError):
+// CHECK:   [[ERROR:%.*]] = init_existential_ref [[ERROR_IN_UNWRAPPED]]
+// CHECK:   [[RESUME_WITH_ERROR:%.*]] = function_ref @{{.*}}resumeUnsafeThrowingContinuationWithError
+// CHECK:   [[ERROR_COPY:%.*]] = copy_value [[ERROR]]
+// CHECK:   apply [[RESUME_WITH_ERROR]]<String>([[CONT]], [[ERROR_COPY]])
+// CHECK:   br [[END_BB]]
+
+// CHECK: sil {{.*}} @[[VOID_COMPLETION_BLOCK]]
+// CHECK:   [[CONT_ADDR:%.*]] = project_block_storage %0
+// CHECK:   [[CONT:%.*]] = load [trivial] [[CONT_ADDR]]
+// CHECK:   [[RESULT_BUF:%.*]] = alloc_stack $()
+// CHECK:   [[RESUME:%.*]] = function_ref @{{.*}}resumeUnsafeContinuation
+// CHECK:   apply [[RESUME]]<()>([[CONT]], [[RESULT_BUF]])
