[sil-capture-promotion] Properly handle generic SILBoxTypes

lib/SILOptimizer/IPO/CapturePromotion.cpp

@@ -350,8 +350,7 @@ computeNewArgInterfaceTypes(SILFunction *F,
     assert(paramBoxTy->getLayout()->getFields().size() == 1
            && "promoting compound box not implemented yet");
     auto paramBoxedTy = paramBoxTy->getFieldType(F->getModule(), 0);
-    auto &paramTL = F->getModule().Types.getTypeLowering(paramBoxedTy);
-
+    auto &paramTL = F->getTypeLowering(paramBoxedTy);
     ParameterConvention convention;
     if (paramTL.isFormallyPassedIndirectly()) {
       convention = ParameterConvention::Indirect_In;
@@ -863,8 +862,12 @@ processPartialApplyInst(PartialApplyInst *PAI, IndicesSet &PromotableIndices,
   // Populate the argument list for a new partial_apply instruction, taking into
   // consideration any captures.
   auto CalleeFunctionTy = PAI->getCallee()->getType().castTo<SILFunctionType>();
-  SILFunctionConventions calleeConv(CalleeFunctionTy, M);
-  auto CalleePInfo = CalleeFunctionTy->getParameters();
+  auto SubstCalleeFunctionTy = CalleeFunctionTy;
+  if (PAI->hasSubstitutions())
+    SubstCalleeFunctionTy =
+        CalleeFunctionTy->substGenericArgs(M, PAI->getSubstitutions());
+  SILFunctionConventions calleeConv(SubstCalleeFunctionTy, M);
+  auto CalleePInfo = SubstCalleeFunctionTy->getParameters();
   SILFunctionConventions paConv(PAI->getType().castTo<SILFunctionType>(), M);
   unsigned FirstIndex = paConv.getNumSILArguments();
   unsigned OpNo = 1, OpCount = PAI->getNumOperands();
@@ -879,10 +882,6 @@ processPartialApplyInst(PartialApplyInst *PAI, IndicesSet &PromotableIndices,
       SILParameterInfo CPInfo = CalleePInfo[Index - NumIndirectResults];
       assert(calleeConv.getSILType(CPInfo) == BoxValue->getType()
              && "SILType of parameter info does not match type of parameter");
-      // Cleanup the captured argument.
-      releasePartialApplyCapturedArg(B, PAI->getLoc(), BoxValue,
-                                     CPInfo);
-
       // Load and copy from the address value, passing the result as an argument
       // to the new closure.
       SILValue Addr;
@@ -906,6 +905,9 @@ processPartialApplyInst(PartialApplyInst *PAI, IndicesSet &PromotableIndices,
       auto &typeLowering = M.getTypeLowering(Addr->getType());
       Args.push_back(
         typeLowering.emitLoadOfCopy(B, PAI->getLoc(), Addr, IsNotTake));
+      // Cleanup the captured argument.
+      releasePartialApplyCapturedArg(B, PAI->getLoc(), BoxValue,
+                                     CPInfo);
       ++NumCapturesPromoted;
     } else {
       Args.push_back(PAI->getOperand(OpNo));

test/SILOptimizer/capture_promotion_generic_context.sil

@@ -67,3 +67,56 @@ entry(%0 : $*T, %1 : $*U, %2 : $Int):
   %k = partial_apply %f<U>(%b) : $@convention(thin) <V> (@in V, <τ_0_0> { var τ_0_0 } <Int>) -> Int
   return %k : $@callee_owned (@in U) -> Int
 }
+
+enum E<X> {
+  case None
+  case Some(X)
+}
+
+struct R<T> {
+}
+
+// Check that the capture promotion took place and the function now
+// take argument of a type  E<(R<T>) -> Builtin.Int32>, which used
+// to be a slot inside a box.
+// CHECK-LABEL: sil @_T023generic_promotable_box2Tf2nni_n : $@convention(thin) <T> (@in R<T>, @in Builtin.Int32, @owned E<(R<T>) -> Builtin.Int32>) -> () 
+// CHECK:       bb0(%0 : $*R<T>, %1 : $*Builtin.Int32, %2 : $E<(R<T>) -> Builtin.Int32>):
+// CHECK-NOT:     project_box
+// CHECK:         switch_enum %2 : $E<(R<T>) -> Builtin.Int32>
+// CHECK-NOT:     project_box
+// CHECK:       } // end sil function '_T023generic_promotable_box2Tf2nni_n'
+sil @generic_promotable_box2 : $@convention(thin) <T> (@in R<T>, @in Int, <τ_0_0> { var  E<(R<τ_0_0>) -> Int> } <T>) -> () {
+entry(%0: $*R<T>, %1 : $*Int, %b : $<τ_0_0> { var E< (R<τ_0_0>)->Int > } <T>):
+  %a = project_box %b : $<τ_0_0> { var E<(R<τ_0_0>)->Int> } <T>, 0
+  %e = load %a : $*E<(R<T>)->Int>
+  switch_enum %e : $E<(R<T>)->Int>, case #E.Some!enumelt.1 : bb1, default bb2
+bb1(%f : $@callee_owned (@in R<T>) -> @out Int):
+  %t = tuple ()
+  apply %f(%1, %0) : $@callee_owned (@in R<T>) -> @out Int
+  br exit
+bb2:
+  br exit
+exit:
+  %r = tuple ()
+  return %r : $()
+}
+
+// Check that alloc_box was eliminated and a specialized version of the
+// closure is invoked.
+// CHECK-LABEL: sil @call_generic_promotable_box_from_different_generic2
+// CHECK:       bb0(%0 : $*R<T>, %1 : $*E<(R<U>) -> Builtin.Int32>, %2 : $*Builtin.Int32):
+// CHECK:         %3 = load %1 : $*E<(R<U>) -> Builtin.Int32>
+// CHECK:         [[F:%.*]] = function_ref @_T023generic_promotable_box2Tf2nni_n : $@convention(thin) <τ_0_0> (@in R<τ_0_0>, @in Builtin.Int32, @owned E<(R<τ_0_0>) -> Builtin.Int32>) -> ()
+// CHECK-NEXT:    [[CLOSURE:%.*]] = partial_apply [[F]]<U>(%2, %3)
+// CHECK-NEXT:    return [[CLOSURE]]
+
+sil @call_generic_promotable_box_from_different_generic2 : $@convention(thin) <T, U: P> (@in R<T>, @in E<(R<U>)->Int>, @in Int) -> @owned @callee_owned (@in R<U>) -> () {
+entry(%0 : $*R<T>, %1 : $*E<(R<U>)->Int>, %2 : $*Int):
+  destroy_addr %0 : $*R<T>
+  %f = function_ref @generic_promotable_box2 : $@convention(thin) <V> (@in R<V>, @in Int, <τ_0_0> { var E<(R<τ_0_0>)->Int> } <V>) -> ()
+  %b = alloc_box $<τ_0_0> { var E<(R<τ_0_0>)->Int> } <U>
+  %a = project_box %b : $<τ_0_0> { var E<(R<τ_0_0>)->Int> } <U>, 0
+  copy_addr [take] %1 to [initialization] %a : $*E<(R<U>)->Int>
+  %k = partial_apply %f<U>(%2, %b) : $@convention(thin) <V> (@in R<V>, @in Int, <τ_0_0> { var E<(R<τ_0_0>)->Int> } <V>) -> ()
+  return %k : $@callee_owned (@in R<U>) -> ()
+}