Merge pull request #68975 from nate-chandler/rdar116255254

[OSSALifetimeCompletion] Handle unavailable blocks in dead-end regions.

Author: nate-chandler

include/swift/SIL/OSSALifetimeCompletion.h

@@ -87,7 +87,7 @@ class OSSALifetimeCompletion {
 
   static void visitUnreachableLifetimeEnds(
       SILValue value, const SSAPrunedLiveness &liveness,
-      llvm::function_ref<void(UnreachableInst *)> visit);
+      llvm::function_ref<void(SILInstruction *)> visit);
 
 protected:
   bool analyzeAndUpdateLifetime(SILValue value, bool forceBoundaryCompletion);

include/swift/SIL/SILBasicBlock.h

@@ -518,13 +518,13 @@ public SwiftObjectHeader {
 
 #ifndef NDEBUG
   /// Print the ID of the block, bbN.
-  void dumpID() const;
+  void dumpID(bool newline = true) const;
 
   /// Print the ID of the block with \p OS, bbN.
-  void printID(llvm::raw_ostream &OS) const;
+  void printID(llvm::raw_ostream &OS, bool newline = true) const;
 
   /// Print the ID of the block with \p Ctx, bbN.
-  void printID(SILPrintContext &Ctx) const;
+  void printID(SILPrintContext &Ctx, bool newline = true) const;
 #endif
 
   /// getSublistAccess() - returns pointer to member of instruction list

include/swift/SILOptimizer/Utils/CanonicalizeOSSALifetime.h

@@ -136,6 +136,10 @@ class CanonicalOSSAConsumeInfo final {
   /// Map blocks on the lifetime boundary to the last consuming instruction.
   llvm::SmallDenseMap<SILBasicBlock *, SILInstruction *, 4> finalBlockConsumes;
 
+  /// The instructions on the availability boundary of the dead-end region where
+  /// this value is not consumed.
+  SmallPtrSet<SILInstruction *, 4> unreachableLifetimeEnds;
+
 public:
   void clear() { finalBlockConsumes.clear(); }
 
@@ -161,6 +165,14 @@ class CanonicalOSSAConsumeInfo final {
     return false;
   }
 
+  void recordUnreachableLifetimeEnd(SILInstruction *inst) {
+    unreachableLifetimeEnds.insert(inst);
+  }
+
+  bool isUnreachableLifetimeEnd(SILInstruction *inst) {
+    return unreachableLifetimeEnds.contains(inst);
+  }
+
   CanonicalOSSAConsumeInfo() {}
   CanonicalOSSAConsumeInfo(CanonicalOSSAConsumeInfo const &) = delete;
   CanonicalOSSAConsumeInfo &
@@ -235,9 +247,12 @@ class CanonicalizeOSSALifetime final {
   SILValue currentDef;
 
   /// Original points in the CFG where the current value's lifetime is consumed
-  /// or destroyed. For guaranteed values it remains empty. A backward walk from
-  /// these blocks must discover all uses on paths that lead to a return or
-  /// throw.
+  /// or destroyed.  Each block either contains a consuming instruction (e.g.
+  /// `destroy_value`) or is on the availability boundary of the value in a
+  /// dead-end region (e.g. `unreachable`).
+  ///
+  /// For guaranteed values it remains empty. A backward walk from these blocks
+  /// must discover all uses on paths that lead to a return or throw.
   ///
   /// These blocks are not necessarily in the pruned live blocks since
   /// pruned liveness does not consider destroy_values.
@@ -406,9 +421,18 @@ class CanonicalizeOSSALifetime final {
   void recordDebugValue(DebugValueInst *dvi) { debugValues.insert(dvi); }
 
   void recordConsumingUse(Operand *use) { recordConsumingUser(use->getUser()); }
+  /// Record that the value is consumed at `user`.
+  ///
+  /// Either `user` is a consuming use (e.g. `destroy_value`) or it is the
+  /// terminator of a block on the availability boundary of the value in a
+  /// dead-end region (e.g. `unreachable`).
   void recordConsumingUser(SILInstruction *user) {
     consumingBlocks.insert(user->getParent());
   }
+  void recordUnreachableLifetimeEnd(SILInstruction *user) {
+    recordConsumingUser(user);
+    consumes.recordUnreachableLifetimeEnd(user);
+  }
   bool computeCanonicalLiveness();
 
   bool endsAccessOverlappingPrunedBoundary(SILInstruction *inst);

lib/SIL/IR/SILPrinter.cpp

@@ -853,8 +853,11 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
   }
 
 #ifndef NDEBUG
-  void printID(const SILBasicBlock *BB) {
-    *this << Ctx.getID(BB) << "\n";
+  void printID(const SILBasicBlock *BB, bool newline) {
+    *this << Ctx.getID(BB);
+    if (newline) {
+      *this << "\n";
+    }
   }
 #endif
 
@@ -3115,17 +3118,17 @@ void SILBasicBlock::print(SILPrintContext &Ctx) const {
 }
 
 #ifndef NDEBUG
-void SILBasicBlock::dumpID() const {
-  printID(llvm::errs());
+void SILBasicBlock::dumpID(bool newline) const {
+  printID(llvm::errs(), newline);
 }
 
-void SILBasicBlock::printID(llvm::raw_ostream &OS) const {
+void SILBasicBlock::printID(llvm::raw_ostream &OS, bool newline) const {
   SILPrintContext Ctx(OS);
-  printID(Ctx);
+  printID(Ctx, newline);
 }
 
-void SILBasicBlock::printID(SILPrintContext &Ctx) const {
-  SILPrinter(Ctx).printID(this);
+void SILBasicBlock::printID(SILPrintContext &Ctx, bool newline) const {
+  SILPrinter(Ctx).printID(this, newline);
 }
 #endif
 

lib/SIL/Utils/OSSALifetimeCompletion.cpp

@@ -51,8 +51,10 @@
 
 #include "swift/SIL/OSSALifetimeCompletion.h"
 #include "swift/SIL/SILBuilder.h"
+#include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/Test.h"
+#include "llvm/ADT/STLExtras.h"
 
 using namespace swift;
 
@@ -99,40 +101,217 @@ static bool endLifetimeAtBoundary(SILValue value,
   return changed;
 }
 
-void OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
-    SILValue value, const SSAPrunedLiveness &liveness,
-    llvm::function_ref<void(UnreachableInst *)> visit) {
+namespace {
+/// Implements OSSALifetimeCompletion::visitUnreachableLifetimeEnds.  Finds
+/// positions as near as possible to unreachables at which `value`'s lifetime
+/// is available.
+///
+/// Finding these positions is a three step process:
+/// 1) computeRegion: Forward CFG walk from non-lifetime-ending boundary to find
+///                   the dead-end region in which the value might be available.
+/// 2) propagateAvailability: Forward iterative dataflow within the region to
+///                           determine which blocks the value is available in.
+/// 3) visitAvailabilityBoundary: Visits the final blocks in the region where
+///                               the value is available--these are the blocks
+///                               without successors or with at least one
+///                               unavailable successor.
+class VisitUnreachableLifetimeEnds {
+  /// The value whose dead-end block lifetime ends are to be visited.
+  SILValue value;
+
+  /// 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()),
+        region(value->getFunction()) {}
+
+  /// Region discovery.
+  ///
+  /// Forward CFG walk from non-lifetime-ending boundary to unreachable
+  /// instructions.
+  void computeRegion(const SSAPrunedLiveness &liveness);
+
+  struct Result;
+
+  /// Iterative dataflow to determine availability for each block in `region`.
+  void propagateAvailablity(Result &result);
+
+  /// Visit the terminators of blocks on the boundary of availability.
+  void
+  visitAvailabilityBoundary(Result const &result,
+                            llvm::function_ref<void(SILInstruction *)> visit);
+
+  struct State {
+    enum class Value : uint8_t {
+      Unavailable = 0,
+      Available,
+      Unknown,
+    };
+    Value value;
+
+    State(Value value) : value(value){};
+    operator Value() const { return value; }
+    State meet(State const other) const {
+      return *this < other ? *this : other;
+    }
+
+    static State Unavailable() { return {Value::Unavailable}; }
+    static State Available() { return {Value::Available}; }
+    static State Unknown() { return {Value::Unknown}; }
+  };
+
+  struct Result {
+    BasicBlockBitfield states;
+
+    Result(SILFunction *function) : states(function, 2) {}
+
+    State getState(SILBasicBlock *block) const {
+      return {(State::Value)states.get(block)};
+    }
+
+    void setState(SILBasicBlock *block, State newState) {
+      states.set(block, (unsigned)newState.value);
+    }
+
+    /// Propagate predecessors' state into `block`.
+    ///
+    /// states[block] ∧= state[predecessor_1] ∧ ... ∧ state[predecessor_n]
+    bool updateState(SILBasicBlock *block) {
+      auto oldState = getState(block);
+      auto state = oldState;
+      for (auto *predecessor : block->getPredecessorBlocks()) {
+        state = state.meet(getState(predecessor));
+      }
+      setState(block, state);
+      return state != oldState;
+    }
+  };
+};
+
+void VisitUnreachableLifetimeEnds::computeRegion(
+    const SSAPrunedLiveness &liveness) {
+  // Find the non-lifetime-ending boundary of `value`.
   PrunedLivenessBoundary boundary;
   liveness.computeBoundary(boundary);
 
-  BasicBlockWorklist deadEndBlocks(value->getFunction());
   for (SILInstruction *lastUser : boundary.lastUsers) {
     if (liveness.isInterestingUser(lastUser)
         != PrunedLiveness::LifetimeEndingUse) {
-      deadEndBlocks.push(lastUser->getParent());
+      region.insert(lastUser->getParent());
+      starts.insert(lastUser->getParent());
     }
   }
   for (SILBasicBlock *edge : boundary.boundaryEdges) {
-    deadEndBlocks.push(edge);
+    region.insert(edge);
+    starts.insert(edge);
   }
   for (SILNode *deadDef : boundary.deadDefs) {
-    deadEndBlocks.push(deadDef->getParentBlock());
+    region.insert(deadDef->getParentBlock());
+    starts.insert(deadDef->getParentBlock());
+  }
+
+  // Forward walk to find the region in which `value` might be available.
+  BasicBlockWorklist regionWorklist(value->getFunction());
+  // Start the forward walk from the non-lifetime-ending boundary.
+  for (auto *start : region) {
+    regionWorklist.push(start);
   }
-  // Forward CFG walk from the non-lifetime-ending boundary to the unreachable
-  // instructions.
-  while (auto *block = deadEndBlocks.pop()) {
+  while (auto *block = regionWorklist.pop()) {
     if (block->succ_empty()) {
       // This assert will fail unless there are already lifetime-ending
       // instruction on all paths to normal function exits.
-      auto *unreachable = cast<UnreachableInst>(block->getTerminator());
-      visit(unreachable);
+      assert(isa<UnreachableInst>(block->getTerminator()));
     }
     for (auto *successor : block->getSuccessorBlocks()) {
-      deadEndBlocks.pushIfNotVisited(successor);
+      regionWorklist.pushIfNotVisited(successor);
+      region.insert(successor);
     }
   }
 }
 
+void VisitUnreachableLifetimeEnds::propagateAvailablity(Result &result) {
+  // Initialize per-block state.
+  // - all blocks outside of the region are ::Unavailable (automatically
+  //   initialized)
+  // - non-initial in-region blocks are Unknown
+  // - start blocks are ::Available
+  for (auto *block : region) {
+    if (starts.contains(block))
+      result.setState(block, State::Available());
+    else
+      result.setState(block, State::Unknown());
+  }
+
+  BasicBlockWorklist worklist(value->getFunction());
+
+  // Initialize worklist with all participating blocks.
+  //
+  // Only perform dataflow in the non-initial region.  Every initial block is
+  // by definition ::Available.
+  for (auto *block : region) {
+    if (starts.contains(block))
+      continue;
+    worklist.push(block);
+  }
+
+  // Iterate over blocks which are successors of blocks whose state changed.
+  while (auto *block = worklist.popAndForget()) {
+    // Only propagate availability in non-initial, in-region blocks.
+    if (!region.contains(block) || starts.contains(block))
+      continue;
+    auto changed = result.updateState(block);
+    if (!changed) {
+      continue;
+    }
+    // The state has changed.  Propagate the new state into successors.
+    for (auto *successor : block->getSuccessorBlocks()) {
+      worklist.pushIfNotVisited(successor);
+    }
+  }
+}
+
+void VisitUnreachableLifetimeEnds::visitAvailabilityBoundary(
+    Result const &result, llvm::function_ref<void(SILInstruction *)> visit) {
+  for (auto *block : region) {
+    auto available = result.getState(block) == State::Available();
+    if (!available) {
+      continue;
+    }
+    auto hasUnreachableSuccessor = [&]() {
+      // 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()) {
+      continue;
+    }
+    assert(hasUnreachableSuccessor() ||
+           isa<UnreachableInst>(block->getTerminator()));
+    visit(block->getTerminator());
+  }
+}
+} // end anonymous namespace
+
+void OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
+    SILValue value, const SSAPrunedLiveness &liveness,
+    llvm::function_ref<void(SILInstruction *)> visit) {
+
+  VisitUnreachableLifetimeEnds visitor(value);
+
+  visitor.computeRegion(liveness);
+
+  VisitUnreachableLifetimeEnds::Result result(value->getFunction());
+
+  visitor.propagateAvailablity(result);
+
+  visitor.visitAvailabilityBoundary(result, visit);
+}
+
 static bool endLifetimeAtUnreachableBlocks(SILValue value,
                                            const SSAPrunedLiveness &liveness) {
   bool changed = false;
@@ -274,4 +453,3 @@ bool UnreachableLifetimeCompletion::completeLifetimes() {
   }
   return changed;
 }
-

lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp

@@ -261,7 +261,7 @@ void CanonicalizeOSSALifetime::extendLivenessToDeinitBarriers() {
 
   OSSALifetimeCompletion::visitUnreachableLifetimeEnds(
       getCurrentDef(), completeLiveness, [&](auto *unreachable) {
-        recordConsumingUser(unreachable);
+        recordUnreachableLifetimeEnd(unreachable);
         unreachable->visitPriorInstructions([&](auto *inst) {
           liveness->extendToNonUse(inst);
           return true;
@@ -905,8 +905,11 @@ static void insertDestroyBeforeInstruction(SILInstruction *nextInstruction,
                                            InstModCallbacks &callbacks) {
   // OSSALifetimeCompletion: This conditional clause can be deleted with
   // complete lifetimes.
-  if (isa<UnreachableInst>(nextInstruction)) {
-    // Don't create a destroy_value if the next instruction is an unreachable.
+  if (consumes.isUnreachableLifetimeEnd(nextInstruction)) {
+    // Don't create a destroy_value if the next instruction is an unreachable
+    // (or a terminator on the availability boundary of the dead-end region
+    // starting from the non-lifetime-ending boundary of `currentDef`).
+    //
     // If there was a destroy here already, it would be reused.  Avoids
     // creating an explicit destroy of a value which might have an unclosed
     // borrow scope.  Doing so would result in