Merge pull request #73386 from nate-chandler/lifetime-completion/20240501/2

[LifetimeCompletion] Require boundary to be specified.

Author: nate-chandler

include/swift/SIL/OSSALifetimeCompletion.h

@@ -57,24 +57,24 @@ 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;
 
     Boundary(Value value) : value(value){};
     operator Value() const { return value; }
-
-    static std::optional<Boundary> getForcingLiveness(bool force) {
-      if (!force)
-        return {};
-      return {Liveness};
-    }
-
-    bool isLiveness() { return value == Liveness; }
-    bool isAvailable() { return !isLiveness(); }
   };
 
   /// Insert a lifetime-ending instruction on every path to complete the OSSA
@@ -95,9 +95,7 @@ class OSSALifetimeCompletion {
   /// lifetime.
   ///
   /// TODO: We also need to complete scoped addresses (e.g. store_borrow)!
-  LifetimeCompletion
-  completeOSSALifetime(SILValue value,
-                       std::optional<Boundary> maybeBoundary = std::nullopt) {
+  LifetimeCompletion completeOSSALifetime(SILValue value, Boundary boundary) {
     if (value->getOwnershipKind() == OwnershipKind::None)
       return LifetimeCompletion::NoLifetime;
 
@@ -112,16 +110,19 @@ class OSSALifetimeCompletion {
     if (!completedValues.insert(value))
       return LifetimeCompletion::AlreadyComplete;
 
-    Boundary boundary = maybeBoundary.value_or(
-        value->isLexical() ? Boundary::Availability : Boundary::Liveness);
-
     return analyzeAndUpdateLifetime(value, boundary)
                ? LifetimeCompletion::WasCompleted
                : 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:
@@ -170,6 +171,22 @@ class UnreachableLifetimeCompletion {
   bool completeLifetimes();
 };
 
+inline llvm::raw_ostream &
+operator<<(llvm::raw_ostream &OS, OSSALifetimeCompletion::Boundary boundary) {
+  switch (boundary) {
+  case OSSALifetimeCompletion::Boundary::Liveness:
+    OS << "liveness";
+    break;
+  case OSSALifetimeCompletion::Boundary::Availability:
+    OS << "availability";
+    break;
+  case OSSALifetimeCompletion::Boundary::AvailabilityWithLeaks:
+    OS << "availability_with_leaks";
+    break;
+  }
+  return OS;
+}
+
 } // namespace swift
 
 #endif

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);
@@ -291,27 +306,33 @@ void VisitUnreachableLifetimeEnds::visitAvailabilityBoundary(
     if (!available) {
       continue;
     }
-    auto hasUnreachableSuccessor = [&]() {
+    auto hasUnavailableSuccessor = [&]() {
       // Use a lambda to avoid checking if possible.
       return llvm::any_of(block->getSuccessorBlocks(), [&result](auto *block) {
         return result.getState(block) == State::Unavailable;
       });
     };
-    if (!block->succ_empty() && !hasUnreachableSuccessor()) {
+    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(hasUnreachableSuccessor() ||
+    assert(hasUnavailableSuccessor() ||
            isa<UnreachableInst>(block->getTerminator()));
     visit(block->getTerminator());
   }
 }
 } // 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;
@@ -342,20 +363,25 @@ bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(SILValue value,
                                                       Boundary boundary) {
   // Called for inner borrows, inner adjacent reborrows, inner reborrows, and
   // scoped addresses.
-  auto handleInnerScope = [this](SILValue innerBorrowedValue) {
-    completeOSSALifetime(innerBorrowedValue);
+  auto handleInnerScope = [this, boundary](SILValue innerBorrowedValue) {
+    completeOSSALifetime(innerBorrowedValue, boundary);
   };
   InteriorLiveness liveness(value);
   liveness.compute(domInfo, handleInnerScope);
 
   bool changed = false;
   switch (boundary) {
-  case Boundary::Availability:
-    changed |= endLifetimeAtAvailabilityBoundary(value, liveness.getLiveness());
-    break;
   case Boundary::Liveness:
     changed |= endLifetimeAtLivenessBoundary(value, liveness.getLiveness());
     break;
+  case Boundary::Availability:
+    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
   // produce guaranteed phis). See FindEnclosingDefs &
@@ -371,16 +397,19 @@ namespace swift::test {
 // Dumps:
 // - function
 static FunctionTest OSSALifetimeCompletionTest(
-    "ossa-lifetime-completion",
+    "ossa_lifetime_completion",
     [](auto &function, auto &arguments, auto &test) {
       SILValue value = arguments.takeValue();
-      std::optional<OSSALifetimeCompletion::Boundary> kind = std::nullopt;
-      if (arguments.hasUntaken()) {
-        kind = arguments.takeBool()
-                   ? OSSALifetimeCompletion::Boundary::Liveness
-                   : OSSALifetimeCompletion::Boundary::Availability;
-      }
-      llvm::outs() << "OSSA lifetime completion: " << value;
+      OSSALifetimeCompletion::Boundary kind =
+          llvm::StringSwitch<OSSALifetimeCompletion::Boundary>(
+              arguments.takeString())
+              .Case("liveness", OSSALifetimeCompletion::Boundary::Liveness)
+              .Case("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);
       completion.completeOSSALifetime(value, kind);
       function.print(llvm::outs());
@@ -462,8 +491,9 @@ bool UnreachableLifetimeCompletion::completeLifetimes() {
 
   bool changed = false;
   for (auto value : incompleteValues) {
-    if (completion.completeOSSALifetime(value)
-        == LifetimeCompletion::WasCompleted) {
+    if (completion.completeOSSALifetime(
+            value, OSSALifetimeCompletion::Boundary::Availability) ==
+        LifetimeCompletion::WasCompleted) {
       changed = true;
     }
   }

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -4014,7 +4014,8 @@ bool MoveOnlyAddressChecker::completeLifetimes() {
                          [](auto *user) { return isa<BranchInst>(user); })) {
           continue;
         }
-        if (completion.completeOSSALifetime(result) ==
+        if (completion.completeOSSALifetime(
+                result, OSSALifetimeCompletion::Boundary::Availability) ==
             LifetimeCompletion::WasCompleted) {
           changed = true;
         }
@@ -4024,7 +4025,8 @@ bool MoveOnlyAddressChecker::completeLifetimes() {
       if (arg->isReborrow()) {
         continue;
       }
-      if (completion.completeOSSALifetime(arg) ==
+      if (completion.completeOSSALifetime(
+              arg, OSSALifetimeCompletion::Boundary::Availability) ==
           LifetimeCompletion::WasCompleted) {
         changed = true;
       }

lib/SILOptimizer/Mandatory/MoveOnlyChecker.cpp

@@ -163,7 +163,8 @@ void MoveOnlyChecker::completeObjectLifetimes(
       for (auto result : inst.getResults()) {
         if (!transitiveValues.isVisited(result))
           continue;
-        if (completion.completeOSSALifetime(result) ==
+        if (completion.completeOSSALifetime(
+                result, OSSALifetimeCompletion::Boundary::Availability) ==
             LifetimeCompletion::WasCompleted) {
           madeChange = true;
         }
@@ -173,7 +174,8 @@ void MoveOnlyChecker::completeObjectLifetimes(
       assert(!arg->isReborrow() && "reborrows not legal at this SIL stage");
       if (!transitiveValues.isVisited(arg))
         continue;
-      if (completion.completeOSSALifetime(arg) ==
+      if (completion.completeOSSALifetime(
+              arg, OSSALifetimeCompletion::Boundary::Availability) ==
           LifetimeCompletion::WasCompleted) {
         madeChange = true;
       }

lib/SILOptimizer/Mandatory/PredictableMemOpt.cpp

@@ -2673,8 +2673,9 @@ bool AllocOptimize::tryToRemoveDeadAllocation() {
     // Lexical enums can have incomplete lifetimes in non payload paths that
     // don't end in unreachable. Force their lifetime to end immediately after
     // the last use instead.
-    auto boundary = OSSALifetimeCompletion::Boundary::getForcingLiveness(
-        v->getType().isOrHasEnum());
+    auto boundary = v->getType().isOrHasEnum()
+                        ? OSSALifetimeCompletion::Boundary::Liveness
+                        : OSSALifetimeCompletion::Boundary::Availability;
     LLVM_DEBUG(llvm::dbgs() << "Completing lifetime of: ");
     LLVM_DEBUG(v->dump());
     completion.completeOSSALifetime(v, boundary);

lib/SILOptimizer/Mandatory/SILGenCleanup.cpp

@@ -117,15 +117,18 @@ bool SILGenCleanup::completeOSSALifetimes(SILFunction *function) {
   for (auto *block : postOrder->getPostOrder()) {
     for (SILInstruction &inst : reverse(*block)) {
       for (auto result : inst.getResults()) {
-        if (completion.completeOSSALifetime(result) ==
+        if (completion.completeOSSALifetime(
+                result,
+                OSSALifetimeCompletion::Boundary::AvailabilityWithLeaks) ==
             LifetimeCompletion::WasCompleted) {
           changed = true;
         }
       }
     }
     for (SILArgument *arg : block->getArguments()) {
       assert(!arg->isReborrow() && "reborrows not legal at this SIL stage");
-      if (completion.completeOSSALifetime(arg) ==
+      if (completion.completeOSSALifetime(
+              arg, OSSALifetimeCompletion::Boundary::AvailabilityWithLeaks) ==
           LifetimeCompletion::WasCompleted) {
         changed = true;
       }

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);