From 3ca33a5c62e7fbc5562a751880c0b260da4cb1c1 Mon Sep 17 00:00:00 2001
From: River Riddle <riverriddle@google.com>
Date: Tue, 25 Jun 2019 12:10:46 -0700
Subject: [PATCH] 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     | 26 ++--------------
 mlir/include/mlir/IR/Block.h          |  7 +++--
 mlir/include/mlir/IR/Operation.h      | 14 +++++----
 mlir/include/mlir/IR/Value.h          | 37 -----------------------
 mlir/include/mlir/Support/STLExtras.h | 43 +++++++++++++++++++++++++++
 mlir/lib/IR/Attributes.cpp            |  5 ++--
 6 files changed, 60 insertions(+), 72 deletions(-)

diff --git a/mlir/include/mlir/IR/Attributes.h b/mlir/include/mlir/IR/Attributes.h
index 95d438b18e72..0cffdc4d9315 100644
--- a/mlir/include/mlir/IR/Attributes.h
+++ b/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;
   };
diff --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index 3dea64d6d79b..f4ecb4ec6d73 100644
--- a/mlir/include/mlir/IR/Block.h
+++ b/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) {}
diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index f39f127338bf..9b310aa0c496 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/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); }
 };
diff --git a/mlir/include/mlir/IR/Value.h b/mlir/include/mlir/IR/Value.h
index a81011da1f68..e90505ec90d4 100644
--- a/mlir/include/mlir/IR/Value.h
+++ b/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
diff --git a/mlir/include/mlir/Support/STLExtras.h b/mlir/include/mlir/Support/STLExtras.h
index c9a5198ad306..4f19a90c4520 100644
--- a/mlir/include/mlir/Support/STLExtras.h
+++ b/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.
diff --git a/mlir/lib/IR/Attributes.cpp b/mlir/lib/IR/Attributes.cpp
index c35e1fdab88e..8c9141bb871c 100644
--- a/mlir/lib/IR/Attributes.cpp
+++ b/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);
 }