[pruned-liveness] Make some small fixes that only are exposed when working with multi-bit code

include/swift/Basic/FrozenMultiMap.h

@@ -246,6 +246,19 @@ class FrozenMultiMap {
     auto optRange = makeOptionalTransformRange(baseRange, ToNonErasedValues());
     return makeTransformRange(optRange, PairWithTypeErasedOptionalSecondElt());
   }
+
+  /// Returns true if all values for all keys have been deleted.
+  ///
+  /// This is intended to be used in use cases where a frozen multi map is
+  /// filled up with a multi-map and then as we process keys, we delete values
+  /// we have handled. In certain cases, one wishes to validate after processing
+  /// that all values for all keys were properly handled. One cannot perform
+  /// this operation with getRange() in a nice way.
+  bool allValuesHaveBeenDeleted() const {
+    return llvm::all_of(storage, [](const std::pair<Key, Optional<Value>> &pair) {
+      return !pair.second.hasValue();
+    });
+  }
 };
 
 template <typename Key, typename Value, typename Storage>

include/swift/Basic/SmallBitVector.h

@@ -0,0 +1,30 @@
+//===--- SmallBitVector.h -------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_BASIC_SMALLBITVECTOR_H
+#define SWIFT_BASIC_SMALLBITVECTOR_H
+
+#include "llvm/ADT/SmallBitVector.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace llvm {
+
+inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
+                                     const SmallBitVector &bv) {
+  for (unsigned i = 0, e = bv.size(); i != e; ++i)
+    os << (bv[i] ? '1' : '0');
+  return os;
+}
+
+} // namespace llvm
+
+#endif

include/swift/SIL/PrunedLiveness.h

