[SIL Optimization] Fix a bug in isAliveAtBeginOfBlock function of ValueLifetimeAnalysis and generalize the utility

include/swift/SILOptimizer/Utils/ValueLifetime.h

@@ -22,13 +22,11 @@
 
 namespace swift {
 
-/// This computes the lifetime of a single SILValue.
-///
-/// This does not compute a set of jointly postdominating use points. Instead it
-/// assumes that the value's existing uses already jointly postdominate the
-/// definition. This makes sense for values that are returned +1 from an
-/// instruction, like partial_apply, and therefore must be released on all paths
-/// via strong_release or apply.
+/// Computes the lifetime frontier for a given value with respect to a
+/// given set of uses. The lifetime frontier is the list of instructions
+/// following the last uses. The set of uses can be passed by the clients
+/// of the analysis and can be a super set of the uses of the SILValue
+/// e.g. it can be the set of transitive uses of the SILValue.
 class ValueLifetimeAnalysis {
 public:
 
@@ -70,28 +68,35 @@ class ValueLifetimeAnalysis {
     UsersMustPostDomDef
   };
 
-  /// Computes and returns the lifetime frontier for the value in \p frontier.
+  /// Computes and returns the lifetime frontier for the value in \p frontier
+  /// with respect to the set of uses in the userSet.
   ///
   /// Returns true if all instructions in the frontier could be found in
   /// non-critical edges.
   /// Returns false if some frontier instructions are located on critical edges.
   /// In this case, if \p mode is AllowToModifyCFG, those critical edges are
-  /// split, otherwise nothing is done and the returned \p frontier is not
-  /// valid.
+  /// split, otherwise the returned \p frontier consists of only those
+  /// instructions of the frontier that are not in the critical edges. Note that
+  /// the method getCriticalEdges can be used to retrieve the critical edges.
   ///
   /// If \p deBlocks is provided, all dead-end blocks are ignored. This
   /// prevents unreachable-blocks to be included in the frontier.
   bool computeFrontier(Frontier &frontier, Mode mode,
                        DeadEndBlocks *deBlocks = nullptr);
 
+  ArrayRef<std::pair<TermInst *, unsigned>> getCriticalEdges() {
+    return criticalEdges;
+  }
+
   /// Returns true if the instruction \p Inst is located within the value's
   /// lifetime.
   /// It is assumed that \p inst is located after the value's definition.
   bool isWithinLifetime(SILInstruction *inst);
 
   /// Returns true if the value is alive at the begin of block \p bb.
   bool isAliveAtBeginOfBlock(SILBasicBlock *bb) {
-    return liveBlocks.count(bb) && bb != defValue->getParent();
+    return liveBlocks.count(bb) &&
+           (bb != defValue->getParent() || hasUsersBeforeDef);
   }
 
   /// Checks if there is a dealloc_ref inside the value's live range.
@@ -112,6 +117,14 @@ class ValueLifetimeAnalysis {
   /// provided with the constructor.
   llvm::SmallPtrSet<SILInstruction *, 16> userSet;
 
+  /// Indicates whether the basic block containing def has users of def that
+  /// precede def. This field is initialized by propagateLiveness.
+  bool hasUsersBeforeDef;
+
+  /// Critical edges that couldn't be split to compute the frontier. This could
+  /// be non-empty when the analysis is invoked with DontModifyCFG mode.
+  llvm::SmallVector<std::pair<TermInst *, unsigned>, 16> criticalEdges;
+
   /// Propagates the liveness information up the control flow graph.
   void propagateLiveness();
 

lib/SILOptimizer/Utils/ValueLifetime.cpp

@@ -18,6 +18,7 @@ using namespace swift;
 
 void ValueLifetimeAnalysis::propagateLiveness() {
   assert(liveBlocks.empty() && "frontier computed twice");
+  assert(!userSet.count(defValue) && "definition cannot be its own use");
 
   auto defBB = defValue->getParentBlock();
   llvm::SmallVector<SILBasicBlock *, 64> worklist;
@@ -42,13 +43,16 @@ void ValueLifetimeAnalysis::propagateLiveness() {
       numUsersBeforeDef--;
   }
 
+  // Initialize the hasUsersBeforeDef field.
+  hasUsersBeforeDef = numUsersBeforeDef > 0;
+
   // Now propagate liveness backwards until we hit the block that defines the
   // value.
   while (!worklist.empty()) {
     auto *bb = worklist.pop_back_val();
 
     // Don't go beyond the definition.
-    if (bb == defBB && numUsersBeforeDef == 0)
+    if (bb == defBB && !hasUsersBeforeDef)
       continue;
 
     for (SILBasicBlock *Pred : bb->getPredecessorBlocks()) {
@@ -93,13 +97,34 @@ bool ValueLifetimeAnalysis::computeFrontier(Frontier &frontier, Mode mode,
 
     bool liveInSucc = false;
     bool deadInSucc = false;
+    bool usedAndRedefinedInSucc = false;
     for (const SILSuccessor &succ : bb->getSuccessors()) {
       if (isAliveAtBeginOfBlock(succ)) {
         liveInSucc = true;
+        if (succ == defValue->getParent()) {
+          // Here, the basic block bb uses the value but also redefines the
+          // value inside bb. The new value could be used by the successors
+          // of succ and therefore could be live at the end of succ as well.
+          usedAndRedefinedInSucc = true;
+        }
       } else if (!deBlocks || !deBlocks->isDeadEnd(succ)) {
         deadInSucc = true;
       }
     }
+    if (usedAndRedefinedInSucc) {
+      // Here, the basic block bb uses the value and later redefines the value.
+      // Therefore, this value's lifetime ends after its last use preceding the
+      // re-definition of the value.
+      auto ii = defValue->getReverseIterator();
+      for (; ii != bb->rend(); ++ii) {
+        if (userSet.count(&*ii)) {
+          frontier.push_back(&*std::next(ii));
+          break;
+        }
+      }
+      assert(ii != bb->rend() &&
+             "There must be a user in bb before definition");
+    }
     if (!liveInSucc) {
       // The value is not live in any of the successor blocks. This means the
       // block contains a last use of the value. The next instruction after
@@ -141,15 +166,17 @@ bool ValueLifetimeAnalysis::computeFrontier(Frontier &frontier, Mode mode,
       }
     }
     if (needSplit) {
-      if (mode == DontModifyCFG)
-        return false;
       // We need to split the critical edge to create a frontier instruction.
       unhandledFrontierBlocks.insert(frontierBB);
     } else {
       // The first instruction of the exit-block is part of the frontier.
       frontier.push_back(&*frontierBB->begin());
     }
   }
+  if (unhandledFrontierBlocks.size() == 0) {
+    return true;
+  }
+
   // Split critical edges from the lifetime region to not yet handled frontier
   // blocks.
   for (SILBasicBlock *frontierPred : liveOutBlocks) {
@@ -163,12 +190,17 @@ bool ValueLifetimeAnalysis::computeFrontier(Frontier &frontier, Mode mode,
 
     for (unsigned i = 0, e = succBlocks.size(); i != e; ++i) {
       if (unhandledFrontierBlocks.count(succBlocks[i])) {
-        assert(mode == AllowToModifyCFG);
         assert(isCriticalEdge(term, i) && "actually not a critical edge?");
+        noCriticalEdges = false;
+        if (mode != AllowToModifyCFG) {
+          // If the CFG need not be modified, just record the critical edge and
+          // continue.
+          this->criticalEdges.push_back({term, i});
+          continue;
+        }
         SILBasicBlock *newBlock = splitEdge(term, i);
         // The single terminator instruction is part of the frontier.
         frontier.push_back(&*newBlock->begin());
-        noCriticalEdges = false;
       }
     }
   }