[MoveOnlyAddressChecker] Fix repr for consumes.

An instruction can consume multiple (discontiguous) fields.  Use a
SmallBitVector to track the fields consumed by an instruction rather
than a TypeTreeLeafRange.

rdar://125103951

Author: nate-chandler

include/swift/SIL/FieldSensitivePrunedLiveness.h

@@ -480,10 +480,10 @@ class FieldSensitivePrunedLiveBlocks {
 
     /// Returns the liveness in \p resultingFoundLiveness. We only return the
     /// bits for endBitNo - startBitNo.
-    void getLiveness(unsigned startBitNo, unsigned endBitNo,
+    void getLiveness(SmallBitVector const &bitsOfInterest,
                      SmallVectorImpl<IsLive> &resultingFoundLiveness) const {
-      for (unsigned i = startBitNo, e = endBitNo; i != e; ++i) {
-        resultingFoundLiveness.push_back(getLiveness(i));
+      for (auto bit : bitsOfInterest.set_bits()) {
+        resultingFoundLiveness.push_back(getLiveness(bit));
       }
     }
 
@@ -601,17 +601,26 @@ class FieldSensitivePrunedLiveBlocks {
   void getBlockLiveness(SILBasicBlock *bb, unsigned startBitNo,
                         unsigned endBitNo,
                         SmallVectorImpl<IsLive> &foundLivenessInfo) const {
+    SmallBitVector bits(*numBitsToTrack);
+    for (auto index = startBitNo; index < endBitNo; ++index) {
+      bits.set(index);
+    }
+    getBlockLiveness(bb, bits, foundLivenessInfo);
+  }
+
+  void getBlockLiveness(SILBasicBlock *bb, SmallBitVector const &bits,
+                        SmallVectorImpl<IsLive> &foundLivenessInfo) const {
     assert(isInitialized());
     auto liveBlockIter = liveBlocks.find(bb);
     if (liveBlockIter == liveBlocks.end()) {
-      for (unsigned i : range(endBitNo - startBitNo)) {
-        (void)i;
+      for (auto bit : bits.set_bits()) {
+        (void)bit;
         foundLivenessInfo.push_back(Dead);
       }
       return;
     }
 
-    liveBlockIter->second.getLiveness(startBitNo, endBitNo, foundLivenessInfo);
+    liveBlockIter->second.getLiveness(bits, foundLivenessInfo);
   }
 
   llvm::StringRef getStringRef(IsLive isLive) const;
@@ -919,6 +928,12 @@ class FieldSensitivePrunedLiveness {
                                 resultingFoundLiveness);
   }
 
+  void getBlockLiveness(SILBasicBlock *bb, SmallBitVector const &bits,
+                        SmallVectorImpl<FieldSensitivePrunedLiveBlocks::IsLive>
+                            &foundLivenessInfo) const {
+    liveBlocks.getBlockLiveness(bb, bits, foundLivenessInfo);
+  }
+
   /// Return the liveness for this specific sub-element of our root value.
   FieldSensitivePrunedLiveBlocks::IsLive
   getBlockLiveness(SILBasicBlock *bb, unsigned subElementNumber) const {
@@ -963,16 +978,16 @@ class FieldSensitivePrunedLiveness {
     return record->isInterestingUser(element);
   }
 
-  /// Whether \p user uses the fields in \p range as indicated by \p kind.
+  /// Whether \p user uses the fields in \p bits as indicated by \p kind.
   bool isInterestingUserOfKind(SILInstruction *user, IsInterestingUser kind,
-                               TypeTreeLeafTypeRange range) const {
+                               SmallBitVector const &bits) const {
     auto *record = getInterestingUser(user);
     if (!record) {
       return kind == IsInterestingUser::NonUser;
     }
 
-    for (auto element : range.getRange()) {
-      if (record->isInterestingUser(element) != kind)
+    for (auto bit : bits.set_bits()) {
+      if (record->isInterestingUser(bit) != kind)
         return false;
     }
     return true;
@@ -1119,10 +1134,10 @@ class FieldSensitivePrunedLiveRange : public FieldSensitivePrunedLiveness {
       : FieldSensitivePrunedLiveness(fn, discoveredBlocks) {}
 
   /// Check if \p inst occurs in between the definition of a def and the
-  /// liveness boundary for bits in \p span.
+  /// liveness boundary for \p bits.
   ///
-  /// NOTE: It is assumed that \p inst is correctly described by span.
-  bool isWithinBoundary(SILInstruction *inst, TypeTreeLeafTypeRange span) const;
+  /// NOTE: It is assumed that \p inst is correctly described by \p bits.
+  bool isWithinBoundary(SILInstruction *inst, SmallBitVector const &bits) const;
 
   /// Customize updateForUse for FieldSensitivePrunedLiveness such that we check
   /// that we consider defs as stopping liveness from being propagated up.

lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp

@@ -875,16 +875,15 @@ static FunctionTest FieldSensitiveSSAUseLivenessTest(
 
 template <typename LivenessWithDefs>
 bool FieldSensitivePrunedLiveRange<LivenessWithDefs>::isWithinBoundary(
-    SILInstruction *inst, TypeTreeLeafTypeRange span) const {
+    SILInstruction *inst, SmallBitVector const &bits) const {
   assert(asImpl().isInitialized());
 
-  PRUNED_LIVENESS_LOG(
-      llvm::dbgs() << "FieldSensitivePrunedLiveRange::isWithinBoundary!\n"
-                   << "Span: ";
-      span.print(llvm::dbgs()); llvm::dbgs() << '\n');
+  PRUNED_LIVENESS_LOG(llvm::dbgs()
+                      << "FieldSensitivePrunedLiveRange::isWithinBoundary!\n"
+                      << "Bits: " << bits << "\n");
 
   // If we do not have any span, return true since we have no counter examples.
-  if (span.empty()) {
+  if (bits.empty()) {
     PRUNED_LIVENESS_LOG(llvm::dbgs() << "    span is empty! Returning true!\n");
     return true;
   }
@@ -894,13 +893,13 @@ bool FieldSensitivePrunedLiveRange<LivenessWithDefs>::isWithinBoundary(
   auto *block = inst->getParent();
 
   SmallVector<IsLive, 8> outVector;
-  getBlockLiveness(block, span, outVector);
+  getBlockLiveness(block, bits, outVector);
 
-  for (auto pair : llvm::enumerate(outVector)) {
-    unsigned bit = span.startEltOffset + pair.index();
+  for (auto bitAndIndex : llvm::enumerate(bits.set_bits())) {
+    unsigned bit = bitAndIndex.value();
     PRUNED_LIVENESS_LOG(llvm::dbgs() << "    Visiting bit: " << bit << '\n');
     bool isLive = false;
-    switch (pair.value()) {
+    switch (outVector[bitAndIndex.index()]) {
     case FieldSensitivePrunedLiveBlocks::DeadToLiveEdge:
     case FieldSensitivePrunedLiveBlocks::Dead:
       PRUNED_LIVENESS_LOG(llvm::dbgs() << "        Dead... continuing!\n");

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -553,7 +553,7 @@ struct UseState {
   /// consider them to be takes since after the transform they must be a take.
   ///
   /// Importantly, these we know are never copied and are only consumed once.
-  llvm::SmallMapVector<SILInstruction *, TypeTreeLeafTypeRange, 4> takeInsts;
+  InstToBitMap takeInsts;
 
   /// A map from a copy_addr, load [copy], or load [take] that we determine
   /// semantically are true copies to the part of the type tree they must copy.
@@ -573,7 +573,7 @@ struct UseState {
   /// We represent these separately from \p takeInsts since:
   ///
   /// 1.
-  llvm::SmallMapVector<SILInstruction *, TypeTreeLeafTypeRange, 4> copyInsts;
+  InstToBitMap copyInsts;
 
   /// A map from an instruction that initializes memory to the description of
   /// the part of the type tree that it initializes.
@@ -654,6 +654,14 @@ struct UseState {
     setAffectedBits(inst, range, initInsts);
   }
 
+  void recordTakeUse(SILInstruction *inst, TypeTreeLeafTypeRange range) {
+    setAffectedBits(inst, range, takeInsts);
+  }
+
+  void recordCopyUse(SILInstruction *inst, TypeTreeLeafTypeRange range) {
+    setAffectedBits(inst, range, copyInsts);
+  }
+
   /// Returns true if this is an instruction that is used by the pass to ensure
   /// that we reinit said argument if we consumed it in a region of code.
   ///
@@ -795,18 +803,18 @@ struct UseState {
     });
   }
 
-  bool isConsume(SILInstruction *inst, TypeTreeLeafTypeRange span) const {
+  bool isConsume(SILInstruction *inst, SmallBitVector const &bits) const {
     {
       auto iter = takeInsts.find(inst);
       if (iter != takeInsts.end()) {
-        if (span.setIntersection(iter->second))
+        if (bits.anyCommon(iter->second))
           return true;
       }
     }
     {
       auto iter = copyInsts.find(inst);
       if (iter != copyInsts.end()) {
-        if (span.setIntersection(iter->second))
+        if (bits.anyCommon(iter->second))
           return true;
       }
     }
@@ -816,34 +824,32 @@ struct UseState {
   bool isCopy(SILInstruction *inst, const SmallBitVector &bv) const {
     auto iter = copyInsts.find(inst);
     if (iter != copyInsts.end()) {
-      for (unsigned index : iter->second.getRange()) {
-        if (bv[index])
-          return true;
-      }
+      if (bv.anyCommon(iter->second))
+        return true;
     }
     return false;
   }
 
-  bool isLivenessUse(SILInstruction *inst, TypeTreeLeafTypeRange span) const {
+  bool isLivenessUse(SILInstruction *inst, SmallBitVector const &bits) const {
     {
       auto iter = nonconsumingUses.find(inst);
       if (iter != nonconsumingUses.end()) {
-        if (span.intersects(iter->second))
+        if (bits.anyCommon(iter->second))
           return true;
       }
     }
     {
       auto iter = borrows.find(inst);
       if (iter != borrows.end()) {
-        if (span.setIntersection(iter->second))
+        if (iter->second.intersects(bits))
           return true;
       }
     }
 
     if (!isReinitToInitConvertibleInst(inst)) {
       auto iter = reinitInsts.find(inst);
       if (iter != reinitInsts.end()) {
-        if (span.intersects(iter->second))
+        if (bits.anyCommon(iter->second))
           return true;
       }
     }
@@ -857,40 +863,19 @@ struct UseState {
     return false;
   }
 
-  bool isInitUse(SILInstruction *inst, TypeTreeLeafTypeRange span) const {
-    {
-      auto iter = initInsts.find(inst);
-      if (iter != initInsts.end()) {
-        if (span.intersects(iter->second))
-          return true;
-      }
-    }
-    if (isReinitToInitConvertibleInst(inst)) {
-      auto iter = reinitInsts.find(inst);
-      if (iter != reinitInsts.end()) {
-        if (span.intersects(iter->second))
-          return true;
-      }
-    }
-    return false;
-  }
-
-  bool isInitUse(SILInstruction *inst, const SmallBitVector &requiredBits,
-                 SmallBitVector &foundInitBits) const {
+  bool isInitUse(SILInstruction *inst, const SmallBitVector &requiredBits) const {
     {
       auto iter = initInsts.find(inst);
       if (iter != initInsts.end()) {
-        foundInitBits = iter->second & requiredBits;
-        if (foundInitBits.any())
+        if (requiredBits.anyCommon(iter->second))
           return true;
       }
     }
 
     if (isReinitToInitConvertibleInst(inst)) {
       auto iter = reinitInsts.find(inst);
       if (iter != reinitInsts.end()) {
-        foundInitBits = iter->second & requiredBits;
-        if (foundInitBits.any())
+        if (requiredBits.anyCommon(iter->second))
           return true;
       }
     }
@@ -2193,10 +2178,10 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
 
     if (copyAddr->isTakeOfSrc()) {
       LLVM_DEBUG(llvm::dbgs() << "Found take: " << *user);
-      useState.takeInsts.insert({user, *leafRange});
+      useState.recordTakeUse(user, *leafRange);
     } else {
       LLVM_DEBUG(llvm::dbgs() << "Found copy: " << *user);
-      useState.copyInsts.insert({user, *leafRange});
+      useState.recordCopyUse(user, *leafRange);
     }
     return true;
   }
@@ -2405,10 +2390,10 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
       // done checking. The load [take] are already complete and good to go.
       if (li->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
         LLVM_DEBUG(llvm::dbgs() << "Found take inst: " << *user);
-        useState.takeInsts.insert({user, *leafRange});
+        useState.recordTakeUse(user, *leafRange);
       } else {
         LLVM_DEBUG(llvm::dbgs() << "Found copy inst: " << *user);
-        useState.copyInsts.insert({user, *leafRange});
+        useState.recordCopyUse(user, *leafRange);
       }
     }
     return true;
@@ -2457,7 +2442,7 @@ bool GatherUsesVisitor::visitUse(Operand *op) {
     }
 
     LLVM_DEBUG(llvm::dbgs() << "Pure consuming use: " << *user);
-    useState.takeInsts.insert({user, *leafRange});
+    useState.recordTakeUse(user, *leafRange);
     return true;
   }
 
@@ -2668,9 +2653,7 @@ struct GlobalLivenessChecker {
   /// Returns true if we emitted any errors.
   bool compute();
 
-  bool testInstVectorLiveness(
-      llvm::SmallMapVector<SILInstruction *, TypeTreeLeafTypeRange, 4>
-          &instsToTest);
+  bool testInstVectorLiveness(UseState::InstToBitMap &instsToTest);
 
   void clear() {
     livenessVector.clear();
@@ -2681,8 +2664,7 @@ struct GlobalLivenessChecker {
 } // namespace
 
 bool GlobalLivenessChecker::testInstVectorLiveness(
-    llvm::SmallMapVector<SILInstruction *, TypeTreeLeafTypeRange, 4>
-        &instsToTest) {
+    UseState::InstToBitMap &instsToTest) {
   bool emittedDiagnostic = false;
 
   for (auto takeInstAndValue : instsToTest) {
@@ -2835,9 +2817,7 @@ bool GlobalLivenessChecker::testInstVectorLiveness(
 #ifndef NDEBUG
           SmallBitVector defBits(addressUseState.getNumSubelements());
           liveness.isDefBlock(block, errorSpan, defBits);
-          SmallBitVector errorSpanBits(addressUseState.getNumSubelements());
-          errorSpan.setBits(errorSpanBits);
-          assert((defBits & errorSpanBits).none() &&
+          assert((defBits & errorSpan).none() &&
                  "If in def block... we are in liveness block");
 #endif
           [[clang::fallthrough]];
@@ -3228,8 +3208,7 @@ void MoveOnlyAddressCheckerPImpl::rewriteUses(
 
   // Check all takes.
   for (auto takeInst : addressUseState.takeInsts) {
-    SmallBitVector bits(liveness.getNumSubElements());
-    takeInst.second.setBits(bits);
+    auto &bits = takeInst.second;
     bool claimedConsume = consumes.claimConsume(takeInst.first, bits);
     (void)claimedConsume;
     if (!claimedConsume) {
@@ -3242,8 +3221,7 @@ void MoveOnlyAddressCheckerPImpl::rewriteUses(
 
   // Then rewrite all copy insts to be takes and claim them.
   for (auto copyInst : addressUseState.copyInsts) {
-    SmallBitVector bits(liveness.getNumSubElements());
-    copyInst.second.setBits(bits);
+    auto &bits = copyInst.second;
     bool claimedConsume = consumes.claimConsume(copyInst.first, bits);
     if (!claimedConsume) {
       llvm::errs()
@@ -3408,7 +3386,7 @@ void ExtendUnconsumedLiveness::run() {
       }
     }
     for (auto pair : addressUseState.takeInsts) {
-      if (pair.second.contains(element)) {
+      if (pair.second.test(element)) {
         destroys[pair.first] = DestroyKind::Take;
       }
     }

lib/SILOptimizer/Mandatory/MoveOnlyBorrowToDestructureUtils.cpp

@@ -604,7 +604,9 @@ void Implementation::checkDestructureUsesOnBoundary() const {
                << "    DestructureNeedingUse: " << *use->getUser());
 
     auto destructureUseSpan = *TypeTreeLeafTypeRange::get(use, getRootValue());
-    if (!liveness.isWithinBoundary(use->getUser(), destructureUseSpan)) {
+    SmallBitVector destructureUseBits(liveness.getNumSubElements());
+    destructureUseSpan.setBits(destructureUseBits);
+    if (!liveness.isWithinBoundary(use->getUser(), destructureUseBits)) {
       LLVM_DEBUG(llvm::dbgs()
                  << "        On boundary or within boundary! No error!\n");
       continue;