Merge pull request #18457 from atrick/fix-unexpected-box-use

Author: swift-ci

lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

@@ -63,70 +63,6 @@ static bool useCaptured(Operand *UI) {
   return true;
 }
 
-/// Given an apply or partial_apply, return the direct callee or
-/// nullptr if this is not a direct call.
-static FunctionRefInst *getDirectCallee(SILInstruction *Call) {
-  if (auto *Apply = dyn_cast<ApplyInst>(Call))
-    return dyn_cast<FunctionRefInst>(Apply->getCallee());
-  else
-    return dyn_cast<FunctionRefInst>(cast<PartialApplyInst>(Call)->getCallee());
-}
-
-/// Given an operand of a direct apply or partial_apply of a function,
-/// return the argument index of the parameter used in the body of the function
-/// to represent this operand.
-static size_t getArgIndexForOperand(Operand *O) {
-  assert(isa<ApplyInst>(O->getUser())
-         || isa<PartialApplyInst>(O->getUser())
-                && "Expected apply or partial_apply!");
-
-  auto OperandIndex = O->getOperandNumber();
-  assert(OperandIndex != 0 && "Operand cannot be the applied function!");
-
-  // The applied function is the first operand.
-  auto ArgIndex = OperandIndex - ApplyInst::getArgumentOperandNumber();
-
-  if (auto *Apply = dyn_cast<ApplyInst>(O->getUser())) {
-    assert(Apply->getSubstCalleeConv().getNumSILArguments()
-               == Apply->getArguments().size()
-           && "Expected all arguments to be supplied!");
-    (void)Apply;
-  } else {
-    auto *PartialApply = cast<PartialApplyInst>(O->getUser());
-    auto fnConv = PartialApply->getSubstCalleeConv();
-    auto ArgCount = PartialApply->getArguments().size();
-    assert(ArgCount <= fnConv.getNumParameters());
-    ArgIndex += (fnConv.getNumSILArguments() - ArgCount);
-  }
-
-  return ArgIndex;
-}
-
-/// Given an operand of a direct apply or partial_apply of a function,
-/// return the parameter used in the body of the function to represent
-/// this operand.
-static SILArgument *getParameterForOperand(SILFunction *F, Operand *O) {
-  assert(F && !F->empty() && "Expected a function with a body!");
-
-  auto &Entry = F->front();
-  size_t ArgIndex = getArgIndexForOperand(O);
-  assert(ArgIndex >= F->getConventions().getSILArgIndexOfFirstParam());
-
-  return Entry.getArgument(ArgIndex);
-}
-
-/// Return a pointer to the SILFunction called by Call if we can
-/// determine which function that is, and we have a body for that
-/// function. Otherwise return nullptr.
-static SILFunction *getFunctionBody(SILInstruction *Call) {
-  if (auto *FRI = getDirectCallee(Call))
-    if (auto *F = FRI->getReferencedFunction())
-      if (!F->empty())
-        return F;
-
-  return nullptr;
-}
-
 //===----------------------------------------------------------------------===//
 //                 Liveness for alloc_box Promotion
 //===----------------------------------------------------------------------===//
@@ -269,16 +205,16 @@ static bool partialApplyEscapes(SILValue V, bool examineApply);
 
 /// Could this operand to an apply escape that function by being
 /// stored or returned?
