[LifetimeCompletion] Add extend_lifetimes.

The new instructions are inserted after every "user" (according to
InteriorLiveness' SSAPrunedLiveness instance) outside the linear
liveness boundary.

Author: nate-chandler

include/swift/SIL/OSSALifetimeCompletion.h

@@ -123,6 +123,8 @@ class OSSALifetimeCompletion {
   enum class LifetimeEnd : uint8_t {
     /// The lifetime ends at the boundary.
     Boundary,
+    /// The lifetime "ends" in a loop.
+    Loop,
   };
 
   static void visitAvailabilityBoundary(

lib/SIL/Utils/OSSALifetimeCompletion.cpp

@@ -63,6 +63,9 @@ static SILInstruction *endOSSALifetime(SILValue value,
                                        SILBuilder &builder) {
   auto loc =
     RegularLocation::getAutoGeneratedLocation(builder.getInsertionPointLoc());
+  if (end == OSSALifetimeCompletion::LifetimeEnd::Loop) {
+    return builder.createExtendLifetime(loc, value);
+  }
   if (value->getOwnershipKind() == OwnershipKind::Owned) {
     if (value->getType().is<SILBoxType>()) {
       return builder.createDeallocBox(loc, value);
@@ -389,6 +392,14 @@ void OSSALifetimeCompletion::visitAvailabilityBoundary(
   visitor.propagateAvailablity(result);
 
   visitor.visitAvailabilityBoundary(result, visit);
+
+  visitUsersOutsideLinearLivenessBoundary(
+      value, liveness, [&](auto *instruction) {
+        instruction->visitSubsequentInstructions([&](auto *next) {
+          visit(next, LifetimeEnd::Loop);
+          return true;
+        });
+      });
 }
 
 static bool endLifetimeAtAvailabilityBoundary(

lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp

@@ -276,6 +276,9 @@ void CanonicalizeOSSALifetime::extendLivenessToDeinitBarriers() {
   OSSALifetimeCompletion::visitAvailabilityBoundary(
       getCurrentDef(), OSSALifetimeCompletion::DoNotAllowLeaks,
       completeLiveness, [&](auto *unreachable, auto end) {
+        if (end == OSSALifetimeCompletion::LifetimeEnd::Boundary) {
+          recordUnreachableLifetimeEnd(unreachable);
+        }
         unreachable->visitPriorInstructions([&](auto *inst) {
           liveness->extendToNonUse(inst);
           return true;

test/SILOptimizer/ossa_lifetime_completion.sil

@@ -545,15 +545,104 @@ right:
 // CHECK-LABEL: begin running test {{.*}} on loopy: liveness_partial_boundary_outside_users with: %o
 // CHECK:       end_borrow
 // CHECK-LABEL: end running test {{.*}} on loopy: liveness_partial_boundary_outside_users with: %o
+// CHECK-LABEL: begin running test {{.*}} on loopy: ossa_lifetime_completion
+// CHECK-LABEL: sil [ossa] @loopy : {{.*}} {
+// CHECK:         [[O:%[^,]+]] = apply undef
+// CHECK:         [[B:%[^,]+]] = begin_borrow [[O]]
+// CHECK:         end_borrow [[B]]
+// CHECK:         extend_lifetime [[O]]
+// CHECK:         br bb1
+// CHECK:       bb1:
+// CHECK:         br bb1
+// CHECK-LABEL: } // end sil function 'loopy'
+// CHECK-LABEL: end running test {{.*}} on loopy: ossa_lifetime_completion
+// CHECK-LABEL: begin running test {{.*}} on loopy: liveness_partial_boundary_outside_users with: %o
+// CHECK-NEXT:  end running test {{.*}} on loopy: liveness_partial_boundary_outside_users with: %o
 sil [ossa] @loopy : $@convention(thin) () -> () {
   %o = apply undef() : $@convention(thin) () -> (@owned C)
   specify_test "liveness_partial_boundary_outside_users %o"
+  specify_test "ossa_lifetime_completion %o availability"
+  specify_test "liveness_partial_boundary_outside_users %o"
+  %b = begin_borrow %o : $C
+  end_borrow %b : $C
+  br loop
+
+loop:
+  br loop
+}
+
+sil [ossa] @loopyComplete : $@convention(thin) () -> () {
+  %o = apply undef() : $@convention(thin) () -> (@owned C)
+  specify_test "ossa_lifetime_completion %o availability"
+  specify_test "liveness_partial_boundary_outside_users %o"
   %b = begin_borrow %o : $C
   end_borrow %b : $C
+  extend_lifetime %o : $C
+  br loop
+
+loop:
+  br loop
+}
+
+// When there are no users in the loop, do not extend the lifetime into it.
+// CHECK-LABEL: begin running test {{.*}} on loopy_sometimes: ossa_lifetime_completion with: %c, availability
+// CHECK-LABEL: sil [ossa] @loopy_sometimes : {{.*}} {
+// CHECK-NOT:     extend_lifetime
+// CHECK-LABEL: } // end sil function 'loopy_sometimes'
+// CHECK-LABEL: end running test {{.*}} on loopy_sometimes: ossa_lifetime_completion with: %c, availability
+// CHECK-LABEL: begin running test {{.*}} on loopy_sometimes: liveness_partial_boundary_outside_users with: %c
+// CHECK-NEXT:  end running test {{.*}} on loopy_sometimes: liveness_partial_boundary_outside_users with: %c
+sil [ossa] @loopy_sometimes : $@convention(thin) (@owned C) -> () {
+entry(%c : @owned $C):
+  specify_test "ossa_lifetime_completion %c availability"
+  specify_test "liveness_partial_boundary_outside_users %c"
+  cond_br undef, header, exit
+
+header:
   br loop
 
 loop:
   br loop
+
+exit:
+  destroy_value %c : $C
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// CHECK-LABEL: begin running test {{.*}} on loopy_sometimes_2: ossa_lifetime_completion with: %c, availability
+// CHECK-LABEL: sil [ossa] @loopy_sometimes_2 : {{.*}} {
+// CHECK:       bb0([[O:%[^,]+]] :
+// CHECK:         cond_br undef, [[HEADER:bb[0-9]+]], [[EXIT:bb[0-9]+]]
+// CHECK:       [[HEADER]]:
+// CHECK:         br [[LOOP:bb[0-9]+]]
+// CHECK:       [[LOOP]]:
+// CHECK:         [[B:%[^,]+]] = begin_borrow [[O]]
+// CHECK:         end_borrow [[B]]
+// CHECK:         extend_lifetime [[O]]
+// CHECK:         br [[LOOP]]
+// CHECK:       [[EXIT]]:
+// CHECK:         destroy_value [[O]]
+// CHECK-LABEL: } // end sil function 'loopy_sometimes_2'
+// CHECK-LABEL: end running test {{.*}} on loopy_sometimes_2: ossa_lifetime_completion with: %c, availability
+sil [ossa] @loopy_sometimes_2 : $@convention(thin) (@owned C) -> () {
+entry(%c : @owned $C):
+  specify_test "ossa_lifetime_completion %c availability"
+  specify_test "liveness_partial_boundary_outside_users %c"
+  cond_br undef, header, exit
+
+header:
+  br loop
+
+loop:
+  %b = begin_borrow %c : $C
+  end_borrow %b : $C
+  br loop
+
+exit:
+  destroy_value %c : $C
+  %retval = tuple ()
+  return %retval : $()
 }
 
 // CHECK-LABEL: begin running test {{.*}} on type_dependent_operand: liveness_partial_boundary_outside_users