[SILGen] Implement async completion bridging via checked continuations

xedin · xedin · commit 6774c0ad1f4b · 2023-10-11T13:46:23.000-07:00
Because `CheckedContinuation` is not a @Frozen struct we have to use `Any` to store it in @block_storage indirectly. If the flag is enabled, we'd emit a block storage with `Any` and initialize the existential with stack allocated `CheckedContinuation` formed from `UnsafeContinuation`. Inside of the completion handler `Any` is going to be projected and cast back to `CheckedContinuation`. (cherry picked from commit 4b22620)
diff --git a/include/swift/AST/KnownSDKTypes.def b/include/swift/AST/KnownSDKTypes.def
@@ -37,6 +37,7 @@ KNOWN_SDK_TYPE_DECL(ObjectiveC, ObjCBool, StructDecl, 0)
 
 // TODO(async): These might move to the stdlib module when concurrency is
 // standardized
+KNOWN_SDK_TYPE_DECL(Concurrency, CheckedContinuation, NominalTypeDecl, 2)
 KNOWN_SDK_TYPE_DECL(Concurrency, UnsafeContinuation, NominalTypeDecl, 2)
 KNOWN_SDK_TYPE_DECL(Concurrency, MainActor, NominalTypeDecl, 0)
 KNOWN_SDK_TYPE_DECL(Concurrency, Job, StructDecl, 0) // TODO: remove in favor of ExecutorJob
diff --git a/lib/SILGen/ResultPlan.cpp b/lib/SILGen/ResultPlan.cpp
@@ -738,6 +738,75 @@ class ForeignAsyncInitializationPlan final : public ResultPlan {
     // A foreign async function shouldn't have any indirect results.
   }
 
