SILGen: Handle pseudogeneric completion-handler-based async APIs.

jckarter · jckarter · commit eea3a6c27f27 · 2021-02-16T16:05:04.000-08:00
Plumb generic signatures through the codegen for invoking foreign APIs as async, so that we
correctly handle APIs declared on ObjC lightweight generic classes. rdar://74361267
diff --git a/docs/ABI/Mangling.rst b/docs/ABI/Mangling.rst
@@ -222,6 +222,9 @@ types where the metadata itself has unknown layout.)
   global ::= from-type to-type generic-signature? 'TR'  // reabstraction thunk
   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 generic-signature? 'TR'  // reabstraction thunk
+  global ::= impl-function-type type generic-signature? 'Tz'     // objc-to-swift-async completion handler block implementation
+  global ::= impl-function-type type generic-signature? '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
@@ -171,6 +171,7 @@ class ASTMangler : public Mangler {
   /// predefined in the Swift runtime for the given type signature.
   std::string mangleObjCAsyncCompletionHandlerImpl(CanSILFunctionType BlockType,
                                                    CanType ResultType,
+                                                   CanGenericSignature Sig,
                                                    bool predefined);
   
   /// Mangle the derivative function (JVP/VJP), or optionally its vtable entry
diff --git a/lib/AST/ASTMangler.cpp b/lib/AST/ASTMangler.cpp
@@ -398,10 +398,13 @@ std::string ASTMangler::mangleReabstractionThunkHelper(
 std::string ASTMangler::mangleObjCAsyncCompletionHandlerImpl(
                                                    CanSILFunctionType BlockType,
                                                    CanType ResultType,
+                                                   CanGenericSignature Sig,
                                                    bool predefined) {
   beginMangling();
   appendType(BlockType);
   appendType(ResultType);
+  if (Sig)
+    appendGenericSignature(Sig);
   appendOperator(predefined ? "TZ" : "Tz");
   return finalize();
 }
diff --git a/lib/Demangling/Demangler.cpp b/lib/Demangling/Demangler.cpp
@@ -2301,12 +2301,16 @@ NodePointer Demangler::demangleThunkOrSpecialization() {
     }
     case 'z':
     case 'Z': {
+      NodePointer sig = popNode(Node::Kind::DependentGenericSignature);
       NodePointer resultType = popNode(Node::Kind::Type);
       NodePointer implType = popNode(Node::Kind::Type);
-      return createWithChildren(c == 'z'
+      auto node = createWithChildren(c == 'z'
                                   ? Node::Kind::ObjCAsyncCompletionHandlerImpl
                                   : Node::Kind::PredefinedObjCAsyncCompletionHandlerImpl,
                                 implType, resultType);
+      if (sig)
+        addChild(node, sig);
+      return node;
     }
     case 'V': {
       NodePointer Base = popNode(isEntity);
diff --git a/lib/SIL/IR/AbstractionPattern.cpp b/lib/SIL/IR/AbstractionPattern.cpp
@@ -608,7 +608,15 @@ AbstractionPattern::getObjCMethodAsyncCompletionHandlerType(
     auto callbackParamTy = getObjCMethod()->parameters()[paramIndex]
                                           ->getType().getTypePtr();
     
-    return AbstractionPattern(swiftCompletionHandlerType, callbackParamTy);
+    CanGenericSignature patternSig;
+    if (auto origSig = getGenericSignature()) {
+      patternSig = origSig;
+    } else if (auto genFnTy = dyn_cast<GenericFunctionType>(getType())) {
+      patternSig = genFnTy->getGenericSignature()->getCanonicalSignature();
+    }
+    
+    return AbstractionPattern(patternSig,
+                              swiftCompletionHandlerType, callbackParamTy);
   }
   case Kind::Opaque:
   case Kind::OpaqueFunction:
diff --git a/lib/SILGen/ResultPlan.cpp b/lib/SILGen/ResultPlan.cpp
@@ -524,17 +524,24 @@ class ForeignAsyncInitializationPlan final : public ResultPlan {
       handlerIsOptional = false;
       impFnTy = cast<SILFunctionType>(impTy.getASTType());
     }
+    auto env = SGF.F.getGenericEnvironment();
+    auto sig = env ? env->getGenericSignature()->getCanonicalSignature()
+                   : CanGenericSignature();
     SILFunction *impl = SGF.SGM
-      .getOrCreateForeignAsyncCompletionHandlerImplFunction(impFnTy,
-                                                continuationTy,
-                                                *calleeTypeInfo.foreign.async);
+      .getOrCreateForeignAsyncCompletionHandlerImplFunction(
+                  cast<SILFunctionType>(impFnTy->mapTypeOutOfContext()
+                                               ->getCanonicalType(sig)),
+                  continuationTy->mapTypeOutOfContext()->getCanonicalType(sig),
+                  sig,
+                  *calleeTypeInfo.foreign.async);
     auto impRef = SGF.B.createFunctionRef(loc, impl);
     
     // Initialize the block object for the completion handler.
     SILValue block = SGF.B.createInitBlockStorageHeader(loc, blockStorage,
-                          impRef, SILType::getPrimitiveObjectType(impFnTy), {});
+                          impRef, SILType::getPrimitiveObjectType(impFnTy),
+                          SGF.getForwardingSubstitutionMap());
     
-    // Wrap it in optional if the callee expects if.
+    // Wrap it in optional if the callee expects it.
     if (handlerIsOptional) {
       block = SGF.B.createOptionalSome(loc, block, impTy);
     }
diff --git a/lib/SILGen/SILGen.h b/lib/SILGen/SILGen.h
@@ -181,6 +181,7 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   SILFunction *getOrCreateForeignAsyncCompletionHandlerImplFunction(
                                            CanSILFunctionType blockType,
                                            CanType continuationTy,
+                                           CanGenericSignature sig,
                                            ForeignAsyncConvention convention);
 
   /// Determine whether the given class has any instance variables that
diff --git a/lib/SILGen/SILGenThunk.cpp b/lib/SILGen/SILGenThunk.cpp
@@ -185,6 +185,7 @@ SILFunction *
 SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
                                          CanSILFunctionType blockType,
                                          CanType continuationTy,
+                                         CanGenericSignature sig,
                                          ForeignAsyncConvention convention) {
   // Extract the result and error types from the continuation type.
   auto resumeType = cast<BoundGenericType>(continuationTy).getGenericArgs()[0];
@@ -207,10 +208,11 @@ SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
       blockType->getClangTypeInfo().getType(),
       getASTContext().getClangTypeForIRGen(blockStorageTy));
   
-  auto implTy = SILFunctionType::get(GenericSignature(),
+  auto implTy = SILFunctionType::get(sig,
          blockType->getExtInfo().intoBuilder()
            .withRepresentation(SILFunctionTypeRepresentation::CFunctionPointer)
            .withClangFunctionType(newClangTy)
+           .withIsPseudogeneric((bool)sig)
            .build(),
          SILCoroutineKind::None,
          ParameterConvention::Direct_Unowned,
@@ -223,6 +225,7 @@ SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
   Mangle::ASTMangler Mangler;
   auto name = Mangler.mangleObjCAsyncCompletionHandlerImpl(blockType,
                                                            resumeType,
+                                                           sig,
                                                            /*predefined*/ false);
   
   SILGenFunctionBuilder builder(*this);
@@ -234,6 +237,9 @@ SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
   
   if (F->empty()) {
     // Emit the implementation.
+    if (sig)
+      F->setGenericEnvironment(sig->getGenericEnvironment());
+
     SILGenFunction SGF(*this, *F, SwiftModule);
     {
       Scope scope(SGF, loc);
@@ -282,6 +288,8 @@ SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
         // Resume the continuation as throwing the given error, bridged to a
         // native Swift error.
         auto nativeError = SGF.emitBridgedToNativeError(loc, matchedError);
+        Type replacementTypes[]
+          = {F->mapTypeIntoContext(resumeType)->getCanonicalType()};
         auto subs = SubstitutionMap::get(errorIntrinsic->getGenericSignature(),
                                          replacementTypes,
                                          ArrayRef<ProtocolConformanceRef>{});
@@ -296,7 +304,7 @@ SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
       }
             
       auto loweredResumeTy = SGF.getLoweredType(AbstractionPattern::getOpaque(),
-                                                resumeType);
+                                            F->mapTypeIntoContext(resumeType));
       
       // Prepare the argument for the resume intrinsic, using the non-error
       // arguments to the callback.
@@ -343,6 +351,8 @@ SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
         
         // Resume the continuation with the composed bridged result.
         ManagedValue resumeArg = SGF.emitManagedBufferWithCleanup(resumeArgBuf);
+        Type replacementTypes[]
+          = {F->mapTypeIntoContext(resumeType)->getCanonicalType()};
         auto subs = SubstitutionMap::get(resumeIntrinsic->getGenericSignature(),
                                          replacementTypes,
                                          ArrayRef<ProtocolConformanceRef>{});
diff --git a/test/Inputs/clang-importer-sdk/usr/include/ObjCConcurrency.h b/test/Inputs/clang-importer-sdk/usr/include/ObjCConcurrency.h
@@ -117,6 +117,11 @@ typedef void ( ^ObjCErrorHandler )( NSError * _Nullable inError );
   - (void) myMethod:(NSInteger)value1 foo:(NSInteger)value2;
 @end
 
+@interface GenericObject<T> : NSObject
+- (void)doSomethingWithCompletionHandler:(void (^)(T _Nullable_result, NSError * _Nullable))completionHandler;
+- (void)doAnotherThingWithCompletionHandler:(void (^)(GenericObject<T> *_Nullable))completionHandler;
+@end
+
 #define MAGIC_NUMBER 42
 
 #pragma clang assume_nonnull end
diff --git a/test/SILGen/objc_async.swift b/test/SILGen/objc_async.swift
@@ -64,6 +64,16 @@ func testSlowServer(slowServer: SlowServer) async throws {
   let _: NSObject = try await slowServer.someObject()
 }
 
+func testGeneric<T: AnyObject>(x: GenericObject<T>) async throws {
+  let _: T? = try await x.doSomething()
+  let _: GenericObject<T>? = await x.doAnotherThing()
+}
+
+func testGeneric2<T: AnyObject, U>(x: GenericObject<T>, y: U) async throws {
+  let _: T? = try await x.doSomething()
+  let _: GenericObject<T>? = await x.doAnotherThing()
+}
+
 // CHECK: sil{{.*}}@[[INT_COMPLETION_BLOCK]]
 // CHECK:   [[CONT_ADDR:%.*]] = project_block_storage %0
 // CHECK:   [[CONT:%.*]] = load [trivial] [[CONT_ADDR]]
