Adding bidirectional iterator support to const_set_bits_iterator_impl

Author: FranciscoThiesen

llvm/include/llvm/ADT/BitVector.h

@@ -40,12 +40,25 @@ template <typename BitVectorT> class const_set_bits_iterator_impl {
     Current = Parent.find_next(Current);
   }
 
+  void retreat() {
+    // For bidirectional iteration to work with reverse_iterator,
+    // we need to handle the case where Current might be at end (-1)
+    // or at a position where we need to find the previous set bit.
+    if (Current == -1) {
+      // If we're at the end, go to the last set bit
+      Current = Parent.find_last();
+    } else {
+      // Otherwise find the previous set bit before Current
+      Current = Parent.find_prev(Current);
+    }
+  }
+
 public:
-  using iterator_category = std::forward_iterator_tag;
+  using iterator_category = std::bidirectional_iterator_tag;
   using difference_type   = std::ptrdiff_t;
-  using value_type        = int;
-  using pointer           = value_type*;
-  using reference         = value_type&;
+  using value_type        = unsigned;
+  using pointer           = const value_type*;
+  using reference         = value_type;
 
   const_set_bits_iterator_impl(const BitVectorT &Parent, int Current)
       : Parent(Parent), Current(Current) {}
@@ -64,6 +77,17 @@ template <typename BitVectorT> class const_set_bits_iterator_impl {
     return *this;
   }
 
+  const_set_bits_iterator_impl operator--(int) {
+    auto Prev = *this;
+    retreat();
+    return Prev;
+  }
+
+  const_set_bits_iterator_impl &operator--() {
+    retreat();
+    return *this;
+  }
+
   unsigned operator*() const { return Current; }
 
   bool operator==(const const_set_bits_iterator_impl &Other) const {

mlir/lib/Transforms/RemoveDeadValues.cpp

@@ -772,7 +772,8 @@ static void cleanUpDeadVals(RDVFinalCleanupList &list) {
   LDBG() << "Cleaning up " << list.operands.size() << " operand lists";
   for (OperationToCleanup &o : list.operands) {
     // Handle call-specific cleanup only when we have a cached callee reference.
-    // This avoids expensive symbol lookup and is defensive against future changes.
+    // This avoids expensive symbol lookup and is defensive against future
+    // changes.
     bool handledAsCall = false;
     if (o.callee && isa<CallOpInterface>(o.op)) {
       auto call = cast<CallOpInterface>(o.op);
@@ -781,11 +782,9 @@ static void cleanUpDeadVals(RDVFinalCleanupList &list) {
         const BitVector &deadArgIdxs = it->second;
         MutableOperandRange args = call.getArgOperandsMutable();
         // First, erase the call arguments corresponding to erased callee
-        // args.
-        for (int i = static_cast<int>(args.size()) - 1; i >= 0; --i) {
-          if (i < static_cast<int>(deadArgIdxs.size()) && deadArgIdxs.test(i))
-            args.erase(i);
-        }
+        // args. We iterate backwards to preserve indices.
+        for (unsigned argIdx : llvm::reverse(deadArgIdxs.set_bits()))
+          args.erase(argIdx);
         // If this operand cleanup entry also has a generic nonLive bitvector,
         // clear bits for call arguments we already erased above to avoid
         // double-erasing (which could impact other segments of ops with