+  std::tuple</*blockStorage=*/SILValue, /*blockStorageType=*/CanType,
+             /*continuationType=*/CanType>
+  emitBlockStorage(SILGenFunction &SGF, SILLocation loc, bool throws) {
+    auto &ctx = SGF.getASTContext();
+
+    // Wrap the Builtin.RawUnsafeContinuation in an
+    // UnsafeContinuation<T, E>.
+    auto *unsafeContinuationDecl = ctx.getUnsafeContinuationDecl();
+    auto errorTy = throws ? ctx.getErrorExistentialType() : ctx.getNeverType();
+    auto continuationTy =
+        BoundGenericType::get(unsafeContinuationDecl, /*parent=*/Type(),
+                              {calleeTypeInfo.substResultType, errorTy})
+            ->getCanonicalType();
+
+    auto wrappedContinuation = SGF.B.createStruct(
+        loc, SILType::getPrimitiveObjectType(continuationTy), {continuation});
+
+    const bool checkedBridging = ctx.LangOpts.UseCheckedAsyncObjCBridging;
+
+    // If checked bridging is enabled, wrap that continuation again in a
+    // CheckedContinuation<T, E>
+    if (checkedBridging) {
+      auto *checkedContinuationDecl = ctx.getCheckedContinuationDecl();
+      continuationTy =
+          BoundGenericType::get(checkedContinuationDecl, /*parent=*/Type(),
+                                {calleeTypeInfo.substResultType, errorTy})
+              ->getCanonicalType();
+    }
+
+    auto blockStorageTy = SILBlockStorageType::get(
+        checkedBridging ? ctx.TheAnyType : continuationTy);
+    auto blockStorage = SGF.emitTemporaryAllocation(
+        loc, SILType::getPrimitiveAddressType(blockStorageTy));
+
+    auto continuationAddr = SGF.B.createProjectBlockStorage(loc, blockStorage);
+
+    // Stash continuation in a buffer for a block object.
+
+    if (checkedBridging) {
+      auto createIntrinsic =
+          throws ? SGF.SGM.getCreateCheckedThrowingContinuation()
+                 : SGF.SGM.getCreateCheckedContinuation();
+
+      // In this case block storage captures `Any` which would be initialized
+      // with an checked continuation.
+      auto underlyingContinuationAddr =
+          SGF.B.createInitExistentialAddr(loc, continuationAddr, continuationTy,
+                                          SGF.getLoweredType(continuationTy),
+                                          /*conformances=*/{});
+
+      auto subs = SubstitutionMap::get(createIntrinsic->getGenericSignature(),
+                                       {calleeTypeInfo.substResultType},
+                                       ArrayRef<ProtocolConformanceRef>{});
+
+      InitializationPtr underlyingInit(
+          new KnownAddressInitialization(underlyingContinuationAddr));
+      auto continuationMV =
+          ManagedValue::forRValueWithoutOwnership(wrappedContinuation);
+      SGF.emitApplyOfLibraryIntrinsic(loc, createIntrinsic, subs,
+                                      {continuationMV}, SGFContext())
+          .forwardInto(SGF, loc, underlyingInit.get());
+    } else {
+      SGF.B.createStore(loc, wrappedContinuation, continuationAddr,
+                        StoreOwnershipQualifier::Trivial);
+    }
+
+    return std::make_tuple(blockStorage, blockStorageTy, continuationTy);
+  }
+
   ManagedValue
   emitForeignAsyncCompletionHandler(SILGenFunction &SGF,
                                     AbstractionPattern origFormalType,
@@ -751,28 +820,11 @@ class ForeignAsyncInitializationPlan final : public ResultPlan {
     continuation = SGF.B.createGetAsyncContinuationAddr(loc, resumeBuf,
                                calleeTypeInfo.substResultType, throws);
 
-    // Wrap the Builtin.RawUnsafeContinuation in an
-    // UnsafeContinuation<T, E>.
-    auto continuationDecl = SGF.getASTContext().getUnsafeContinuationDecl();
-
-    auto errorTy = throws
-      ? SGF.getASTContext().getErrorExistentialType()
-      : SGF.getASTContext().getNeverType();
-    auto continuationTy = BoundGenericType::get(continuationDecl, Type(),
-                                                { calleeTypeInfo.substResultType, errorTy })
-      ->getCanonicalType();
-    auto wrappedContinuation =
-        SGF.B.createStruct(loc,
-                           SILType::getPrimitiveObjectType(continuationTy),
-                           {continuation});
-
-    // Stash it in a buffer for a block object.
-    auto blockStorageTy = SILBlockStorageType::get(continuationTy);
-    auto blockStorage = SGF.emitTemporaryAllocation(
-        loc, SILType::getPrimitiveAddressType(blockStorageTy));
-    auto continuationAddr = SGF.B.createProjectBlockStorage(loc, blockStorage);
-    SGF.B.createStore(loc, wrappedContinuation, continuationAddr,
-                      StoreOwnershipQualifier::Trivial);
+    SILValue blockStorage;
+    CanType blockStorageTy;
+    CanType continuationTy;
+    std::tie(blockStorage, blockStorageTy, continuationTy) =
+        emitBlockStorage(SGF, loc, throws);
 
     // Get the block invocation function for the given completion block type.
     auto completionHandlerIndex = calleeTypeInfo.foreign.async
@@ -797,6 +849,7 @@ class ForeignAsyncInitializationPlan final : public ResultPlan {
             cast<SILFunctionType>(
                 impFnTy->mapTypeOutOfContext()->getReducedType(sig)),
             blockStorageTy->mapTypeOutOfContext()->getReducedType(sig),
+            continuationTy->mapTypeOutOfContext()->getReducedType(sig),
             origFormalType, sig, calleeTypeInfo);
     auto impRef = SGF.B.createFunctionRef(loc, impl);
 
diff --git a/lib/SILGen/SILGen.h b/lib/SILGen/SILGen.h
@@ -190,8 +190,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// as `async` in Swift.
   SILFunction *getOrCreateForeignAsyncCompletionHandlerImplFunction(
       CanSILFunctionType blockType, CanType blockStorageType,
-      AbstractionPattern origFormalType, CanGenericSignature sig,
-      CalleeTypeInfo &calleeInfo);
+      CanType continuationType, AbstractionPattern origFormalType,
+      CanGenericSignature sig, CalleeTypeInfo &calleeInfo);
 
   /// Determine whether the given class has any instance variables that
   /// need to be destroyed.
diff --git a/lib/SILGen/SILGenThunk.cpp b/lib/SILGen/SILGenThunk.cpp
@@ -200,8 +200,8 @@ static const clang::Type *prependParameterType(
 
 SILFunction *SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
     CanSILFunctionType blockType, CanType blockStorageType,
-    AbstractionPattern origFormalType, CanGenericSignature sig,
-    CalleeTypeInfo &calleeInfo) {
+    CanType continuationType, AbstractionPattern origFormalType,
+    CanGenericSignature sig, CalleeTypeInfo &calleeInfo) {
   auto convention = *calleeInfo.foreign.async;
   auto resumeType =
       calleeInfo.substResultType->mapTypeOutOfContext()->getReducedType(sig);
@@ -293,11 +293,33 @@ SILFunction *SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
 
       // 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::forObjectRValueWithoutOwnership(continuationVal);
+      SILValue continuationAddr =
+          SGF.B.createProjectBlockStorage(loc, blockStorage);
+
+      auto &ctx = SGF.getASTContext();
+
+      bool checkedBridging = ctx.LangOpts.UseCheckedAsyncObjCBridging;
+
+      ManagedValue continuation;
+      if (checkedBridging) {
+        FormalEvaluationScope scope(SGF);
+
+        auto underlyingValueTy = OpenedArchetypeType::get(ctx.TheAnyType, sig);
+
+        auto underlyingValueAddr = SGF.emitOpenExistential(
+            loc, ManagedValue::forTrivialAddressRValue(continuationAddr),
+            SGF.getLoweredType(underlyingValueTy), AccessKind::Read);
+
+        continuation = SGF.B.createUncheckedAddrCast(
+            loc, underlyingValueAddr,
+            SILType::getPrimitiveAddressType(
+                F->mapTypeIntoContext(continuationType)->getCanonicalType()));
+      } else {
+        auto continuationVal = SGF.B.createLoad(
+            loc, continuationAddr, LoadOwnershipQualifier::Trivial);
+        continuation =
+            ManagedValue::forObjectRValueWithoutOwnership(continuationVal);
+      }
 
       // Check for an error if the convention includes one.
       // Increment the error and flag indices if present.  They do not account
@@ -311,8 +333,12 @@ SILFunction *SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
 
       SILBasicBlock *returnBB = nullptr;
       if (errorIndex) {
-        resumeIntrinsic = getResumeUnsafeThrowingContinuation();
-        auto errorIntrinsic = getResumeUnsafeThrowingContinuationWithError();
+        resumeIntrinsic = checkedBridging
+                              ? getResumeCheckedThrowingContinuation()
+                              : getResumeUnsafeThrowingContinuation();
+        auto errorIntrinsic =
+            checkedBridging ? getResumeCheckedThrowingContinuationWithError()
+                            : getResumeUnsafeThrowingContinuationWithError();
 
         auto errorArgument = params[*errorIndex];
         auto someErrorBB = SGF.createBasicBlock(FunctionSection::Postmatter);
@@ -385,9 +411,12 @@ SILFunction *SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
         SGF.B.createBranch(loc, returnBB);
         SGF.B.emitBlock(noneErrorBB);
       } else if (auto foreignError = calleeInfo.foreign.error) {
-        resumeIntrinsic = getResumeUnsafeThrowingContinuation();
+        resumeIntrinsic = checkedBridging
+                              ? getResumeCheckedThrowingContinuation()
+                              : getResumeUnsafeThrowingContinuation();
       } else {
-        resumeIntrinsic = getResumeUnsafeContinuation();
+        resumeIntrinsic = checkedBridging ? getResumeCheckedContinuation()
+                                          : getResumeUnsafeContinuation();
       }
 
       auto loweredResumeTy = SGF.getLoweredType(AbstractionPattern::getOpaque(),
diff --git a/test/SILGen/objc_async_checked.swift b/test/SILGen/objc_async_checked.swift
diff --git a/test/SILGen/objc_effectful_properties_checked.swift b/test/SILGen/objc_effectful_properties_checked.swift
