CapturePromotion: Handle closures with indirect results correctly (#3169)

The code was not taking them into account when translating between arguments of
the partial_apply and the parameter info.

rdar://26214143

Author: aschwaighofer

lib/SILOptimizer/IPO/CapturePromotion.cpp

@@ -332,19 +332,26 @@ static void
 computeNewArgInterfaceTypes(SILFunction *F,
                             IndicesSet &PromotableIndices,
                             SmallVectorImpl<SILParameterInfo> &OutTys) {
-  auto Parameters = F->getLoweredFunctionType()->getParameters();
+  auto FunctionTy = F->getLoweredFunctionType();
+  auto Parameters = FunctionTy->getParameters();
+  auto NumIndirectResults = FunctionTy->getNumIndirectResults();
 
   DEBUG(llvm::dbgs() << "Preparing New Args!\n");
 
   // For each parameter in the old function...
   for (unsigned Index : indices(Parameters)) {
     auto &param = Parameters[Index];
 
+    // The PromotableIndices index is expressed as the argument index (num
+    // indirect result + param index). Add back the num indirect results to get
+    // the arg index when working with PromotableIndices.
+    unsigned ArgIndex = Index + NumIndirectResults;
+
     DEBUG(llvm::dbgs() << "Index: " << Index << "; PromotableIndices: "
-          << (PromotableIndices.count(Index)?"yes":"no")
+          << (PromotableIndices.count(ArgIndex)?"yes":"no")
           << " Param: "; param.dump());
 
-    if (!PromotableIndices.count(Index)) {
+    if (!PromotableIndices.count(ArgIndex)) {
       OutTys.push_back(param);
       continue;
     }
@@ -377,8 +384,14 @@ static std::string getSpecializedName(SILFunction *F,
   CanSILFunctionType FTy = F->getLoweredFunctionType();
 
   ArrayRef<SILParameterInfo> Parameters = FTy->getParameters();
+  auto NumIndirectResults = FTy->getNumIndirectResults();
+
   for (unsigned Index : indices(Parameters)) {
-    if (!PromotableIndices.count(Index))
+    // The PromotableIndices index is expressed as the argument index (num
+    // indirect result + param index). Add back the num indirect results to get
+    // the arg index when working with PromotableIndices.
+    unsigned ArgIndex = Index + NumIndirectResults;
+    if (!PromotableIndices.count(ArgIndex))
       continue;
     FSSM.setArgumentBoxToValue(Index);
   }
@@ -747,11 +760,6 @@ examineAllocBoxInst(AllocBoxInst *ABI, ReachabilityInfo &RI,
       if (signatureHasDependentTypes(*CalleeTy))
         return false;
 
-      // The code currently does not handle indirect results correctly.
-      // Conservatively, bail
-      if (CalleeTy->hasIndirectResults())
-        return false;
-
       auto closureType = PAI->getType().castTo<SILFunctionType>();
 
       // Calculate the index into the closure's argument list of the captured
@@ -912,19 +920,20 @@ processPartialApplyInst(PartialApplyInst *PAI, IndicesSet &PromotableIndices,
 
   // Populate the argument list for a new partial_apply instruction, taking into
   // consideration any captures.
-  auto CalleePInfo =
-      PAI->getCallee()->getType().castTo<SILFunctionType>()->getParameters();
-  auto PInfo = PAI->getType().castTo<SILFunctionType>()->getParameters();
-  unsigned FirstIndex = PInfo.size();
+  auto CalleeFunctionTy = PAI->getCallee()->getType().castTo<SILFunctionType>();
+  auto CalleePInfo = CalleeFunctionTy->getParameters();
+  unsigned FirstIndex =
+      PAI->getType().castTo<SILFunctionType>()->getNumSILArguments();
   unsigned OpNo = 1, OpCount = PAI->getNumOperands();
   SmallVector<SILValue, 16> Args;
+  auto NumIndirectResults = CalleeFunctionTy->getNumIndirectResults();
   while (OpNo != OpCount) {
     unsigned Index = OpNo - 1 + FirstIndex;
     if (PromotableIndices.count(Index)) {
       SILValue BoxValue = PAI->getOperand(OpNo);
       AllocBoxInst *ABI = cast<AllocBoxInst>(BoxValue);
 
-      SILParameterInfo CPInfo = CalleePInfo[Index];
+      SILParameterInfo CPInfo = CalleePInfo[Index - NumIndirectResults];
       assert(CPInfo.getSILType() == BoxValue->getType() &&
              "SILType of parameter info does not match type of parameter");
       // Cleanup the captured argument.

test/SILOptimizer/capture_promotion.sil

@@ -110,6 +110,29 @@ bb0:
   %15 = apply %12(%14, %13) : $@convention(thin) (Builtin.Word, @thin Int.Type) -> Int
   store %15 to %11a : $*Int
 
+// CHECK-NOT: function_ref @closure_indirect_result :
+// CHECK: [[CLOSURE_PROMOTE:%.*]] = function_ref @_TTSf2i_i_i_n__closure_indirect_result
+// CHECK-NOT: function_ref @closure_indirect_result :
+
+// The three strong retains are removed
+
+// The Foo variable is loaded from and retained, because it is a reference type
+// CHECK-NEXT: [[LOADFOO:%.*]] = load {{.*}} : $*Foo
+// CHECK-NEXT: strong_retain [[LOADFOO]] : $Foo
+
+// The Baz variable is loaded and retain_value'd, because it is a non-trivial
+// aggregate type
+// CHECK-NEXT: [[LOADBAZ:%.*]] = load {{.*}} : $*Baz
+// CHECK-NEXT: retain_value [[LOADBAZ:%.*]] : $Baz
+
+// The Int variable is loaded only, because it is trivial
+// CHECK-NEXT: [[LOADINT:%.*]] = load {{.*}} : $*Int
+
+// The partial apply has one value argument for each pair of arguments that was
+// previously used to capture and pass the variable by reference
+// CHECK-NEXT: {{.*}} = partial_apply [[CLOSURE_PROMOTE]]([[LOADFOO]], [[LOADBAZ]], [[LOADINT]])
+
+
   %17 = function_ref @closure_indirect_result : $@convention(thin) (@owned @box Foo, @owned @box Baz, @owned @box Int) -> @out Int
   strong_retain %1 : $@box Foo
   strong_retain %6 : $@box Baz