[LifetimeCompletion] Add availability_with_leaks.

The new boundary allows for invalid OSSA where values are not consumed
on paths leading to blocks that exit the function normally.  In such
cases, the value is allowed to continue leaking as before.

Author: nate-chandler

include/swift/SIL/OSSALifetimeCompletion.h

@@ -57,10 +57,19 @@ class OSSALifetimeCompletion {
   // Availability: "As late as possible."  Consume the value in the last blocks
   //               beyond the non-consuming uses in which the value has been
   //               consumed on no incoming paths.
+  // AvailabilityWithLeaks: "As late as possible or later."  Consume the value
+  //                        in the last blocks beyond the non-consuming uses in
+  //                        which the value has been consumed on no incoming
+  //                        paths, unless that block's terminator isn't an
+  //                        unreachable, in which case, don't consume it there.
+  //
+  //                        This boundary works around bugs where SILGen emits
+  //                        illegal OSSA lifetimes.
   struct Boundary {
     enum Value : uint8_t {
       Liveness,
       Availability,
+      AvailabilityWithLeaks,
     };
     Value value;
 
@@ -106,8 +115,14 @@ class OSSALifetimeCompletion {
                : LifetimeCompletion::AlreadyComplete;
   }
 
+  enum AllowLeaks_t : bool {
+    AllowLeaks = true,
+    DoNotAllowLeaks = false,
+  };
+
   static void visitUnreachableLifetimeEnds(
-      SILValue value, const SSAPrunedLiveness &liveness,
+      SILValue value, AllowLeaks_t allowLeaks,
+      const SSAPrunedLiveness &liveness,
       llvm::function_ref<void(SILInstruction *)> visit);
 
 protected:
@@ -165,6 +180,9 @@ operator<<(llvm::raw_ostream &OS, OSSALifetimeCompletion::Boundary boundary) {
   case OSSALifetimeCompletion::Boundary::Availability:
     OS << "availability";
     break;
+  case OSSALifetimeCompletion::Boundary::AvailabilityWithLeaks:
+    OS << "availability_with_leaks";
+    break;
   }
   return OS;
 }

lib/SIL/Utils/OSSALifetimeCompletion.cpp

@@ -122,14 +122,21 @@ class VisitUnreachableLifetimeEnds {
   /// The value whose dead-end block lifetime ends are to be visited.
   SILValue value;
 
+  /// Whether to allow leaks.
+  ///
+  /// Here, that entails allowing walks to reach non-unreachable terminators and
+  /// not creating lifetime ends before them.
+  OSSALifetimeCompletion::AllowLeaks_t allowLeaks;
+
   /// The non-lifetime-ending boundary of `value`.
   BasicBlockSet starts;
   /// The region between (inclusive) the `starts` and the unreachable blocks.
   BasicBlockSetVector region;
 
 public:
-  VisitUnreachableLifetimeEnds(SILValue value)
-      : value(value), starts(value->getFunction()),
+  VisitUnreachableLifetimeEnds(SILValue value,
+                               OSSALifetimeCompletion::AllowLeaks_t allowLeaks)
+      : value(value), allowLeaks(allowLeaks), starts(value->getFunction()),
         region(value->getFunction()) {}
 
   /// Region discovery.
@@ -232,9 +239,17 @@ void VisitUnreachableLifetimeEnds::computeRegion(
   // (3) Forward walk to find the region in which `value` might be available.
   while (auto *block = regionWorklist.pop()) {
     if (block->succ_empty()) {
-      // This assert will fail unless there is already a lifetime-ending
-      // instruction on each path to normal function exits.
-      assert(isa<UnreachableInst>(block->getTerminator()));
+      // This is a function-exiting block.
+      //
+      // In valid-but-lifetime-incomplete OSSA there must be a lifetime-ending
+      // instruction on each path from the def that exits the function normally.
+      // Thus finding a value available at the end of such a block means that
+      // the block does _not_ must not exits the function normally; in other
+      // words its terminator must be an UnreachableInst.
+      //
+      // In invalid OSSA, indicated by the `allowLeaks` flag, no such guarantee
+      // exists.
+      assert(isa<UnreachableInst>(block->getTerminator()) || allowLeaks);
     }
     for (auto *successor : block->getSuccessorBlocks()) {
       regionWorklist.pushIfNotVisited(successor);
@@ -300,6 +315,12 @@ void VisitUnreachableLifetimeEnds::visitAvailabilityBoundary(
     if (!block->succ_empty() && !hasUnavailableSuccessor()) {
       continue;
     }
+    if (allowLeaks && block->succ_empty() &&
+        !isa<UnreachableInst>(block->getTerminator())) {
+      // Availability extends to the end of a function-exiting-normally block.
+      // If leaks are allowed, don't visit.
+      continue;
+    }
     assert(hasUnavailableSuccessor() ||
            isa<UnreachableInst>(block->getTerminator()));
     visit(block->getTerminator());
@@ -308,10 +329,10 @@ void VisitUnreachableLifetimeEnds::visitAvailabilityBoundary(
 } // end anonymous namespace
 
 void OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
-    SILValue value, const SSAPrunedLiveness &liveness,
+    SILValue value, AllowLeaks_t allowLeaks, const SSAPrunedLiveness &liveness,
     llvm::function_ref<void(SILInstruction *)> visit) {
 
-  VisitUnreachableLifetimeEnds visitor(value);
+  VisitUnreachableLifetimeEnds visitor(value, allowLeaks);
 
   visitor.computeRegion(liveness);
 
@@ -322,12 +343,12 @@ void OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
   visitor.visitAvailabilityBoundary(result, visit);
 }
 
-static bool
-endLifetimeAtAvailabilityBoundary(SILValue value,
-                                  const SSAPrunedLiveness &liveness) {
+static bool endLifetimeAtAvailabilityBoundary(
+    SILValue value, OSSALifetimeCompletion::AllowLeaks_t allowLeaks,
+    const SSAPrunedLiveness &liveness) {
   bool changed = false;
   OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
-      value, liveness, [&](auto *unreachable) {
+      value, allowLeaks, liveness, [&](auto *unreachable) {
         SILBuilderWithScope builder(unreachable);
         endOSSALifetime(value, builder);
         changed = true;
@@ -354,7 +375,12 @@ bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(SILValue value,
     changed |= endLifetimeAtLivenessBoundary(value, liveness.getLiveness());
     break;
   case Boundary::Availability:
-    changed |= endLifetimeAtAvailabilityBoundary(value, liveness.getLiveness());
+    changed |= endLifetimeAtAvailabilityBoundary(value, DoNotAllowLeaks,
+                                                 liveness.getLiveness());
+    break;
+  case Boundary::AvailabilityWithLeaks:
+    changed |= endLifetimeAtAvailabilityBoundary(value, AllowLeaks,
+                                                 liveness.getLiveness());
     break;
   }
   // TODO: Rebuild outer adjacent phis on demand (SILGen does not currently
@@ -379,7 +405,9 @@ static FunctionTest OSSALifetimeCompletionTest(
               arguments.takeString())
               .Case("liveness", OSSALifetimeCompletion::Boundary::Liveness)
               .Case("availability",
-                    OSSALifetimeCompletion::Boundary::Availability);
+                    OSSALifetimeCompletion::Boundary::Availability)
+              .Case("availability_with_leaks",
+                    OSSALifetimeCompletion::Boundary::AvailabilityWithLeaks);
       llvm::outs() << "OSSA lifetime completion on " << kind
                    << " boundary: " << value;
       OSSALifetimeCompletion completion(&function, /*domInfo*/ nullptr);

lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp

@@ -274,7 +274,8 @@ void CanonicalizeOSSALifetime::extendLivenessToDeinitBarriers() {
   }
 
   OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
-      getCurrentDef(), completeLiveness, [&](auto *unreachable) {
+      getCurrentDef(), OSSALifetimeCompletion::DoNotAllowLeaks,
+      completeLiveness, [&](auto *unreachable) {
         recordUnreachableLifetimeEnd(unreachable);
         unreachable->visitPriorInstructions([&](auto *inst) {
           liveness->extendToNonUse(inst);