[SIL] Only visit final partial_apply [on_stack] lifetime ends.

lib/SIL/IR/SILInstruction.cpp

@@ -15,14 +15,17 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/SIL/SILInstruction.h"
-#include "swift/Basic/Assertions.h"
 #include "swift/Basic/AssertImplements.h"
+#include "swift/Basic/Assertions.h"
 #include "swift/Basic/Unicode.h"
 #include "swift/Basic/type_traits.h"
 #include "swift/SIL/ApplySite.h"
 #include "swift/SIL/DynamicCasts.h"
+#include "swift/SIL/InstWrappers.h"
 #include "swift/SIL/InstructionUtils.h"
+#include "swift/SIL/NodeDatastructures.h"
 #include "swift/SIL/OwnershipUtils.h"
+#include "swift/SIL/PrunedLiveness.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILCloner.h"
 #include "swift/SIL/SILDebugScope.h"
@@ -1828,6 +1831,134 @@ bool SILInstruction::maySuspend() const {
   return false;
 }
 
+static SILValue lookThroughOwnershipAndForwardingInsts(SILValue value) {
+  auto current = value;
+  while (true) {
+    if (auto *inst = current->getDefiningInstruction()) {
+      switch (inst->getKind()) {
+      case SILInstructionKind::MoveValueInst:
+      case SILInstructionKind::CopyValueInst:
+      case SILInstructionKind::BeginBorrowInst:
+        current = inst->getOperand(0);
+        continue;
+      default:
+        break;
+      }
+      auto forward = ForwardingOperation(inst);
+      Operand *op = nullptr;
+      if (forward && (op = forward.getSingleForwardingOperand())) {
+        current = op->get();
+        continue;
+      }
+    } else if (auto *result = SILArgument::isTerminatorResult(current)) {
+      auto *op = result->forwardedTerminatorResultOperand();
+      if (!op) {
+        break;
+      }
+      current = op->get();
+      continue;
+    }
+    break;
+  }
+  return current;
+}
+
+bool
+PartialApplyInst::visitOnStackLifetimeEnds(
+                             llvm::function_ref<bool (Operand *)> func) const {
+  assert(getFunction()->hasOwnership()
+         && isOnStack()
+         && "only meaningful for OSSA stack closures");
+  bool noUsers = true;
+
+  auto *function = getFunction();
+
+  SmallVector<SILBasicBlock *, 32> discoveredBlocks;
+  SSAPrunedLiveness liveness(function, &discoveredBlocks);
+  liveness.initializeDef(this);
+
+  StackList<SILValue> values(function);
+  values.push_back(this);
+
+  while (!values.empty()) {
+    auto value = values.pop_back_val();
+    for (auto *use : value->getUses()) {
+      if (!use->isConsuming()) {
+        if (auto *cvi = dyn_cast<CopyValueInst>(use->getUser())) {
+          values.push_back(cvi);
+        }
+        continue;
+      }
+      noUsers = false;
+      if (isa<DestroyValueInst>(use->getUser())) {
+        liveness.updateForUse(use->getUser(), /*lifetimeEnding=*/true);
+        continue;
+      }
+      auto forward = ForwardingOperand(use);
+      if (!forward) {
+        // There shouldn't be any non-forwarding consumptions of a nonescaping
+        // partial_apply that don't forward it along, aside from destroy_value.
+        //
+        // On-stack partial_apply cannot be cloned, so it should never be used
+        // by a BranchInst.
+        //
+        // This is a fatal error because it performs SIL verification that is
+        // not separately checked in the verifier. It is the only check that
+        // verifies the structural requirements of on-stack partial_apply uses.
+        if (lookThroughOwnershipAndForwardingInsts(use->get()) !=
+            SILValue(this)) {
+          // Consumes of values which aren't "essentially" the
+          //   partial_apply [on_stack]
+          // are okay.  For example, a not-on_stack partial_apply that captures
+          // it.
+          continue;
+        }
+        llvm::errs() << "partial_apply [on_stack] use:\n";
+        auto *user = use->getUser();
+        user->printInContext(llvm::errs());
+        if (isa<BranchInst>(user)) {
+          llvm::report_fatal_error("partial_apply [on_stack] cannot be cloned");
+        }
+        llvm::report_fatal_error("partial_apply [on_stack] must be directly "
+                                 "forwarded to a destroy_value");
+      }
+      forward.visitForwardedValues([&values](auto value) {
+        values.push_back(value);
+        return true;
+      });
+    }
+  }
+  PrunedLivenessBoundary boundary;
+  liveness.computeBoundary(boundary);
+
+  for (auto *inst : boundary.lastUsers) {
+    // Only destroy_values were added to liveness, so only destroy_values can be
+    // the last users.
+    auto *dvi = cast<DestroyValueInst>(inst);
+    auto keepGoing = func(&dvi->getOperandRef());
+    if (!keepGoing) {
+      return false;
+    }
+  }
+  return !noUsers;
+}
+
+namespace swift::test {
+FunctionTest PartialApplyPrintOnStackLifetimeEnds(
+    "partial_apply_print_on_stack_lifetime_ends",
+    [](auto &function, auto &arguments, auto &test) {
+      auto *inst = arguments.takeInstruction();
+      auto *pai = cast<PartialApplyInst>(inst);
+      function.print(llvm::outs());
+      auto result = pai->visitOnStackLifetimeEnds([](auto *operand) {
+        operand->print(llvm::outs());
+        return true;
+      });
+      const char *resultString = result ? "true" : "false";
+      llvm::outs() << "returned: " << resultString << "\n";
+    });
+} // end namespace swift::test
+
 static bool
 visitRecursivelyLifetimeEndingUses(
   SILValue i, bool &noUsers,
@@ -1869,41 +2000,6 @@ visitRecursivelyLifetimeEndingUses(
   return true;
 }
 
-bool
-PartialApplyInst::visitOnStackLifetimeEnds(
-                             llvm::function_ref<bool (Operand *)> func) const {
-  assert(getFunction()->hasOwnership()
-         && isOnStack()
-         && "only meaningful for OSSA stack closures");
-  bool noUsers = true;
-
-  auto visitUnknownUse = [](Operand *unknownUse) {
-    // There shouldn't be any dead-end consumptions of a nonescaping
-    // partial_apply that don't forward it along, aside from destroy_value.
-    //
-    // On-stack partial_apply cannot be cloned, so it should never be used by a
-    // BranchInst.
-    //
-    // This is a fatal error because it performs SIL verification that is not
-    // separately checked in the verifier. It is the only check that verifies
-    // the structural requirements of on-stack partial_apply uses.
-    llvm::errs() << "partial_apply [on_stack] use:\n";
-    auto *user = unknownUse->getUser();
-    user->printInContext(llvm::errs());
-    if (isa<BranchInst>(user)) {
-      llvm::report_fatal_error("partial_apply [on_stack] cannot be cloned");
-    }
-    llvm::report_fatal_error("partial_apply [on_stack] must be directly "
-                             "forwarded to a destroy_value");
-    return false;
-  };
-  if (!visitRecursivelyLifetimeEndingUses(this, noUsers, func,
-                                          visitUnknownUse)) {
-    return false;
-  }
-  return !noUsers;
-}
-
 // FIXME: Rather than recursing through all results, this should only recurse
 // through ForwardingInstruction and OwnershipTransitionInstruction and the
 // client should prove that any other uses cannot be upstream from a consume of

lib/SIL/Utils/OwnershipLiveness.cpp

@@ -414,7 +414,7 @@ static FunctionTest LinearLivenessTest("linear_liveness", [](auto &function,
 // - the computed pruned liveness
 // - the liveness boundary
 static FunctionTest
-    InteriorLivenessTest("interior-liveness",
+    InteriorLivenessTest("interior_liveness",
                          [](auto &function, auto &arguments, auto &test) {
                            SILValue value = arguments.takeValue();
                            function.print(llvm::outs());

test/SIL/Instruction/partial_apply_on_stack_lifetime_ends.sil

@@ -0,0 +1,40 @@
+// RUN: %target-sil-opt \
+// RUN:     -test-runner \
+// RUN:     %s \
+// RUN:     -o /dev/null \
+// RUN: 2>&1 | %FileCheck %s
+
+sil_stage canonical
+
+class C {}
+
+sil @borrowC : $@convention(thin) (@guaranteed C) -> ()
+
+// CHECK-LABEL: begin running test {{.*}} on copied_partial_apply
+// CHECK-LABEL: sil [ossa] @copied_partial_apply : {{.*}} {
+// CHECK:       bb0([[C:%[^,]+]] :
+// CHECK:         [[BORROW_C:%[^,]+]] = function_ref @borrowC
+// CHECK:         [[PA:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[BORROW_C]]([[C]])
+// CHECK:         [[PA2:%[^,]+]] = copy_value [[PA]]
+// CHECK:         destroy_value [[PA]]
+// CHECK:         destroy_value [[PA2]]
+// CHECK:         destroy_value [[C]]
+// CHECK-LABEL: } // end sil function 'copied_partial_apply'
+// CHECK:       Operand.
+// CHECK:       Owner:   destroy_value [[PA2]]
+// CHECK:       returned: true
+// CHECK-LABEL: end running test {{.*}} on copied_partial_apply
+sil [ossa] @copied_partial_apply : $@convention(thin) (@owned C) -> () {
+entry(%c: @owned $C):
+  specify_test "partial_apply_print_on_stack_lifetime_ends %pa"
+
+  %callee = function_ref @borrowC : $@convention(thin) (@guaranteed C) -> ()
+  %pa = partial_apply [callee_guaranteed] [on_stack] %callee(%c) : $@convention(thin) (@guaranteed C) -> ()
+  %pa2 = copy_value %pa
+  destroy_value %pa
+
+  destroy_value %pa2
+  destroy_value %c
+  %retval = tuple ()
+  return %retval
+}

test/SILOptimizer/canonicalize_ossa_lifetime_unit.sil

@@ -162,9 +162,9 @@ exit(%phi : @owned $C, %typhi : $S):
 sil @empty : $@convention(thin) () -> () {
 [global: ]
 bb0:
-  %0 = tuple ()                                   
-  return %0 : $()                                 
-} 
+  %0 = tuple ()
+  return %0 : $()
+}
 
 // Even though the apply of %empty is not a deinit barrier, verify that the
 // destroy is not hoisted, because MoS is move-only.
@@ -568,16 +568,16 @@ entry(%c1 : @owned $C):
 // CHECK:       bb0([[C1:%[^,]+]] : @owned $C):
 // CHECK:         [[TAKE_C:%[^,]+]] = function_ref @takeC
 // CHECK:         [[BARRIER:%[^,]+]] = function_ref @barrier
-// CHECK:         cond_br undef, [[LEFT]], [[RIGHT]]                         
-// CHECK:       [[LEFT]]:                                              
+// CHECK:         cond_br undef, [[LEFT]], [[RIGHT]]
+// CHECK:       [[LEFT]]:
 // CHECK:         [[M:%[^,]+]] = move_value [[C1]]
 // CHECK:         apply [[TAKE_C]]([[M]])
-// CHECK:         br [[EXIT]]                                          
-// CHECK:       [[RIGHT]]:                                              
+// CHECK:         br [[EXIT]]
+// CHECK:       [[RIGHT]]:
 // CHECK:         apply [[BARRIER]]()
 // CHECK:         destroy_value [[C1]]
-// CHECK:         br [[EXIT]]                                          
-// CHECK:       [[EXIT]]:                                              
+// CHECK:         br [[EXIT]]
+// CHECK:       [[EXIT]]:
 // CHECK:         apply [[BARRIER]]()
 // CHECK-LABEL: } // end sil function 'lexical_end_at_end_2'
 // CHECK-LABEL: end running test {{.*}} on lexical_end_at_end_2: canonicalize_ossa_lifetime
@@ -884,3 +884,37 @@ die:
   apply undef(%reload) : $@convention(thin) (@guaranteed C) -> ()
   unreachable
 }
+
+// The destroy of a value must not be hoisted over a destroy of a copy of a
+// partial_apply [on_stack] which captures the value.
+// CHECK-LABEL: begin running test {{.*}} on destroy_after_pa_copy_destroy
+// CHECK-LABEL: sil [ossa] @destroy_after_pa_copy_destroy : {{.*}} {
+// CHECK:       bb0([[A:%[^,]+]] :
+// CHECK:         [[C:%[^,]+]] = move_value [[A]]
+// CHECK:         [[BORROW_C:%[^,]+]] = function_ref @borrowC
+// CHECK:         [[PA:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[BORROW_C]]([[C]])
+// CHECK:         [[PA2:%[^,]+]] = copy_value [[PA]]
+// CHECK:         destroy_value [[PA]]
+// CHECK:         destroy_value [[PA2]]
+// CHECK:         destroy_value [[C]]
+// CHECK:         apply undef()
+// CHECK-LABEL: } // end sil function 'destroy_after_pa_copy_destroy'
+// CHECK-LABEL: end running test {{.*}} on destroy_after_pa_copy_destroy
+sil [ossa] @destroy_after_pa_copy_destroy : $@convention(thin) (@owned C) -> () {
+entry(%a: @owned $C):
+  // To defeat deinit barriers (apply undef).
+  %c = move_value %a
+  specify_test "canonicalize_ossa_lifetime true false true %c"
+
+  %callee = function_ref @borrowC : $@convention(thin) (@guaranteed C) -> ()
+  %pa = partial_apply [callee_guaranteed] [on_stack] %callee(%c) : $@convention(thin) (@guaranteed C) -> ()
+  %pa2 = copy_value %pa
+  destroy_value %pa
+
+  destroy_value %pa2
+  // To defeat ignoredByDestroyHoisting.
+  apply undef() : $@convention(thin) () -> ()
+  destroy_value %c
+  %retval = tuple ()
+  return %retval
+}