[SILGen] Precompute "returned" completion handler indices.

When generating the completion handler that is passed to ObjC methods
that are being called as async Swift functions, the parameters taken by
the generated completion handler are matched up with the values that are
returned from the async function.

Previously, this process was done by starting at the index into the list
of values expected to be returned, adding 1 if the index matched first
the error index and second the flag index, and finally adding 1 to
account for the fact that the initial parameter to the completion
handler is the block_storage parameter.

That strategy was problematic: it failed to increment indices
appropriately because it did not skip past the block_storage parameter
to begin with; it also failed to increment indices appropriately in
certain cases where the flag and error parameters appeared in an
unexpected order (flag first, then error).

Here, the process is made more robust.  Now, the indices which
correspond to the block_storage, flag, and error parameters are filtered
out to begin with.  The remaining indices can then be indexed into
using the index into the result tuple (or 0 if the result is not a
tuple).

rdar://80847020

Author: nate-chandler

lib/SILGen/SILGenThunk.cpp

@@ -255,19 +255,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);
@@ -276,8 +280,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());
@@ -372,31 +376,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())
+  }
+}