[OSSACanonicalizeOwned] Dead-end extend base lifetime.

include/swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h

@@ -237,6 +237,8 @@ class CanonicalizeOSSALifetime final {
   /// copies.
   const MaximizeLifetime_t maximizeLifetime;
 
+  SILFunction *function;
+
   // If present, will be used to ensure that the lifetime is not shortened to
   // end inside an access scope which it previously enclosed.  (Note that ending
   // before such an access scope is fine regardless.)
@@ -354,7 +356,7 @@ class CanonicalizeOSSALifetime final {
       DeadEndBlocksAnalysis *deadEndBlocksAnalysis, DominanceInfo *domTree,
       BasicCalleeAnalysis *calleeAnalysis, InstructionDeleter &deleter)
       : pruneDebugMode(pruneDebugMode), maximizeLifetime(maximizeLifetime),
-        accessBlockAnalysis(accessBlockAnalysis),
+        function(function), accessBlockAnalysis(accessBlockAnalysis),
         deadEndBlocksAnalysis(deadEndBlocksAnalysis), domTree(domTree),
         calleeAnalysis(calleeAnalysis), deleter(deleter) {}
 
@@ -368,9 +370,7 @@ class CanonicalizeOSSALifetime final {
     currentDef = def;
     currentLexicalLifetimeEnds = lexicalLifetimeEnds;
 
-    if (maximizeLifetime || respectsDeinitBarriers()) {
-      liveness->initializeDiscoveredBlocks(&discoveredBlocks);
-    }
+    liveness->initializeDiscoveredBlocks(&discoveredBlocks);
     liveness->initializeDef(getCurrentDef());
   }
 
@@ -499,6 +499,7 @@ class CanonicalizeOSSALifetime final {
                             PrunedLivenessBoundary &boundary);
 
   void extendLivenessToDeadEnds();
+  void extendLexicalLivenessToDeadEnds();
   void extendLivenessToDeinitBarriers();
 
   void extendUnconsumedLiveness(PrunedLivenessBoundary const &boundary);

lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp

@@ -239,6 +239,13 @@ bool CanonicalizeOSSALifetime::computeCanonicalLiveness() {
         break;
       case OperandOwnership::InteriorPointer:
       case OperandOwnership::AnyInteriorPointer:
+        if (liveness->checkAndUpdateInteriorPointer(use) !=
+            AddressUseKind::NonEscaping) {
+          LLVM_DEBUG(llvm::dbgs()
+                     << "      Inner address use is escaping! Giving up\n");
+          return false;
+        }
+        break;
       case OperandOwnership::GuaranteedForwarding:
       case OperandOwnership::EndBorrow:
         // Guaranteed values are exposed by inner adjacent reborrows. If user is
@@ -274,6 +281,49 @@ bool CanonicalizeOSSALifetime::computeCanonicalLiveness() {
   return true;
 }
 
+/// Extend liveness to the availability boundary of currentDef.  Even if a copy
+/// is consumed on a path to the dead-end, if the def stays live through to the
+/// dead-end, its lifetime must not be shrunk back from it (eventually we'll
+/// support shrinking it back to deinit barriers).
+///
+/// Example:
+///     %def is lexical
+///     %copy = copy_value %def
+///     consume %copy
+///     apply %foo() // deinit barrier
+///     // Must extend lifetime of %def up to this point per language rules.
+///     unreachable
+void CanonicalizeOSSALifetime::extendLexicalLivenessToDeadEnds() {
+  // TODO: OSSALifetimeCompletion: Once lifetimes are always complete, delete
+  //                               this method.
+  SmallVector<SILBasicBlock *, 32> directDiscoverdBlocks;
+  SSAPrunedLiveness directLiveness(function, &directDiscoverdBlocks);
+  directLiveness.initializeDef(getCurrentDef());
+  directLiveness.computeSimple();
+  OSSALifetimeCompletion::visitAvailabilityBoundary(
+      getCurrentDef(), directLiveness, [&](auto *unreachable, auto end) {
+        if (end == OSSALifetimeCompletion::LifetimeEnd::Boundary) {
+          recordUnreachableLifetimeEnd(unreachable);
+        }
+        unreachable->visitPriorInstructions([&](auto *inst) {
+          liveness->extendToNonUse(inst);
+          return true;
+        });
+      });
+}
+
+/// Extend liveness to the copy-extended availability boundary of currentDef.
+/// Prevents destroys from being inserted between borrows of (copies of) the
+/// def and dead-ends.
+///
+/// Example:
+///     %def need not be lexical
+///     %c = copy_value %def
+///     %sb = store_borrow %c to %addr
+///     // Must extend lifetime of %def up to this point.  Otherwise, a
+///     // destroy_value could be inserted within a borrow scope or interior
+///     // pointer use.
+///     unreachable
 void CanonicalizeOSSALifetime::extendLivenessToDeadEnds() {
   // TODO: OSSALifetimeCompletion: Once lifetimes are always complete, delete
   //                               this method.
@@ -1357,7 +1407,10 @@ bool CanonicalizeOSSALifetime::computeLiveness() {
     return false;
   }
   if (respectsDeinitBarriers()) {
-    extendLivenessToDeadEnds();
+    extendLexicalLivenessToDeadEnds();
+  }
+  extendLivenessToDeadEnds();
+  if (respectsDeinitBarriers()) {
     extendLivenessToDeinitBarriers();
   }
   if (accessBlockAnalysis) {

test/SILOptimizer/assemblyvision_remark/chacha.swift

@@ -28,8 +28,8 @@ func checkResult(_ plaintext: [UInt8]) {
 
 @_semantics("optremark.sil-assembly-vision-remark-gen")
 public func run_ChaCha(_ N: Int) {
-  let key = Array(repeating: UInt8(1), count: 32) // expected-remark {{release of type '}}
-  let nonce = Array(repeating: UInt8(2), count: 12) // expected-remark {{release of type '}}
+  let key = Array(repeating: UInt8(1), count: 32) // expected-note {{of 'key}}
+  let nonce = Array(repeating: UInt8(2), count: 12) // expected-note {{of 'nonce}}
 
   var checkedtext = Array(repeating: UInt8(0), count: 1024) // expected-note {{of 'checkedtext}}
   ChaCha20.encrypt(bytes: &checkedtext, key: key, nonce: nonce)
@@ -44,3 +44,5 @@ public func run_ChaCha(_ N: Int) {
   }
 } // expected-remark {{release of type '}}
   // expected-remark @-1 {{release of type '}}
+  // expected-remark @-2 {{release of type '}}
+  // expected-remark @-3 {{release of type '}}

test/SILOptimizer/canonicalize_ossa_lifetime_unit.sil

@@ -1,13 +1,27 @@
 // RUN: %target-sil-opt -sil-print-types -test-runner %s -o /dev/null 2>&1 | %FileCheck %s
 
+import Builtin
 import Swift
 
 class C {}
+
+enum Never {}
+
+struct Pointer {
+  var rawValue: Builtin.RawPointer
+}
+
+struct Int32 {
+  var int: Builtin.Int32
+}
+sil @run : $@convention(thin) () -> Never
 sil @getOwned : $@convention(thin) () -> @owned C
 sil @barrier : $@convention(thin) () -> ()
 sil [ossa] @getC : $@convention(thin) () -> @owned C
 sil [ossa] @borrowC : $@convention(thin) (@guaranteed C) -> ()
 sil [ossa] @takeC : $@convention(thin) (@owned C) -> ()
+sil @getPointer : $@convention(thin) () -> (@owned Pointer)
+
 struct S {}
 
 struct MoS: ~Copyable {}
@@ -885,6 +899,33 @@ die:
   unreachable
 }
 
+// CHECK-LABEL: begin running test {{.*}} on consume_copy_before_use_in_dead_end_2
+// CHECK-LABEL: sil [ossa] @consume_copy_before_use_in_dead_end_2 : {{.*}} {
+// CHECK-NOT:     destroy_value [dead_end]
+// CHECK-LABEL: } // end sil function 'consume_copy_before_use_in_dead_end_2'
+// CHECK-LABEL: end running test {{.*}} on consume_copy_before_use_in_dead_end_2
+sil [ossa] @consume_copy_before_use_in_dead_end_2 : $@convention(thin) (@owned C) -> () {
+entry(%c : @owned $C):
+  cond_br undef, exit, die
+
+exit:
+  destroy_value %c
+  %retval = tuple ()
+  return %retval
+
+die:
+  %move = move_value %c
+  %copy = copy_value %move
+  specify_test "canonicalize_ossa_lifetime true false true %move"
+  apply undef(%move) : $@convention(thin) (@owned C) -> ()
+  %addr = alloc_stack $C
+  %token = store_borrow %copy to %addr
+  apply undef() : $@convention(thin) () -> ()
+  %reload = load_borrow %token
+  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
@@ -918,3 +959,43 @@ entry(%a: @owned $C):
   %retval = tuple ()
   return %retval
 }
+
+// CHECK-LABEL: begin running test {{.*}} on extend_lifetime_to_deadend_despite_copy_consume
+// CHECK-LABEL: sil [ossa] @extend_lifetime_to_deadend_despite_copy_consume : {{.*}} {
+// CHECK:         copy_value
+// CHECK-LABEL: } // end sil function 'extend_lifetime_to_deadend_despite_copy_consume'
+// CHECK-LABEL: end running test {{.*}} on extend_lifetime_to_deadend_despite_copy_consume
+sil [ossa] @extend_lifetime_to_deadend_despite_copy_consume : $@convention(thin) (@owned C) -> () {
+entry(%c : @owned $C):
+  specify_test "canonicalize_ossa_lifetime true false true @argument"
+  %cc = copy_value %c
+  %takeC = function_ref @takeC : $@convention(thin) (@owned C) -> ()
+  apply %takeC(%cc) : $@convention(thin) (@owned C) -> ()
+  %run = function_ref @run : $@convention(thin) () -> Never
+  unreachable
+}
+
+/// If there's a pointer escape (here, mark_dependence) of the value, we can't
+/// canonicalize its lifetime.
+// CHECK-LABEL: begin running test {{.*}} on dont_canonicalize_on_pointer_escape
+// CHECK-LABEL: sil [ossa] @dont_canonicalize_on_pointer_escape : {{.*}} {
+// CHECK:         load
+// CHECK-NEXT:    destroy_value
+// CHECK-LABEL: } // end sil function 'dont_canonicalize_on_pointer_escape'
+// CHECK-LABEL: end running test {{.*}} on dont_canonicalize_on_pointer_escape
+sil [ossa] @dont_canonicalize_on_pointer_escape : $@convention(thin) () -> () {
+entry:
+  %getC = function_ref @getC : $@convention(thin) () -> (@owned C)
+  %getPointer = function_ref @getPointer : $@convention(thin) () -> (@owned Pointer)
+  %c = apply %getC() : $@convention(thin) () -> (@owned C)
+  %pointer = apply %getPointer() : $@convention(thin) () -> (@owned Pointer)
+  specify_test "canonicalize_ossa_lifetime true false true %c"
+  %rawPointer = struct_extract %pointer, #Pointer.rawValue
+  %o = pointer_to_address %rawPointer to [strict] $*Int32
+  %dependent = mark_dependence [nonescaping] %o on %c
+  %int_addr = struct_element_addr %dependent, #Int32.int
+  %int = load [trivial] %int_addr
+  destroy_value %c
+  %retval = tuple ()
+  return %retval : $()
+}