[pruned-liveness] Change PrunedLiveBlocks so that it can propagate multiple bits of liveness rather than just one.

This is useful for making a form of PrunedLiveness that is sensitive to address
fields. I wired up the current PrunedLiveness to just use PrunedLiveBlocks with
num bits set to 1.

Author: gottesmm

include/swift/SIL/PrunedLiveness.h

@@ -115,6 +115,11 @@ class DeadEndBlocks;
 /// which the def block's predecessors incorrectly remain dead. This situation
 /// could be handled by adding an updateForUseBeforeFirstDef() API.
 ///
+/// We allow for multiple bits of liveness information to be tracked by
+/// internally using a SmallBitVector. We default to only tracking a single
+/// bit. The multiple bit tracking is useful when tracking state for multiple
+/// fields of the same root value.
+///
 /// TODO: This can be made space-efficient if all clients can maintain a block
 /// numbering so liveness info can be represented as bitsets across the blocks.
 class PrunedLiveBlocks {
@@ -136,19 +141,24 @@ class PrunedLiveBlocks {
   enum IsLive { Dead, LiveWithin, LiveOut };
 
 private:
-  // Map all blocks in which current def is live to a flag indicating whether
-  // the value is also liveout of the block.
-  llvm::SmallDenseMap<SILBasicBlock *, bool, 4> liveBlocks;
+  /// Map all blocks in which current def is live to a SmallBitVector indicating
+  /// whether the value represented by said bit is also liveout of the block.
+  llvm::SmallDenseMap<SILBasicBlock *, SmallBitVector, 4> liveBlocks;
+
+  /// Number of bits of liveness to track. By default 1. Used to track multiple
+  /// liveness bits.
+  unsigned numBitsToTrack;
 
-  // Optional vector of live blocks for clients that deterministically iterate.
+  /// Optional vector of live blocks for clients that deterministically iterate.
   SmallVectorImpl<SILBasicBlock *> *discoveredBlocks;
 
-  // Once the first use has been seen, no definitions can be added.
+  /// Once the first use has been seen, no definitions can be added.
   SWIFT_ASSERT_ONLY_DECL(bool seenUse = false);
 
 public:
-  PrunedLiveBlocks(SmallVectorImpl<SILBasicBlock *> *discoveredBlocks = nullptr)
-      : discoveredBlocks(discoveredBlocks) {
+  PrunedLiveBlocks(unsigned numBitsToTrack,
+                   SmallVectorImpl<SILBasicBlock *> *discoveredBlocks = nullptr)
+      : numBitsToTrack(numBitsToTrack), discoveredBlocks(discoveredBlocks) {
     assert(!discoveredBlocks || discoveredBlocks->empty());
   }
 
@@ -167,36 +177,38 @@ class PrunedLiveBlocks {
     return *discoveredBlocks;
   }
 
-  void initializeDefBlock(SILBasicBlock *defBB) {
-    markBlockLive(defBB, LiveWithin);
+  void initializeDefBlock(SILBasicBlock *defBB, unsigned bitNo) {
+    markBlockLive(defBB, LiveWithin, bitNo);
   }
 
   /// Update this liveness result for a single use.
-  IsLive updateForUse(SILInstruction *user);
+  IsLive updateForUse(SILInstruction *user, unsigned bitNo);
 
-  IsLive getBlockLiveness(SILBasicBlock *bb) const {
+  IsLive getBlockLiveness(SILBasicBlock *bb, unsigned bitNo) const {
     auto liveBlockIter = liveBlocks.find(bb);
     if (liveBlockIter == liveBlocks.end())
       return Dead;
-    return liveBlockIter->second ? LiveOut : LiveWithin;
+    assert(liveBlockIter->second.size() == 1);
+    return liveBlockIter->second[bitNo] ? LiveOut : LiveWithin;
   }
 
 protected:
-  void markBlockLive(SILBasicBlock *bb, IsLive isLive) {
+  void markBlockLive(SILBasicBlock *bb, IsLive isLive, unsigned bitNo) {
     assert(isLive != Dead && "erasing live blocks isn't implemented.");
     bool isLiveOut = (isLive == LiveOut);
     auto iterAndInserted =
-      liveBlocks.insert(std::make_pair(bb, isLiveOut));
+        liveBlocks.insert(std::make_pair(bb, SmallBitVector(numBitsToTrack)));
     if (iterAndInserted.second) {
+      iterAndInserted.first->getSecond()[bitNo] = isLiveOut;
       if (discoveredBlocks)
         discoveredBlocks->push_back(bb);
     } else if (isLiveOut) {
       // Update the existing entry to be live-out.
-      iterAndInserted.first->getSecond() = true;
+      iterAndInserted.first->getSecond()[bitNo] = true;
     }
   }
 
-  void computeUseBlockLiveness(SILBasicBlock *userBB);
+  void computeUseBlockLiveness(SILBasicBlock *userBB, unsigned bitNo);
 };
 
 /// PrunedLiveness tracks PrunedLiveBlocks along with "interesting" use
@@ -252,7 +264,7 @@ class PrunedLiveness {
   PrunedLiveness(SmallVectorImpl<SILBasicBlock *> *discoveredBlocks = nullptr,
                  SmallSetVector<SILInstruction *, 8>
                      *nonLifetimeEndingUsesInLiveOut = nullptr)
-      : liveBlocks(discoveredBlocks),
+      : liveBlocks(1 /*num bits*/, discoveredBlocks),
         nonLifetimeEndingUsesInLiveOut(nonLifetimeEndingUsesInLiveOut) {}
 
   bool empty() const {
@@ -317,7 +329,7 @@ class PrunedLiveness {
   }
 
   void initializeDefBlock(SILBasicBlock *defBB) {
-    liveBlocks.initializeDefBlock(defBB);
+    liveBlocks.initializeDefBlock(defBB, 0);
   }
 
   /// For flexibility, \p lifetimeEnding is provided by the
@@ -335,7 +347,7 @@ class PrunedLiveness {
   void extendAcrossLiveness(PrunedLiveness &otherLiveness);
 
   PrunedLiveBlocks::IsLive getBlockLiveness(SILBasicBlock *bb) const {
-    return liveBlocks.getBlockLiveness(bb);
+    return liveBlocks.getBlockLiveness(bb, 0);
   }
 
   enum IsInterestingUser { NonUser, NonLifetimeEndingUse, LifetimeEndingUse };

lib/SIL/Utils/PrunedLiveness.cpp

@@ -19,10 +19,11 @@ using namespace swift;
 
 /// Mark blocks live during a reverse CFG traversal from one specific block
 /// containing a user.
-void PrunedLiveBlocks::computeUseBlockLiveness(SILBasicBlock *userBB) {
-  // If we are visiting this block, then it is not already LiveOut. Mark it
+void PrunedLiveBlocks::computeUseBlockLiveness(SILBasicBlock *userBB,
+                                               unsigned bitNo) {
+  // If, we are visiting this block, then it is not already LiveOut. Mark it
   // LiveWithin to indicate a liveness boundary within the block.
-  markBlockLive(userBB, LiveWithin);
+  markBlockLive(userBB, LiveWithin, bitNo);
 
   SmallVector<SILBasicBlock *, 8> predBBWorklist({userBB});
   while (!predBBWorklist.empty()) {
@@ -33,12 +34,12 @@ void PrunedLiveBlocks::computeUseBlockLiveness(SILBasicBlock *userBB) {
     // Traversal terminates at any previously visited block, including the
     // blocks initialized as definition blocks.
     for (auto *predBB : bb->getPredecessorBlocks()) {
-      switch (getBlockLiveness(predBB)) {
+      switch (getBlockLiveness(predBB, bitNo)) {
       case Dead:
         predBBWorklist.push_back(predBB);
         LLVM_FALLTHROUGH;
       case LiveWithin:
-        markBlockLive(predBB, LiveOut);
+        markBlockLive(predBB, LiveOut, bitNo);
         break;
       case LiveOut:
         break;
@@ -52,20 +53,21 @@ void PrunedLiveBlocks::computeUseBlockLiveness(SILBasicBlock *userBB) {
 /// Return the updated liveness of the \p use block (LiveOut or LiveWithin).
 ///
 /// Terminators are not live out of the block.
-PrunedLiveBlocks::IsLive PrunedLiveBlocks::updateForUse(SILInstruction *user) {
+PrunedLiveBlocks::IsLive PrunedLiveBlocks::updateForUse(SILInstruction *user,
+                                                        unsigned bitNo) {
   SWIFT_ASSERT_ONLY(seenUse = true);
 
   auto *bb = user->getParent();
-  switch (getBlockLiveness(bb)) {
+  switch (getBlockLiveness(bb, bitNo)) {
   case LiveOut:
     return LiveOut;
   case LiveWithin:
     return LiveWithin;
   case Dead: {
     // This use block has not yet been marked live. Mark it and its predecessor
     // blocks live.
-    computeUseBlockLiveness(bb);
-    return getBlockLiveness(bb);
+    computeUseBlockLiveness(bb, bitNo);
+    return getBlockLiveness(bb, bitNo);
   }
   }
   llvm_unreachable("covered switch");
@@ -76,7 +78,7 @@ PrunedLiveBlocks::IsLive PrunedLiveBlocks::updateForUse(SILInstruction *user) {
 //===----------------------------------------------------------------------===//
 
 void PrunedLiveness::updateForUse(SILInstruction *user, bool lifetimeEnding) {
-  auto useBlockLive = liveBlocks.updateForUse(user);
+  auto useBlockLive = liveBlocks.updateForUse(user, 0);
   // Record all uses of blocks on the liveness boundary. For blocks marked
   // LiveWithin, the boundary is considered to be the last use in the block.
   //