Merge pull request #41535 from jckarter/objc-async-bridging-multi-result-with-any

SILGen: Handle objc async bridging when there is a combination of `Any` and other results.

Author: jckarter

lib/SILGen/SILGenBridging.cpp

@@ -1593,13 +1593,17 @@ void SILGenFunction::emitNativeToForeignThunk(SILDeclRef thunk) {
     }
   }
 
-  // If we are bridging a Swift method with an Any return value, create a
-  // stack allocation to hold the result, since Any is address-only.
+  // If we are bridging a Swift method with Any return value(s), create a
+  // stack allocation to hold the result(s), since Any is address-only.
   SmallVector<SILValue, 4> args;
-
   if (substConv.hasIndirectSILResults()) {
-    args.push_back(emitTemporaryAllocation(
-        loc, substConv.getSingleSILResultType(getTypeExpansionContext())));
+    for (auto result : substConv.getResults()) {
+      if (!substConv.isSILIndirect(result)) {
+        continue;
+      }
+      args.push_back(emitTemporaryAllocation(
+                loc, substConv.getSILType(result, getTypeExpansionContext())));
+    }
   }
 
   // If the '@objc' was inferred due to deprecated rules,
@@ -1792,14 +1796,29 @@ void SILGenFunction::emitNativeToForeignThunk(SILDeclRef thunk) {
           pushErrorFlag(/*has error*/ false, completionHandlerArgs);
           continue;
         }