-static bool partialApplyArgumentEscapes(Operand *O) {
-  SILFunction *F = getFunctionBody(O->getUser());
+static bool applyArgumentEscapes(FullApplySite Apply, Operand *O) {
+  SILFunction *F = Apply.getReferencedFunction();
   // If we cannot examine the function body, assume the worst.
-  if (!F)
+  if (!F || F->empty())
     return true;
 
   // Check the uses of the operand, but do not recurse down into other
   // apply instructions.
-  auto Param = SILValue(getParameterForOperand(F, O));
-  return partialApplyEscapes(Param, /* examineApply = */ false);
+  auto calleeArg = F->getArgument(Apply.getCalleeArgIndex(*O));
+  return partialApplyEscapes(calleeArg, /* examineApply = */ false);
 }
 
 static bool partialApplyEscapes(SILValue V, bool examineApply) {
@@ -301,9 +237,7 @@ static bool partialApplyEscapes(SILValue V, bool examineApply) {
       continue;
     }
 
-    if (auto AI = dyn_cast<ApplyInst>(User)) {
-      ApplySite Apply(AI);
-
+    if (auto Apply = FullApplySite::isa(User)) {
       // Applying a function does not cause the function to escape.
       if (!Apply.isArgumentOperand(*Op))
         continue;
@@ -315,7 +249,7 @@ static bool partialApplyEscapes(SILValue V, bool examineApply) {
 
       // Optionally drill down into an apply to see if the operand is
       // captured in or returned from the apply.
-      if (examineApply && !partialApplyArgumentEscapes(Op))
+      if (examineApply && !applyArgumentEscapes(Apply, Op))
         continue;
     }
 
@@ -349,16 +283,16 @@ static SILInstruction *findUnexpectedBoxUse(SILValue Box,
 /// disqualify it from being promoted to a stack location.  Return
 /// true if this partial apply will not block our promoting the box.
 static bool checkPartialApplyBody(Operand *O) {
-  SILFunction *F = getFunctionBody(O->getUser());
+  SILFunction *F = ApplySite(O->getUser()).getReferencedFunction();
   // If we cannot examine the function body, assume the worst.
-  if (!F)
+  if (!F || F->empty())
     return false;
 
   // We don't actually use these because we're not recursively
   // rewriting the partial applies we find.
   llvm::SmallVector<Operand *, 1> PromotedOperands;
-  auto Param = SILValue(getParameterForOperand(F, O));
-  return !findUnexpectedBoxUse(Param, /* examinePartialApply = */ false,
+  auto calleeArg = F->getArgument(ApplySite(O->getUser()).getCalleeArgIndex(*O));
+  return !findUnexpectedBoxUse(calleeArg, /* examinePartialApply = */ false,
                                /* inAppliedFunction = */ true,
                                PromotedOperands);
 }
@@ -802,7 +736,7 @@ void PromotedParamCloner::visitProjectBoxInst(ProjectBoxInst *Inst) {
 /// references.
 static PartialApplyInst *
 specializePartialApply(PartialApplyInst *PartialApply,
-                       ArgIndexList &PromotedArgIndices,
+                       ArgIndexList &PromotedCalleeArgIndices,
                        AllocBoxToStackState &pass) {
   auto *FRI = cast<FunctionRefInst>(PartialApply->getCallee());
   assert(FRI && "Expected a direct partial_apply!");
@@ -813,7 +747,8 @@ specializePartialApply(PartialApplyInst *PartialApply,
   if (PartialApply->getFunction()->isSerialized())
     Serialized = IsSerializable;
 
-  std::string ClonedName = getClonedName(F, Serialized, PromotedArgIndices);
+  std::string ClonedName =
+    getClonedName(F, Serialized, PromotedCalleeArgIndices);
 
   auto &M = PartialApply->getModule();
 
@@ -823,7 +758,8 @@ specializePartialApply(PartialApplyInst *PartialApply,
     ClonedFn = PrevFn;
   } else {
     // Clone the function the existing partial_apply references.
-    PromotedParamCloner Cloner(F, Serialized, PromotedArgIndices, ClonedName);
+    PromotedParamCloner Cloner(F, Serialized, PromotedCalleeArgIndices,
+                               ClonedName);
     Cloner.populateCloned();
     ClonedFn = Cloner.getCloned();
     pass.T->notifyAddFunction(ClonedFn, F);
@@ -836,8 +772,8 @@ specializePartialApply(PartialApplyInst *PartialApply,
 
   // Promote the arguments that need promotion.
   for (auto &O : PartialApply->getArgumentOperands()) {
-    auto ArgIndex = getArgIndexForOperand(&O);
-    if (!count(PromotedArgIndices, ArgIndex)) {
+    auto CalleeArgIndex = ApplySite(O.getUser()).getCalleeArgIndex(O);
+    if (!count(PromotedCalleeArgIndices, CalleeArgIndex)) {
       Args.push_back(O.get());
       continue;
     }
@@ -886,18 +822,18 @@ static void rewritePartialApplies(AllocBoxToStackState &pass) {
   // Build a map from partial_apply to the indices of the operands
   // that will be promoted in our rewritten version.
   for (auto *O : pass.PromotedOperands) {
-    auto ArgIndexNumber = getArgIndexForOperand(O);
+    auto CalleeArgIndexNumber = ApplySite(O->getUser()).getCalleeArgIndex(*O);
 
     Indices.clear();
-    Indices.push_back(ArgIndexNumber);
+    Indices.push_back(CalleeArgIndexNumber);
 
     auto *PartialApply = cast<PartialApplyInst>(O->getUser());
     llvm::DenseMap<PartialApplyInst *, ArgIndexList>::iterator It;
     bool Inserted;
     std::tie(It, Inserted) = IndexMap.insert(std::make_pair(PartialApply,
                                                             Indices));
     if (!Inserted)
-      It->second.push_back(ArgIndexNumber);
+      It->second.push_back(CalleeArgIndexNumber);
   }
 
   // Clone the referenced function of each partial_apply, removing the

test/SILOptimizer/allocbox_to_stack.sil

@@ -1029,3 +1029,34 @@ bb3:
   unreachable
 
 }
+
+// Verify the case in which a box used by a partial apply which is then used by a try_apply.
+//
+// <rdar://problem/42815224> Swift CI: 2. Swift Source Compatibility Suite (master): Kitura project fails with compiler crash
+// CHECK-LABEL: sil @testPAUsedByTryApplyClosure : $@convention(thin) (@guaranteed { var Int }) -> () {
+// CHECK-LABEL: } // end sil function 'testPAUsedByTryApplyClosure'
+sil @testPAUsedByTryApplyClosure : $@convention(thin) (@guaranteed { var Int }) -> () {
+bb0(%0: ${ var Int }):
+  %v = tuple ()
+  return %v : $()
+}
+
+// CHECK-LABEL: sil private @testPAUsedByTryApply : $@convention(method) (@callee_guaranteed (@guaranteed @callee_guaranteed () -> ()) -> @error Error) -> () {
+// CHECK-LABEL: } // end sil function 'testPAUsedByTryApply'
+sil private @testPAUsedByTryApply : $@convention(method) (@callee_guaranteed (@guaranteed @callee_guaranteed () -> ()) -> @error Error) -> () {
+bb0(%0: $@callee_guaranteed (@guaranteed @callee_guaranteed () -> ()) -> @error Error):
+  %box = alloc_box ${ var Int }, var
+  %closure = function_ref @testPAUsedByTryApplyClosure : $@convention(thin) (@guaranteed { var Int }) -> ()
+  %pa = partial_apply [callee_guaranteed] %closure(%box) : $@convention(thin) (@guaranteed { var Int }) -> ()
+  try_apply %0(%pa) : $@callee_guaranteed (@guaranteed @callee_guaranteed () -> ()) -> @error Error, normal bb2, error bb3
+
+bb2(%result : $()):
+  br bb4
+
+bb3(%error : $Error):
+  br bb4
+
+bb4:
+  %v = tuple ()
+  return %v : $()
+}