Merge pull request #38518 from nate-chandler/rdar80847020

[SILGen] Precompute "returned" completion handler indices.

Author: nate-chandler

lib/SILGen/SILGenThunk.cpp

@@ -261,19 +261,23 @@ SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
       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();
-      auto flagIndex = convention.completionHandlerFlagParamIndex();
+      // Increment the error and flag indices if present.  They do not account
+      // for the fact that they are preceded by the block_storage arguments.
+      auto errorIndex = convention.completionHandlerErrorParamIndex().map(
+          [](auto original) { return original + 1; });
+      auto flagIndex = convention.completionHandlerFlagParamIndex().map(
+          [](auto original) { return original + 1; });
 
       FuncDecl *resumeIntrinsic;
 
       SILBasicBlock *returnBB = nullptr;
       if (errorIndex) {
         resumeIntrinsic = getResumeUnsafeThrowingContinuation();
         auto errorIntrinsic = getResumeUnsafeThrowingContinuationWithError();
-        
-        auto errorArgument = params[*errorIndex + 1];
+
+        auto errorArgument = params[*errorIndex];
         auto someErrorBB = SGF.createBasicBlock(FunctionSection::Postmatter);
         auto noneErrorBB = SGF.createBasicBlock();
         returnBB = SGF.createBasicBlockAfter(noneErrorBB);
@@ -282,8 +286,8 @@ SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
         // Check whether there's an error, based on the presence of a flag
         // parameter. If there is a flag parameter, test it against zero.
         if (flagIndex) {
-          auto flagArgument = params[*flagIndex + 1];
-          
+          auto flagArgument = params[*flagIndex];
+
           // The flag must be an integer type. Get the underlying builtin
           // integer field from it.
           auto builtinFlagArg = SGF.emitUnwrapIntegerResult(loc, flagArgument.getValue());
@@ -378,31 +382,31 @@ SILGenModule::getOrCreateForeignAsyncCompletionHandlerImplFunction(
           bridgedArg.forwardInto(SGF, loc, destBuf);
         };
 
+        // Collect the indices which correspond to the values to be returned.
+        SmallVector<unsigned long, 4> paramIndices;
+        for (auto index : indices(params)) {
+          // The first index is the block_storage parameter.
+          if (index == 0)
+            continue;
+          if (errorIndex && index == *errorIndex)
+            continue;
+          if (flagIndex && index == *flagIndex)
+            continue;
+          paramIndices.push_back(index);
+        }
         if (auto resumeTuple = dyn_cast<TupleType>(resumeType)) {
+          assert(paramIndices.size() == resumeTuple->getNumElements());
           assert(params.size() == resumeTuple->getNumElements()
                                    + 1 + (bool)errorIndex + (bool)flagIndex);
           for (unsigned i : indices(resumeTuple.getElementTypes())) {
-            unsigned paramI = i;
-            if (errorIndex && paramI >= *errorIndex) {
-              ++paramI;
-            }
-            if (flagIndex && paramI >= *flagIndex) {
-              ++paramI;
-            }
             auto resumeEltBuf = SGF.B.createTupleElementAddr(loc,
                                                              resumeArgBuf, i);
-            prepareArgument(resumeEltBuf, params[paramI + 1]);
+            prepareArgument(resumeEltBuf, params[paramIndices[i]]);
           }
         } else {
+          assert(paramIndices.size() == 1);
           assert(params.size() == 2 + (bool)errorIndex + (bool)flagIndex);
-          unsigned paramI = 0;
-          if (errorIndex && paramI >= *errorIndex) {
-            ++paramI;
-          }
-          if (flagIndex && paramI >= *flagIndex) {
-            ++paramI;
-          }
-          prepareArgument(resumeArgBuf, params[paramI + 1]);
+          prepareArgument(resumeArgBuf, params[paramIndices[0]]);
         }
 
         // Resume the continuation with the composed bridged result.

test/Interpreter/Inputs/rdar80847020.h

@@ -0,0 +1,9 @@
+#import <Foundation/Foundation.h>
+
+#pragma clang assume_nonnull begin
+
+@interface Clazz : NSObject
+-(void)doSomethingMultiResultFlaggyWithCompletionHandler:(void (^)(BOOL, NSString *_Nullable, NSError *_Nullable, NSString *_Nullable))completionHandler __attribute__((swift_async_error(zero_argument, 1)));
+@end
+
+#pragma clang assume_nonnull end

test/Interpreter/Inputs/rdar80847020.m

@@ -0,0 +1,11 @@
+#include "rdar80847020.h"
+
+#pragma clang assume_nonnull begin
+
+@implementation Clazz
+-(void)doSomethingMultiResultFlaggyWithCompletionHandler:(void (^)(BOOL, NSString *_Nullable, NSError *_Nullable, NSString *_Nullable))completionHandler __attribute__((swift_async_error(zero_argument, 1))) {
+    completionHandler(YES, @"hi", NULL, @"bye");
+}
+@end
+
+#pragma clang assume_nonnull end

test/Interpreter/rdar80847020.swift

@@ -0,0 +1,20 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-clang %S/Inputs/rdar80847020.m -I %S/Inputs -c -o %t/rdar80847020.o
+// RUN: %target-build-swift -import-objc-header %S/Inputs/rdar80847020.h -Xlinker %t/rdar80847020.o -parse-as-library %s -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: objc_interop
+
+func run(_ s: Clazz) async throws {
+    let res: (String, String) = try await s.doSomethingMultiResultFlaggy()
+    // CHECK: ("hi", "bye")
+    print(res)
+}
+
+@main struct Main {
+  static func main() async throws {
+    try await run(Clazz())
+  }
+}