Merge pull request #11173 from aschwaighofer/swift-4.0-branch-irgen_fix_partial_apply_generic_class_capture_empty_type

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

Author: aschwaighofer

lib/IRGen/GenFunc.cpp

@@ -698,6 +698,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
@@ -919,11 +954,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;
@@ -1095,13 +1128,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")