@@ -196,26 +196,19 @@ class PrunedLiveBlocks {
 
     unsigned size() const { return bits.size() / 2; }
 
-    // FIXME: specialize this for scalar liveness, which is the critical path
-    // for all OSSA utilities.
     IsLive getLiveness(unsigned bitNo) const {
-      SmallVector<IsLive, 1> foundLiveness;
-      getLiveness(bitNo, bitNo + 1, foundLiveness);
-      return foundLiveness[0];
+      if (!bits[bitNo * 2])
+        return IsLive::Dead;
+      return bits[bitNo * 2 + 1] ? LiveOut : LiveWithin;
     }
 
+    /// Returns the liveness in \p resultingFoundLiveness. We only return the
+    /// bits for endBitNo - startBitNo.
     void getLiveness(unsigned startBitNo, unsigned endBitNo,
                      SmallVectorImpl<IsLive> &resultingFoundLiveness) const {
       unsigned actualStartBitNo = startBitNo * 2;
       unsigned actualEndBitNo = endBitNo * 2;
 
-      // NOTE: We pad both before/after with Dead to ensure that we are
-      // returning an array that acts as a bit mask and thus can be directly
-      // compared against other such bitmasks. This invariant is used when
-      // computing boundaries.
-      for (unsigned i = 0; i != startBitNo; ++i) {
-        resultingFoundLiveness.push_back(Dead);
-      }
       for (unsigned i = actualStartBitNo, e = actualEndBitNo; i != e; i += 2) {
         if (!bits[i]) {
           resultingFoundLiveness.push_back(Dead);
@@ -224,9 +217,6 @@ class PrunedLiveBlocks {
 
         resultingFoundLiveness.push_back(bits[i + 1] ? LiveOut : LiveWithin);
       }
-      for (unsigned i = endBitNo, e = size(); i != e; ++i) {
-        resultingFoundLiveness.push_back(Dead);
-      }
     }
 
     void setLiveness(unsigned startBitNo, unsigned endBitNo, IsLive isLive) {
@@ -291,9 +281,12 @@ class PrunedLiveBlocks {
 
   /// Update this liveness result for a single use.
   IsLive updateForUse(SILInstruction *user, unsigned bitNo) {
-    SmallVector<IsLive, 1> resultingLiveness;
-    updateForUse(user, bitNo, bitNo + 1, resultingLiveness);
-    return resultingLiveness[0];
+    auto *block = user->getParent();
+    auto liveness = getBlockLiveness(block, bitNo);
+    if (liveness != Dead)
+      return liveness;
+    computeScalarUseBlockLiveness(block, bitNo);
+    return getBlockLiveness(block, bitNo);
   }
 
   /// Update this range of liveness results for a single use.
@@ -302,19 +295,22 @@ class PrunedLiveBlocks {
                     SmallVectorImpl<IsLive> &resultingLiveness);
 
   IsLive getBlockLiveness(SILBasicBlock *bb, unsigned bitNo) const {
-    SmallVector<IsLive, 1> isLive;
-    getBlockLiveness(bb, bitNo, bitNo + 1, isLive);
-    return isLive[0];
+    auto liveBlockIter = liveBlocks.find(bb);
+    if (liveBlockIter == liveBlocks.end()) {
+      return Dead;
+    }
+
+    return liveBlockIter->second.getLiveness(bitNo);
   }
 
-  // FIXME: This API should directly return the live bitset. The live bitset
-  // type should have an api for querying and iterating over the live fields.
+  /// FIXME: This API should directly return the live bitset. The live bitset
+  /// type should have an api for querying and iterating over the live fields.
   void getBlockLiveness(SILBasicBlock *bb, unsigned startBitNo,
                         unsigned endBitNo,
                         SmallVectorImpl<IsLive> &foundLivenessInfo) const {
     auto liveBlockIter = liveBlocks.find(bb);
     if (liveBlockIter == liveBlocks.end()) {
-      for (unsigned i : range(numBitsToTrack)) {
+      for (unsigned i : range(endBitNo - startBitNo)) {
         (void)i;
         foundLivenessInfo.push_back(Dead);
       }
@@ -330,11 +326,6 @@ class PrunedLiveBlocks {
 
 protected:
   void markBlockLive(SILBasicBlock *bb, unsigned bitNo, IsLive isLive) {
-    markBlockLive(bb, bitNo, bitNo + 1, isLive);
-  }
-
-  void markBlockLive(SILBasicBlock *bb, unsigned startBitNo, unsigned endBitNo,
-                     IsLive isLive) {
     assert(isLive != Dead && "erasing live blocks isn't implemented.");
     auto iterAndInserted =
         liveBlocks.insert(std::make_pair(bb, LivenessSmallBitVector()));
@@ -344,18 +335,44 @@ class PrunedLiveBlocks {
       // we have more than SmallBitVector's small size number of bits.
       auto &insertedBV = iterAndInserted.first->getSecond();
       insertedBV.init(numBitsToTrack);
-      insertedBV.setLiveness(startBitNo, endBitNo, isLive);
+      insertedBV.setLiveness(bitNo, bitNo + 1, isLive);
       if (discoveredBlocks)
         discoveredBlocks->push_back(bb);
-    } else if (isLive == LiveOut) {
-      // Update the existing entry to be live-out.
-      iterAndInserted.first->getSecond().setLiveness(startBitNo, endBitNo,
-                                                     LiveOut);
+    } else {
+      // If we are dead, always update to the new liveness.
+      switch (iterAndInserted.first->getSecond().getLiveness(bitNo)) {
+      case Dead:
+        iterAndInserted.first->getSecond().setLiveness(bitNo, bitNo + 1,
+                                                       isLive);
+        break;
+      case LiveWithin:
+        if (isLive == LiveOut) {
+          // Update the existing entry to be live-out.
+          iterAndInserted.first->getSecond().setLiveness(bitNo, bitNo + 1,
+                                                         LiveOut);
+        }
+        break;
+      case LiveOut:
+        break;
+      }
     }
   }
 
-  void computeUseBlockLiveness(SILBasicBlock *userBB, unsigned startBitNo,
-                               unsigned endBitNo);
+  void markBlockLive(SILBasicBlock *bb, unsigned startBitNo, unsigned endBitNo,
+                     IsLive isLive) {
+    for (unsigned index : range(startBitNo, endBitNo)) {
+      markBlockLive(bb, index, isLive);
+    }
+  }
+
+private:
+  /// A helper routine that as a fast path handles the scalar case. We do not
+  /// handle the mult-bit case today since the way the code is written today
+  /// assumes we process a bit at a time.
+  ///
+  /// TODO: Make a multi-bit query for efficiency reasons.
+  void computeScalarUseBlockLiveness(SILBasicBlock *userBB,
+                                     unsigned startBitNo);
 };
 
 /// If inner borrows are 'Contained', then liveness is fully described by the

lib/SIL/Utils/PrunedLiveness.cpp

@@ -21,16 +21,12 @@
 
 using namespace swift;
 
-/// Mark blocks live during a reverse CFG traversal from one specific block
-/// containing a user.
-void PrunedLiveBlocks::computeUseBlockLiveness(SILBasicBlock *userBB,
-                                               unsigned startBitNo,
-                                               unsigned endBitNo) {
+void PrunedLiveBlocks::computeScalarUseBlockLiveness(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, startBitNo, endBitNo, LiveWithin);
+  markBlockLive(userBB, bitNo, LiveWithin);
 
-  SmallVector<IsLive, 8> predLivenessInfo;
   BasicBlockWorklist worklist(userBB->getFunction());
   worklist.push(userBB);
 
@@ -40,19 +36,15 @@ void PrunedLiveBlocks::computeUseBlockLiveness(SILBasicBlock *userBB,
     // Traversal terminates at any previously visited block, including the
     // blocks initialized as definition blocks.
     for (auto *predBlock : block->getPredecessorBlocks()) {
-      SWIFT_DEFER { predLivenessInfo.clear(); };
-      getBlockLiveness(predBlock, startBitNo, endBitNo, predLivenessInfo);
-      for (unsigned i : indices(predLivenessInfo)) {
-        switch (predLivenessInfo[i]) {
-        case Dead:
-          worklist.pushIfNotVisited(predBlock);
-          LLVM_FALLTHROUGH;
-        case LiveWithin:
-          markBlockLive(predBlock, startBitNo, endBitNo, LiveOut);
-          break;
-        case LiveOut:
-          break;
-        }
+      switch (getBlockLiveness(predBlock, bitNo)) {
+      case Dead:
+        worklist.pushIfNotVisited(predBlock);
+        LLVM_FALLTHROUGH;
+      case LiveWithin:
+        markBlockLive(predBlock, bitNo, LiveOut);
+        break;
+      case LiveOut:
+        break;
       }
     }
   }
@@ -66,22 +58,26 @@ void PrunedLiveBlocks::computeUseBlockLiveness(SILBasicBlock *userBB,
 void PrunedLiveBlocks::updateForUse(
     SILInstruction *user, unsigned startBitNo, unsigned endBitNo,
     SmallVectorImpl<IsLive> &resultingLivenessInfo) {
+  resultingLivenessInfo.clear();
+
   SWIFT_ASSERT_ONLY(seenUse = true);
 
   auto *bb = user->getParent();
   getBlockLiveness(bb, startBitNo, endBitNo, resultingLivenessInfo);
 
-  for (auto isLive : resultingLivenessInfo) {
-    switch (isLive) {
+  for (auto pair : llvm::enumerate(resultingLivenessInfo)) {
+    unsigned index = pair.index();
+    unsigned specificBitNo = startBitNo + index;
+    switch (pair.value()) {
     case LiveOut:
     case LiveWithin:
       continue;
     case Dead: {
       // This use block has not yet been marked live. Mark it and its
       // predecessor blocks live.
-      computeUseBlockLiveness(bb, startBitNo, endBitNo);
-      resultingLivenessInfo.clear();
-      return getBlockLiveness(bb, startBitNo, endBitNo, resultingLivenessInfo);
+      computeScalarUseBlockLiveness(bb, specificBitNo);
+      resultingLivenessInfo.push_back(getBlockLiveness(bb, specificBitNo));
+      continue;
     }
     }
     llvm_unreachable("covered switch");
@@ -179,10 +175,13 @@ void PrunedLiveBlocks::print(llvm::raw_ostream &OS) const {
   if (!discoveredBlocks) {
     OS << "No deterministic live block list\n";
   }
+  SmallVector<IsLive, 8> isLive;
   for (auto *block : *discoveredBlocks) {
     block->printAsOperand(OS);
-    OS
-      << ": " << getStringRef(this->getBlockLiveness(block, 0)) << "\n";
+    OS << ": ";
+    for (unsigned i : range(getNumBitsToTrack()))
+      OS << getStringRef(this->getBlockLiveness(block, i)) << ", ";
+    OS << "\n";
   }
 }
 

lib/SILOptimizer/Mandatory/MoveKillsCopyableAddressesChecker.cpp

@@ -141,6 +141,7 @@
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/FrozenMultiMap.h"
 #include "swift/Basic/GraphNodeWorklist.h"
+#include "swift/Basic/SmallBitVector.h"
 #include "swift/SIL/BasicBlockBits.h"
 #include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SIL/Consumption.h"
@@ -183,15 +184,6 @@ static void diagnose(ASTContext &Context, SourceLoc loc, Diag<T...> diag,
   Context.Diags.diagnose(loc, diag, std::forward<U>(args)...);
 }
 
-namespace llvm {
-llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const SmallBitVector &bv) {
-  for (unsigned index : range(bv.size())) {
-    os << (bv[index] ? 1 : 0);
-  }
-  return os;
-}
-} // namespace llvm
-
 static SourceLoc getSourceLocFromValue(SILValue value) {
   if (auto *defInst = value->getDefiningInstruction())
     return defInst->getLoc().getSourceLoc();