IRGen: Fix partial applies of generic functions capturing the generic parameters in an argument

The code assumed if there is no context object but rather the captured argument
is reused as the context object that the parameter index for this argument is
going to be the last one. This is not true if there are empty types in the
parameter list.

rdar://33502272

Author: aschwaighofer

lib/IRGen/GenFunc.cpp

@@ -705,6 +705,41 @@ static CanType getArgumentLoweringType(CanType type,
   }
 }
 
+static bool isABIIgnoredParameterWithoutStorage(IRGenModule &IGM,
+                                                IRGenFunction &IGF,
+                                                CanSILFunctionType substType,
+                                                unsigned paramIdx) {
+  auto param = substType->getParameters()[paramIdx];
+  SILType argType = IGM.silConv.getSILType(param);
+  auto argLoweringTy =
+    getArgumentLoweringType(argType.getSwiftRValueType(), param);
+  auto &ti = IGF.getTypeInfoForLowered(argLoweringTy);
+  // Empty values don't matter.
+  return ti.getSchema().size() == 0 && !param.isFormalIndirect();
+}
+
+/// Find the parameter index for the one (assuming there was only one) partially
+/// applied argument ignoring empty types that are not passed as part of the
+/// ABI.
+static unsigned findSinglePartiallyAppliedParameterIndexIgnoringEmptyTypes(
+    IRGenFunction &IGF, CanSILFunctionType substType,
+    CanSILFunctionType outType) {
+  auto substParameters = substType->getParameters();
+  auto outParamters = outType->getParameters();
+  unsigned firstNonEmpty = -1U;
+  for (unsigned paramIdx = outParamters.size() ; paramIdx != substParameters.size(); ++paramIdx) {
+    bool isEmpty =
+        isABIIgnoredParameterWithoutStorage(IGF.IGM, IGF, substType, paramIdx);
+    assert((isEmpty || firstNonEmpty == -1U) && "Expect at most one partially "
+                                                "applied that is passed as an "
+                                                "ABI argument");
+    if (!isEmpty)
+      firstNonEmpty = paramIdx;
+  }
+  assert(firstNonEmpty != -1U);
+  return firstNonEmpty;
+}
+
 /// Emit the forwarding stub function for a partial application.
 ///
 /// If 'layout' is null, there is a single captured value of