-        pushArg(asyncResult,
-                nativeFormalResultType,
-                completionTy->getParameters()[i]);
+        
+        // Use the indirect return argument if the result is indirect.
+        if (substConv.hasIndirectSILResults()) {
+          pushArg(emitManagedRValueWithCleanup(args[0]),
+                  nativeFormalResultType,
+                  completionTy->getParameters()[i]);
+        } else {
+          pushArg(asyncResult,
+                  nativeFormalResultType,
+                  completionTy->getParameters()[i]);
+        }
       }
     } else {
       // A tuple return maps to multiple completion handler parameters.
       auto formalTuple = cast<TupleType>(nativeFormalResultType);
 
+      unsigned indirectResultI = 0;
+      unsigned directResultI = 0;
+      
+      auto directResults = substConv.getDirectSILResults();
+      auto hasMultipleDirectResults
+        = std::next(directResults.begin()) != directResults.end();
+      
       for (unsigned paramI : indices(completionTy->getParameters())) {
         if (errorParamIndex && paramI == *errorParamIndex) {
           pushErrorPlaceholder();
@@ -1813,7 +1832,21 @@ void SILGenFunction::emitNativeToForeignThunk(SILDeclRef thunk) {
                                - (errorFlagIndex && paramI > *errorFlagIndex);
         auto param = completionTy->getParameters()[paramI];
         auto formalTy = formalTuple.getElementType(elementI);
-        auto argPiece = B.createTupleExtract(loc, asyncResult, elementI);
+        ManagedValue argPiece;
+        
+        auto result = substConv.getResults()[elementI];
+        if (substConv.isSILIndirect(result)) {
+          // Take the arg piece from the indirect return arguments.
+          argPiece = emitManagedRValueWithCleanup(args[indirectResultI++]);
+        } else if (hasMultipleDirectResults) {
+          // Take the arg piece from one of the tuple elements of the direct
+          // result tuple from the apply.
+          argPiece = B.createTupleExtract(loc, asyncResult, directResultI++);
+        } else {
+          // Take the entire direct result from the apply as the arg piece.
+          argPiece = asyncResult;
+        }
+        
         pushArg(argPiece, formalTy, param);
       }
     }
@@ -1829,16 +1862,11 @@ void SILGenFunction::emitNativeToForeignThunk(SILDeclRef thunk) {
     // The immediate function result is an empty tuple.
     return SILUndef::get(SGM.Types.getEmptyTupleType(), F);
   };
-  
+    
   if (!substTy->hasErrorResult()) {
     // Create the apply.
     result = B.createApply(loc, nativeFn, subs, args);
-
-    if (substConv.hasIndirectSILResults()) {
-      assert(substTy->getNumResults() == 1);
-      result = args[0];
-    }
-
+  
     // Leave the argument cleanup scope immediately.  This isn't really
     // necessary; it just limits lifetimes a little bit more.
     argScope.pop();
@@ -1849,6 +1877,10 @@ void SILGenFunction::emitNativeToForeignThunk(SILDeclRef thunk) {
     if (foreignAsync) {
       result = passResultToCompletionHandler(result);
     } else {
+      if (substConv.hasIndirectSILResults()) {
+        assert(substTy->getNumResults() == 1);
+        result = args[0];
+      }
       result = emitBridgeReturnValue(*this, loc, result, nativeFormalResultType,
                                      bridgedFormalResultType, objcResultTy);
     }
@@ -1864,17 +1896,16 @@ void SILGenFunction::emitNativeToForeignThunk(SILDeclRef thunk) {
       SILValue nativeResult =
           normalBB->createPhiArgument(swiftResultTy, OwnershipKind::Owned);
 
-      if (substConv.hasIndirectSILResults()) {
-        assert(substTy->getNumResults() == 1);
-        nativeResult = args[0];
-      }
-      
       if (foreignAsync) {
         // If the function is async, pass the results as the success argument(s)
         // to the completion handler, with a nil error.
         passResultToCompletionHandler(nativeResult);
         B.createBranch(loc, contBB);
       } else {
+        if (substConv.hasIndirectSILResults()) {
+          assert(substTy->getNumResults() == 1);
+          nativeResult = args[0];
+        }
         // In this branch, the eventual return value is mostly created
         // by bridging the native return value, but we may need to
         // adjust it slightly.

test/SILGen/Inputs/objc_async_tuple_88949633.h

@@ -0,0 +1,14 @@
+#import <Foundation.h>
+
+#pragma clang assume_nonnull begin
+
+@protocol NSButz
+
+- (void)idString:(NSObject *)x replyHandler:(void (^)(id _Nullable reply, NSString *_Nullable errorMessage))replyHandler __attribute__((swift_async(not_swift_private, 2)));
+- (void)idStringID:(NSObject *)x replyHandler:(void (^)(id _Nullable reply, NSString *_Nullable errorMessage, id _Nullable reply2))replyHandler __attribute__((swift_async(not_swift_private, 2)));
+- (void)stringIDString:(NSObject *)x replyHandler:(void (^)(NSString *_Nullable errorMessage, id _Nullable reply, NSString *_Nullable errorMessage2))replyHandler __attribute__((swift_async(not_swift_private, 2)));
+- (void)idStringIDString:(NSObject *)x replyHandler:(void (^)(id _Nullable reply, NSString *_Nullable errorMessage, id _Nullable reply2, NSString *_Nullable errorMessage2))replyHandler __attribute__((swift_async(not_swift_private, 2)));
+
+@end
+
+#pragma clang assume_nonnull end

test/SILGen/objc_async_tuple_return_88949633.swift

@@ -0,0 +1,25 @@
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -emit-silgen -I %S/Inputs/custom-modules -import-objc-header %S/Inputs/objc_async_tuple_88949633.h -disable-availability-checking %s -verify
+// REQUIRES: concurrency
+// REQUIRES: objc_interop
+
+// rdar://88949633
+
+import Foundation
+import ObjCConcurrency
+
+class Foo: NSObject, NSButz {
+    func idString(_: NSObject) async -> (Any?, String?) {
+        return (nil, nil)
+    }
+    func idStringID(_: NSObject) async -> (Any?, String?, Any?) {
+        return (nil, nil, nil)
+    }
+    func stringIDString(_: NSObject) async -> (String?, Any?, String?) {
+        return (nil, nil, nil)
+    }
+    func idStringIDString(_: NSObject) async -> (Any?, String?, Any?, String?) {
+        return (nil, nil, nil, nil)
+    }
+}
+
+