Move the IndexedAccessorIterator to STLExtras to allow for reuse.

This iterator is useful for implementing random access iterators based upon an index and an object pointer. Moving it to STLExtras allows for reuse elsewhere throughout the codebase, e.g. simplifying the DenseElementsAttr iterators.

PiperOrigin-RevId: 255020377

mlir/include/mlir/IR/Attributes.h

@@ -549,40 +549,18 @@ public:
 
   /// A utility iterator that allows walking over the internal raw APInt values.
   class IntElementIterator
-      : public llvm::iterator_facade_base<IntElementIterator,
-                                          std::bidirectional_iterator_tag,
-                                          APInt, std::ptrdiff_t, APInt, APInt> {
+      : public indexed_accessor_iterator<IntElementIterator, const char *,
+                                         APInt, APInt, APInt> {
   public:
-    /// Iterator movement.
-    IntElementIterator &operator++() {
-      ++index;
-      return *this;
-    }
-    IntElementIterator &operator--() {
-      --index;
-      return *this;
-    }
-
     /// Accesses the raw APInt value at this iterator position.
     APInt operator*() const;
 
-    /// Iterator equality.
-    bool operator==(const IntElementIterator &rhs) const {
-      return rawData == rhs.rawData && index == rhs.index;
-    }
-
   private:
     friend DenseElementsAttr;
 
     /// Constructs a new iterator.
     IntElementIterator(DenseElementsAttr attr, size_t index);
 
-    /// The base address of the raw data buffer.
-    const char *rawData;
-
-    /// The current element index.
-    size_t index;
-
     /// The bitwidth of the element type.
     size_t bitWidth;
   };

mlir/include/mlir/IR/Block.h

@@ -374,12 +374,13 @@ inline auto Block::getPredecessors() -> llvm::iterator_range<pred_iterator> {
 
 /// This template implements the successor iterators for Block.
 class SuccessorIterator final
-    : public IndexedAccessorIterator<SuccessorIterator, Block, Block> {
+    : public indexed_accessor_iterator<SuccessorIterator, Block *, Block *,
+                                       Block *, Block *> {
 public:
   /// Initializes the result iterator to the specified index.
   SuccessorIterator(Block *object, unsigned index)
-      : IndexedAccessorIterator<SuccessorIterator, Block, Block>(object,
-                                                                 index) {}
+      : indexed_accessor_iterator<SuccessorIterator, Block *, Block *, Block *,
+                                  Block *>(object, index) {}
 
   SuccessorIterator(const SuccessorIterator &other)
       : SuccessorIterator(other.object, other.index) {}

mlir/include/mlir/IR/Operation.h

@@ -522,12 +522,13 @@ inline raw_ostream &operator<<(raw_ostream &os, Operation &op) {
 /// This class implements the const/non-const operand iterators for the
 /// Operation class in terms of getOperand(idx).
 class OperandIterator final
-    : public IndexedAccessorIterator<OperandIterator, Operation, Value> {
+    : public indexed_accessor_iterator<OperandIterator, Operation *, Value *,
+                                       Value *, Value *> {
 public:
   /// Initializes the operand iterator to the specified operand index.
   OperandIterator(Operation *object, unsigned index)
-      : IndexedAccessorIterator<OperandIterator, Operation, Value>(object,
-                                                                   index) {}
+      : indexed_accessor_iterator<OperandIterator, Operation *, Value *,
+                                  Value *, Value *>(object, index) {}
 
   Value *operator*() const { return this->object->getOperand(this->index); }
 };
@@ -575,12 +576,13 @@ inline auto Operation::getOperandTypes() -> operand_type_range {
 /// This class implements the result iterators for the Operation class
 /// in terms of getResult(idx).
 class ResultIterator final
-    : public IndexedAccessorIterator<ResultIterator, Operation, Value> {
+    : public indexed_accessor_iterator<ResultIterator, Operation *, Value *,
+                                       Value *, Value *> {
 public:
   /// Initializes the result iterator to the specified index.
   ResultIterator(Operation *object, unsigned index)
-      : IndexedAccessorIterator<ResultIterator, Operation, Value>(object,
-                                                                  index) {}
+      : indexed_accessor_iterator<ResultIterator, Operation *, Value *, Value *,
+                                  Value *>(object, index) {}
 
   Value *operator*() const { return this->object->getResult(this->index); }
 };

mlir/include/mlir/IR/Value.h

@@ -171,43 +171,6 @@ private:
   Operation *const owner;
 };
 
-/// This is a helper template used to implement an iterator that contains a
-/// pointer to some object and an index into it.  The iterator moves the
-/// index but keeps the object constant.
-template <typename ConcreteType, typename ObjectType, typename ElementType>
-class IndexedAccessorIterator
-    : public llvm::iterator_facade_base<
-          ConcreteType, std::random_access_iterator_tag, ElementType *,
-          std::ptrdiff_t, ElementType *, ElementType *> {
-public:
-  ptrdiff_t operator-(const IndexedAccessorIterator &rhs) const {
-    assert(object == rhs.object && "incompatible iterators");
-    return index - rhs.index;
-  }
-  bool operator==(const IndexedAccessorIterator &rhs) const {
-    return object == rhs.object && index == rhs.index;
-  }
-  bool operator<(const IndexedAccessorIterator &rhs) const {
-    assert(object == rhs.object && "incompatible iterators");
-    return index < rhs.index;
-  }
-
-  ConcreteType &operator+=(ptrdiff_t offset) {
-    this->index += offset;
-    return static_cast<ConcreteType &>(*this);
-  }
-  ConcreteType &operator-=(ptrdiff_t offset) {
-    this->index -= offset;
-    return static_cast<ConcreteType &>(*this);
-  }
-
-protected:
-  IndexedAccessorIterator(ObjectType *object, unsigned index)
-      : object(object), index(index) {}
-  ObjectType *object;
-  unsigned index;
-};
-
 } // namespace mlir
 
 #endif

mlir/include/mlir/Support/STLExtras.h

@@ -24,6 +24,7 @@
 #define MLIR_SUPPORT_STLEXTRAS_H
 
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/iterator.h"
 #include <tuple>
 
 namespace mlir {
@@ -105,6 +106,48 @@ struct detector<void_t<Op<Args...>>, Op, Args...> {
 
 template <template <class...> class Op, class... Args>
 using is_detected = typename detail::detector<void, Op, Args...>::value_t;
+
+//===----------------------------------------------------------------------===//
+//     Extra additions to <iterator>
+//===----------------------------------------------------------------------===//
+
+/// A utility class used to implement an iterator that contains some object and
+/// an index. The iterator moves the index but keeps the object constant.
+template <typename DerivedT, typename ObjectType, typename T,
+          typename PointerT = T *, typename ReferenceT = T &>
+class indexed_accessor_iterator
+    : public llvm::iterator_facade_base<DerivedT,
+                                        std::random_access_iterator_tag, T,
+                                        std::ptrdiff_t, PointerT, ReferenceT> {
+public:
+  ptrdiff_t operator-(const indexed_accessor_iterator &rhs) const {
+    assert(object == rhs.object && "incompatible iterators");
+    return index - rhs.index;
+  }
+  bool operator==(const indexed_accessor_iterator &rhs) const {
+    return object == rhs.object && index == rhs.index;
+  }
+  bool operator<(const indexed_accessor_iterator &rhs) const {
+    assert(object == rhs.object && "incompatible iterators");
+    return index < rhs.index;
+  }
+
+  DerivedT &operator+=(ptrdiff_t offset) {
+    this->index += offset;
+    return static_cast<DerivedT &>(*this);
+  }
+  DerivedT &operator-=(ptrdiff_t offset) {
+    this->index -= offset;
+    return static_cast<DerivedT &>(*this);
+  }
+
+protected:
+  indexed_accessor_iterator(ObjectType object, ptrdiff_t index)
+      : object(object), index(index) {}
+  ObjectType object;
+  ptrdiff_t index;
+};
+
 } // end namespace mlir
 
 // Allow tuples to be usable as DenseMap keys.

mlir/lib/IR/Attributes.cpp

@@ -464,12 +464,13 @@ static bool hasSameElementsOrSplat(ShapedType type, const Values &values) {
 /// Constructs a new iterator.
 DenseElementsAttr::IntElementIterator::IntElementIterator(
     DenseElementsAttr attr, size_t index)
-    : rawData(attr.getRawData().data()), index(index),
+    : indexed_accessor_iterator<IntElementIterator, const char *, APInt, APInt,
+                                APInt>(attr.getRawData().data(), index),
       bitWidth(getDenseElementBitwidth(attr.getType().getElementType())) {}
 
 /// Accesses the raw APInt value at this iterator position.
 APInt DenseElementsAttr::IntElementIterator::operator*() const {
-  return readBits(rawData, index * getDenseElementStorageWidth(bitWidth),
+  return readBits(object, index * getDenseElementStorageWidth(bitWidth),
                   bitWidth);
 }