[CopyProp] Don't extend lifetime over end_access.

If there is a consuming use of the value inside an access scope, the
lifetime should not be extended over the end_access instruction.

Author: nate-chandler

include/swift/SILOptimizer/Utils/CanonicalOSSALifetime.h

@@ -298,6 +298,11 @@ class CanonicalizeOSSALifetime {
   /// ending". end_borrows do not end a liverange that may include owned copies.
   PrunedLiveness liveness;
 
+  /// The destroys of the value.  These are not uses, but need to be recorded so
+  /// that we know when the last use in a consuming block is (without having to
+  /// repeatedly do use-def walks from destroys).
+  SmallPtrSet<DestroyValueInst *, 8> destroys;
+
   /// remnantLiveOutBlocks are part of the original extended lifetime that are
   /// not in canonical pruned liveness. There is a path from a PrunedLiveness
   /// boundary to an original destroy that passes through a remnant block.

lib/SILOptimizer/Utils/CanonicalOSSALifetime.cpp

@@ -162,8 +162,10 @@ bool CanonicalizeOSSALifetime::computeCanonicalLiveness() {
         liveness.updateForUse(user, /*lifetimeEnding*/ true);
         break;
       case OperandOwnership::DestroyingConsume:
-        // destroy_value does not force pruned liveness (but store etc. does).
-        if (!isa<DestroyValueInst>(user)) {
+        if (auto *dvi = dyn_cast<DestroyValueInst>(user)) {
+          destroys.insert(dvi);
+        } else {
+          // destroy_value does not force pruned liveness (but store etc. does).
           liveness.updateForUse(user, /*lifetimeEnding*/ true);
         }
         recordConsumingUse(use);
@@ -335,18 +337,33 @@ void CanonicalizeOSSALifetime::extendLivenessThroughOverlappingAccess() {
   bool changed = true;
   while (changed) {
     changed = false;
-    blockWorklist.initializeRange(consumingBlocks);
-    while (auto *bb = blockWorklist.pop()) {
+    // The blocks in which we may have to extend liveness over access scopes.
+    //
+    // It must be populated first so that we can test membership during the loop
+    // (see findLastConsume).
+    BasicBlockSetVector blocksToVisit(currentDef->getFunction());
+    for (auto *block : consumingBlocks) {
+      blocksToVisit.insert(block);
+    }
+    for (auto iterator = blocksToVisit.begin(); iterator != blocksToVisit.end();
+         ++iterator) {
+      auto *bb = *iterator;
+      // If the block isn't dead, then we won't need to extend liveness within
+      // any of its predecessors (though we may within it).
+      if (liveness.getBlockLiveness(bb) != PrunedLiveBlocks::Dead)
+        continue;
+      // Continue searching upward to find the pruned liveness boundary.
+      for (auto *predBB : bb->getPredecessorBlocks()) {
+        blocksToVisit.insert(predBB);
+      }
+    }
+    for (auto *bb : blocksToVisit) {
       auto blockLiveness = liveness.getBlockLiveness(bb);
       // Ignore blocks within pruned liveness.
       if (blockLiveness == PrunedLiveBlocks::LiveOut) {
         continue;
       }
       if (blockLiveness == PrunedLiveBlocks::Dead) {
-        // Continue searching upward to find the pruned liveness boundary.
-        for (auto *predBB : bb->getPredecessorBlocks()) {
-          blockWorklist.insert(predBB);
-        }
         // Otherwise, ignore dead blocks with no nonlocal end_access.
         if (!accessBlocks->containsNonLocalEndAccess(bb)) {
           continue;
@@ -356,7 +373,43 @@ void CanonicalizeOSSALifetime::extendLivenessThroughOverlappingAccess() {
       // Find the latest partially overlapping access scope, if one exists:
       //     use %def // pruned liveness ends here
       //     end_access
+
+      // Whether to look for the last consume in the block.
+      //
+      // If the block is consuming, we need to determine whether any subsequent
+      // block contains a destroy_value.  If any does, then we need to count
+      // every end_access in this block as a use.  If none do, then we only
+      // need to count end_accesses which occur before the consuming use in
+      // this block as uses.
+      bool findLastConsume =
+          consumingBlocks.contains(bb) &&
+          llvm::none_of(bb->getSuccessorBlocks(), [&](auto *successor) {
+            // If a block was originally live, then it either contained or was
+            // backwards reachable from a destroy or a use.  If it was dead,
+            // then it was only backward reachable from a destroy.
+            //
+            // If we've seen a destroy, then we've seen the last consuming use.
+            return originalLiveBlocks.contains(successor) &&
+                   liveness.getBlockLiveness(successor) ==
+                       PrunedLiveBlocks::Dead;
+          });
       for (auto &inst : llvm::reverse(*bb)) {
+        if (findLastConsume) {
+          // The value is consumed in the block, and we haven't already seen the
+          // consuming use.
+          //
+          // If this instruction is the consuming use, note that we've seen it.
+          //
+          // We only need to check for destroys because we exit the loop if the
+          // use is interesting.
+          if (auto *dvi = dyn_cast<DestroyValueInst>(&inst))
+            findLastConsume = !destroys.contains(dvi);
+          // Either this instruction is the consuming use or we haven't seen the
+          // consuming use yet.  If it's the consuming use, it's not an
+          // overlapping end_access.  If we haven't seen it yet, if it's an
+          // end_access, it doesn't overlap.
+          continue;
+        }
         // Stop at the latest use. An earlier end_access does not overlap.
         if (blockHasUse && liveness.isInterestingUser(&inst)) {
           break;

test/SILOptimizer/copy_propagation.sil

@@ -616,24 +616,21 @@ bb2:
 
 // Original Overlapping (unnecessarilly extends pruned liveness):
 //
-// TODO: this copy could be avoided but is probably an unusual case,
-// and sinking the destroy outside the access scope might help to
-// optimize the access itself.
-//
 //     %def
 //     begin_access // access scope unrelated to def
 //     use %def     // pruned liveness ends here
 //     destroy %def
 //     end_access
 //
 // CHECK-LABEL: sil [ossa] @testOriginalOverlapInLiveBlock : $@convention(thin) () -> () {
-// CHECK:   [[DEF:%.*]] = apply %{{.*}}() : $@convention(thin) () -> @owned AnyObject
-// CHECK:   begin_access
-// CHECK:   [[COPY:%.*]] = copy_value [[DEF]] : $AnyObject
-// CHECK:   store [[COPY]] to [init] %{{.*}} : $*AnyObject
-// CHECK:   end_access
-// CHECK:   destroy_value [[DEF]] : $AnyObject
-// CHECK: br bb1
+// CHECK:         [[DEF:%.*]] = apply %{{.*}}() : $@convention(thin) () -> @owned AnyObject
+// CHECK:         begin_access
+// CHECK-OPT:     store [[DEF]] to [init] %{{.*}} : $*AnyObject
+// CHECK-ONONE:   [[COPY:%[^,]+]] = copy_value [[DEF]]
+// CHECK-ONONE:   store [[COPY]] to [init] %{{.*}} : $*AnyObject
+// CHECK-ONONE:   destroy_value [[DEF]]
+// CHECK:         end_access
+// CHECK:         br bb1
 // CHECK-LABEL: } // end sil function 'testOriginalOverlapInLiveBlock'
 sil [ossa] @testOriginalOverlapInLiveBlock : $@convention(thin) () -> () {
 bb0:
@@ -901,3 +898,44 @@ exit:
     %retval = tuple ()
     return %retval : $()
 }
+
+// CHECK-LABEL: sil [ossa] @dont_extend_beyond_nonoverlapping_end_access_after_store_in_consuming_block : {{.*}} {
+// CHECK:       {{bb[0-9]+}}([[INSTANCE:%[^,]+]] : @owned $C
+// CHECK:         [[COPY:%[^,]+]] = copy_value [[INSTANCE]]
+// CHECK:         begin_access
+// CHECK-NOT:     copy_value
+// CHECK:         store [[COPY]] to [init] {{%[^,]+}}
+// CHECK-LABEL: } // end sil function 'dont_extend_beyond_nonoverlapping_end_access_after_store_in_consuming_block'
+sil [ossa] @dont_extend_beyond_nonoverlapping_end_access_after_store_in_consuming_block : $@convention(thin) (@owned C) -> () {
+bb0(%instance : @owned $C):
+  %addr = alloc_stack [lexical] $C, var
+  %borrow = begin_borrow %instance : $C
+  %copy = copy_value %borrow : $C
+  end_borrow %borrow : $C
+  destroy_value %instance : $C
+  %access = begin_access [static] [modify] %addr : $*C
+  store %copy to [init] %addr : $*C
+  end_access %access : $*C
+  destroy_addr %addr : $*C
+  dealloc_stack %addr : $*C
+  %retval = tuple()
+  return %retval : $()
+}
+
+// CHECK-LABEL: sil [ossa] @dont_extend_beyond_nonoverlapping_endaccess_after_destroyvalue_in_consuming_block : $@convention(thin) () -> () {
+// CHECK:         begin_access
+// CHECK:         destroy_value
+// CHECK:         end_access
+// CHECK-NOT:     destroy_value
+// CHECK-LABEL: } // end sil function 'dont_extend_beyond_nonoverlapping_endaccess_after_destroyvalue_in_consuming_block'
+sil [ossa] @dont_extend_beyond_nonoverlapping_endaccess_after_destroyvalue_in_consuming_block : $@convention(thin) () -> () {
+  %instance = apply undef() : $@convention(thin) () -> @owned C
+  %addr = alloc_stack [lexical] $C, var
+  %access = begin_access [static] [modify] %addr : $*C
+  apply undef(%instance) : $@convention(thin) (@guaranteed C) -> ()
+  destroy_value %instance : $C
+  end_access %access : $*C
+  dealloc_stack %addr : $*C
+  %retval = tuple()
+  return %retval : $()
+}