[ShrinkBorrowScope] Leave under rewritten copies.

lib/SILOptimizer/UtilityPasses/UnitTestRunner.cpp

@@ -13,7 +13,7 @@
 // TO ADD A NEW TEST, search in this file for [new_tests]
 //
 // Provides a mechanism for doing faux unit tests.  The idea is to emulate the
-// basic function of calling a function and validating its results/effects.
+// basic functionality of calling a function and validating its results/effects.
 //
 // This is done via the test_specification instruction.  Using one or more
 // instances of it in your function, you can specify which test (i.e. UnitTest
@@ -75,6 +75,7 @@
 #include "swift/SILOptimizer/Analysis/BasicCalleeAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
+#include "swift/SILOptimizer/Utils/CanonicalizeBorrowScope.h"
 #include "swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "swift/SILOptimizer/Utils/InstructionDeleter.h"
@@ -274,6 +275,36 @@ struct IsDeinitBarrierTest : UnitTest {
   }
 };
 
+struct ShrinkBorrowScopeTest : UnitTest {
+  ShrinkBorrowScopeTest(UnitTestRunner *pass) : UnitTest(pass) {}
+  void invoke(Arguments &arguments) override {
+    auto instruction = arguments.takeValue();
+    auto expected = arguments.takeBool();
+    auto *bbi = cast<BeginBorrowInst>(instruction);
+    auto *analysis = getAnalysis<BasicCalleeAnalysis>();
+    SmallVector<CopyValueInst *, 4> modifiedCopyValueInsts;
+    InstructionDeleter deleter(
+        InstModCallbacks().onDelete([&](auto *instruction) {
+          llvm::errs() << "DELETED:\n";
+          instruction->dump();
+        }));
+    auto shrunk =
+        shrinkBorrowScope(*bbi, deleter, analysis, modifiedCopyValueInsts);
+    unsigned index = 0;
+    for (auto *cvi : modifiedCopyValueInsts) {
+      auto expectedCopy = arguments.takeValue();
+      llvm::errs() << "rewritten copy " << index << ":\n";
+      llvm::errs() << "expected:\n";
+      expectedCopy->print(llvm::errs());
+      llvm::errs() << "got:\n";
+      cvi->dump();
+      assert(cvi == expectedCopy);
+      ++index;
+    }
+    assert(expected == shrunk && "didn't shrink expectedly!?");
+  }
+};
+
 /// [new_tests] Add the new UnitTest subclass above this line.
 
 class UnitTestRunner : public SILFunctionTransform {
@@ -314,6 +345,7 @@ class UnitTestRunner : public SILFunctionTransform {
     ADD_UNIT_TEST_SUBCLASS(
         "pruned-liveness-boundary-with-list-of-last-users-insertion-points",
         PrunedLivenessBoundaryWithListOfLastUsersInsertionPointsTest)
+    ADD_UNIT_TEST_SUBCLASS("shrink-borrow-scope", ShrinkBorrowScopeTest)
     ADD_UNIT_TEST_SUBCLASS("is-deinit-barrier", IsDeinitBarrierTest)
     /// [new_tests] Add the new mapping from string to subclass above this line.
 

lib/SILOptimizer/Utils/ShrinkBorrowScope.cpp

@@ -347,6 +347,7 @@ class Rewriter final {
   bool run();
 
 private:
+  EndBorrowInst *findPreexistingEndBorrow(SILInstruction *instruction);
   bool createEndBorrow(SILInstruction *insertionPoint);
 };
 
@@ -422,12 +423,31 @@ bool Rewriter::run() {
   return madeChange;
 }
 
-bool Rewriter::createEndBorrow(SILInstruction *insertionPoint) {
-  if (auto *ebi = dyn_cast<EndBorrowInst>(insertionPoint)) {
-    if (llvm::find(uses.ends, insertionPoint) != uses.ends.end()) {
-      reusedEndBorrowInsts.insert(insertionPoint);
-      return false;
+EndBorrowInst *
+Rewriter::findPreexistingEndBorrow(SILInstruction *insertionPoint) {
+  for (auto *instruction = insertionPoint; instruction;
+       instruction = instruction->getNextInstruction()) {
+    if (auto *ebi = dyn_cast<EndBorrowInst>(instruction)) {
+      if (llvm::find(uses.ends, ebi) != uses.ends.end())
+        return ebi;
+    }
+    if (auto *cvi = dyn_cast<CopyValueInst>(instruction)) {
+      if (llvm::is_contained(barriers.copies, cvi)) {
+        continue;
+      }
     }
+    /// Otherwise, this is an "interesting" instruction.  We want to record that
+    /// we were able to hoist the end of the borrow scope over it, so we stop
+    /// looking for the preexisting end_borrow.
+    return nullptr;
+  }
+  return nullptr;
+}
+
+bool Rewriter::createEndBorrow(SILInstruction *insertionPoint) {
+  if (auto *ebi = findPreexistingEndBorrow(insertionPoint)) {
+    reusedEndBorrowInsts.insert(ebi);
+    return false;
   }
   auto builder = SILBuilderWithScope(insertionPoint);
   builder.createEndBorrow(

test/SILOptimizer/copy_propagation_opaque.sil

@@ -413,8 +413,8 @@ bb0:
 
 // CHECK-TRACE-LABEL: *** CopyPropagation: testBorrowCopy
 // CHECK-TRACE:        Canonicalizing: %1
-// CHECK-TRACE:        Removing   destroy_value %4 : $T
-// CHECK-TRACE:        Removing   %4 = copy_value %1 : $T
+// CHECK-TRACE:        Removing   destroy_value %3 : $T
+// CHECK-TRACE:        Removing   %3 = copy_value %1 : $T
 //
 // CHECK-LABEL: sil [ossa] @testBorrowCopy : {{.*}} {
 // CHECK-LABEL: bb0:

test/SILOptimizer/shrink_borrow_scope.sil

@@ -804,12 +804,12 @@ entry(%instance : @owned $C):
   %_ = apply %callee_guaranteed(%lifetime_1) : $@convention(thin) (@guaranteed C) -> ()
   %synchronized = apply %synchronization_point() : $@convention(thin) () -> ()
   %copy_2 = copy_value %copy_1 : $C
+  %result = tuple ()
   end_borrow %lifetime_1 : $C
   %0 = apply %callee_owned(%instance) : $@convention(thin) (@owned C) -> ()
   %1 = apply %callee_owned(%copy_1) : $@convention(thin) (@owned C) -> ()
   %2 = apply %callee_owned(%copy_2) : $@convention(thin) (@owned C) -> ()
 
-  %result = tuple ()
   return %result : $()
 }
 
@@ -1122,6 +1122,54 @@ entry:
   return %retval : $()
 }
 
+// Don't hoist over a copy_value that is rewritten.
+//
+// The borrowee's lifetime will be canonicalized again regardless.
+//
+// CHECK-LABEL: sil [ossa] @nohoist_over_rewritten_copy : {{.*}} { 
+// CHECK:       [[INSTANCE:%[^,]+]] = apply
+// CHECK:       [[LIFETIME:%[^,]+]] = begin_borrow [lexical] [[INSTANCE]]
+// CHECK:       [[COPY:%[^,]+]] = copy_value [[INSTANCE]]
+// CHECK:       end_borrow [[LIFETIME]]
+// CHECK:       tuple ([[COPY]] : $C, [[INSTANCE]] : $C)
+// CHECK-LABEL: } // end sil function 'nohoist_over_rewritten_copy'
+sil [ossa] @nohoist_over_rewritten_copy : $@convention(thin) () -> (@owned C, @owned C) {
+  %get_owned_c = function_ref @get_owned_c : $@convention(thin) () -> (@owned C)
+  %instance = apply %get_owned_c() : $@convention(thin) () -> (@owned C)
+  %lifetime = begin_borrow [lexical] %instance : $C
+  %callee_guaranteed = function_ref @callee_guaranteed : $@convention(thin) (@guaranteed C) -> ()
+  apply %callee_guaranteed(%lifetime) : $@convention(thin) (@guaranteed C) -> ()
+  %copy = copy_value %lifetime : $C
+  end_borrow %lifetime : $C
+  %retval = tuple (%copy : $C, %instance : $C)
+  return %retval : $(C, C)
+}
+
+// Don't hoist over multiple copy_values that are rewritten.
+//
+// The borrowee's lifetime will be canonicalized again regardless.
+//
+// CHECK-LABEL: sil [ossa] @nohoist_over_rewritten_copy_multiple : {{.*}} { 
+// CHECK:       [[INSTANCE:%[^,]+]] = apply
+// CHECK:       [[LIFETIME:%[^,]+]] = begin_borrow [lexical] [[INSTANCE]]
+// CHECK:       [[COPY1:%[^,]+]] = copy_value [[INSTANCE]]
+// CHECK:       [[COPY2:%[^,]+]] = copy_value [[INSTANCE]]
+// CHECK:       end_borrow [[LIFETIME]]
+// CHECK:       tuple ([[COPY2]] : $C, [[COPY1]] : $C, [[INSTANCE]] : $C)
+// CHECK-LABEL: } // end sil function 'nohoist_over_rewritten_copy_multiple'
+sil [ossa] @nohoist_over_rewritten_copy_multiple : $@convention(thin) () -> (@owned C, @owned C, @owned C) {
+  %get_owned_c = function_ref @get_owned_c : $@convention(thin) () -> (@owned C)
+  %instance = apply %get_owned_c() : $@convention(thin) () -> (@owned C)
+  %lifetime = begin_borrow [lexical] %instance : $C
+  %callee_guaranteed = function_ref @callee_guaranteed : $@convention(thin) (@guaranteed C) -> ()
+  apply %callee_guaranteed(%lifetime) : $@convention(thin) (@guaranteed C) -> ()
+  %copy1 = copy_value %lifetime : $C
+  %copy2 = copy_value %lifetime : $C
+  end_borrow %lifetime : $C
+  %retval = tuple (%copy2 : $C, %copy1 : $C, %instance : $C)
+  return %retval : $(C, C, C)
+}
+
 // =============================================================================
 // instruction tests                                                          }}
 // =============================================================================

test/SILOptimizer/shrink_borrow_scope.unit.sil

@@ -0,0 +1,30 @@
+// RUN: %target-sil-opt -unit-test-runner %s 2>&1 | %FileCheck %s
+
+import Builtin
+import Swift
+
+class C {}
+
+sil [ossa] @get_owned_c : $@convention(thin) () -> (@owned C)
+sil [ossa] @callee_guaranteed: $@convention(thin) (@guaranteed C) -> ()
+
+// CHECK-LABEL: begin running test 1 of {{[^,]+}} on nohoist_over_rewritten_copy
+// CHECK-NOT: DELETED:
+// CHECK-NOT: end_borrow
+// CHECK-LABEL: end running test 1 of {{[^,]+}} on nohoist_over_rewritten_copy
+sil [ossa] @nohoist_over_rewritten_copy : $@convention(thin) () -> (@owned C, @owned C) {
+entry:
+  test_specification "shrink-borrow-scope @trace true @trace[1]"
+  %get_owned_c = function_ref @get_owned_c : $@convention(thin) () -> (@owned C)
+  %instance = apply %get_owned_c() : $@convention(thin) () -> (@owned C)
+  %lifetime = begin_borrow [lexical] %instance : $C
+  debug_value [trace] %lifetime : $C
+  %callee_guaranteed = function_ref @callee_guaranteed : $@convention(thin) (@guaranteed C) -> ()
+  apply %callee_guaranteed(%lifetime) : $@convention(thin) (@guaranteed C) -> ()
+  %copy = copy_value %lifetime : $C
+  debug_value [trace] %copy : $C
+  end_borrow %lifetime : $C
+  %retval = tuple (%copy : $C, %instance : $C)
+  return %retval : $(C, C)
+}
+