@@ -924,11 +959,9 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
     // care for the former for the purpose of reconstructing polymorphic
     // parameters from regular arguments.
     if (!calleeHasContext) {
-      unsigned paramI = substType->getParameters().size() - 1;
-      assert(substType->getParameters().size() -
-                     outType->getParameters().size() ==
-                 1 &&
-             "Expect one partially applied argument");
+      unsigned paramI =
+          findSinglePartiallyAppliedParameterIndexIgnoringEmptyTypes(
+              subIGF, substType, outType);
       auto paramInfo = substType->getParameters()[paramI];
       auto &ti = IGM.getTypeInfoForLowered(paramInfo.getType());
       Explosion param;
@@ -1097,13 +1130,8 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
 
       // Skip empty parameters.
       while (origParamI < origType->getParameters().size()) {
-        auto param = substType->getParameters()[origParamI];
-        SILType argType = IGM.silConv.getSILType(param);
-        auto argLoweringTy =
-            getArgumentLoweringType(argType.getSwiftRValueType(), param);
-        auto &ti = subIGF.getTypeInfoForLowered(argLoweringTy);
-        // Empty values don't matter.
-        if (ti.getSchema().size() != 0 || param.isFormalIndirect())
+        if (!isABIIgnoredParameterWithoutStorage(IGM, subIGF, substType,
+                                                 origParamI))
           break;
         origParamI++;
       }

test/IRGen/partial_apply_forwarder.sil

@@ -70,7 +70,11 @@ public struct BaseProducer<T> : Q {
 
 public class WeakBox<T> {}
 
+public struct EmptyType {}
+
 sil hidden_external @takingQ : $@convention(thin) <τ_0_0 where  τ_0_0 : Q> (@owned WeakBox<τ_0_0>) -> ()
+sil hidden_external @takingQAndEmpty : $@convention(thin) <τ_0_0 where  τ_0_0 : Q> (@owned WeakBox<τ_0_0>, EmptyType) -> ()
+sil hidden_external @takingEmptyAndQ : $@convention(thin) <τ_0_0 where  τ_0_0 : Q> (EmptyType, @owned WeakBox<τ_0_0>) -> ()
 
 // CHECK-LABEL: define{{( protected)?}} swiftcc { i8*, %swift.refcounted* } @bind_polymorphic_param_from_context(%swift.opaque* noalias nocapture, %swift.type* %"\CF\84_0_1")
 // CHECK: entry:
@@ -103,6 +107,35 @@ bb0(%0 : $*τ_0_1):
   return %9 : $@callee_owned () -> ()
 }
 
+// CHECK-LABEL: define{{( protected)?}} swiftcc { i8*, %swift.refcounted* } @bind_polymorphic_param_from_context_2(%swift.opaque* noalias nocapture, %swift.type* %"\CF\84_0_1")
+// CHECK:   [[OBJ:%.*]] = call {{.*}} @swift_rt_swift_allocObject
+// CHECK:   [[WB:%.*]] = bitcast %swift.refcounted* [[OBJ]]
+// CHECK:   [[REF:%.*]] = bitcast {{.*}}* [[WB]] to %swift.refcounted*
+// CHECK:   [[CLOSURE:%.*]] = insertvalue { i8*, %swift.refcounted* } { i8* bitcast (void (%swift.refcounted*)* @_T015takingQAndEmptyTA to i8*), %swift.refcounted* undef }, %swift.refcounted* [[REF]], 1
+// CHECK: }
+sil public @bind_polymorphic_param_from_context_2 : $@convention(thin) <τ_0_1>(@in τ_0_1, EmptyType) -> @owned @callee_owned () -> () {
+bb0(%0 : $*τ_0_1, %2: $EmptyType):
+  %1 = alloc_ref $WeakBox<BaseProducer<τ_0_1>>
+  %8 = function_ref @takingQAndEmpty : $@convention(thin) <τ_0_0 where τ_0_0 : Q> (@owned WeakBox<τ_0_0>, EmptyType) -> ()
+  %9 = partial_apply %8<BaseProducer<τ_0_1>>(%1, %2) : $@convention(thin) <τ_0_0 where τ_0_0 : Q> (@owned WeakBox<τ_0_0>, EmptyType) -> ()
+  return %9 : $@callee_owned () -> ()
+}
+
+// CHECK-LABEL: define{{( protected)?}} swiftcc { i8*, %swift.refcounted* } @bind_polymorphic_param_from_context_3(%swift.opaque* noalias nocapture, %swift.type* %"\CF\84_0_1")
+// CHECK:   [[OBJ:%.*]] = call {{.*}} @swift_rt_swift_allocObject
+// CHECK:   [[WB:%.*]] = bitcast %swift.refcounted* [[OBJ]]
+// CHECK:   [[REF:%.*]] = bitcast {{.*}}* [[WB]] to %swift.refcounted*
+// CHECK:   [[CLOSURE:%.*]] = insertvalue { i8*, %swift.refcounted* } { i8* bitcast (void (%swift.refcounted*)* @_T015takingEmptyAndQTA to i8*), %swift.refcounted* undef }, %swift.refcounted* [[REF]], 1
+// CHECK: }
+
+sil public @bind_polymorphic_param_from_context_3 : $@convention(thin) <τ_0_1>(@in τ_0_1, EmptyType) -> @owned @callee_owned () -> () {
+bb0(%0 : $*τ_0_1, %2: $EmptyType):
+  %1 = alloc_ref $WeakBox<BaseProducer<τ_0_1>>
+  %8 = function_ref @takingEmptyAndQ : $@convention(thin) <τ_0_0 where τ_0_0 : Q> (EmptyType, @owned WeakBox<τ_0_0>) -> ()
+  %9 = partial_apply %8<BaseProducer<τ_0_1>>(%2, %1) : $@convention(thin) <τ_0_0 where τ_0_0 : Q> (EmptyType, @owned WeakBox<τ_0_0>) -> ()
+  return %9 : $@callee_owned () -> ()
+}
+
 // CHECK-LABEL: define{{( protected)?}} swiftcc void @bind_polymorphic_param_from_forwarder_parameter(%swift.opaque* noalias nocapture, %swift.type* %"\CF\84_0_1")
 // CHECK: entry:
 // CHECK:   [[BPTY:%.*]] = call %swift.type* @_T023partial_apply_forwarder12BaseProducerVMa(%swift.type* %"\CF\84_0_1")