Add new indexed_accessor_range_base and indexed_accessor_range classes that simplify defining index-able ranges.

Many ranges want similar functionality from a range type(e.g. slice/drop_front/operator[]/etc.), so these classes provide a generic implementation that may be used by many different types of ranges. This removes some code duplication, and also empowers many of the existing range types in MLIR(e.g. result type ranges, operand ranges, ElementsAttr ranges, etc.). This change only updates RegionRange and ValueRange, more ranges will be updated in followup commits.

PiperOrigin-RevId: 284615679

3 files changed, 36 insertions, 49 deletions

diff --git a/mlir/lib/IR/Attributes.cpp b/mlir/lib/IR/Attributes.cpp
index f2f3d41f980..b546643837b 100644
--- a/mlir/lib/IR/Attributes.cpp
+++ b/mlir/lib/IR/Attributes.cpp
@@ -527,7 +527,7 @@ DenseElementsAttr::AttributeElementIterator::AttributeElementIterator(
 
 /// Accesses the Attribute value at this iterator position.
 Attribute DenseElementsAttr::AttributeElementIterator::operator*() const {
-  auto owner = getFromOpaquePointer(object).cast<DenseElementsAttr>();
+  auto owner = getFromOpaquePointer(base).cast<DenseElementsAttr>();
   Type eltTy = owner.getType().getElementType();
   if (auto intEltTy = eltTy.dyn_cast<IntegerType>()) {
     if (intEltTy.getWidth() == 1)
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index ae635d108b2..0483c27e968 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -750,60 +750,41 @@ Operation *Operation::clone() {
 //===----------------------------------------------------------------------===//
 
 ValueRange::ValueRange(ArrayRef<Value *> values)
-    : owner(values.data()), count(values.size()) {}
+    : ValueRange(values.data(), values.size()) {}
 ValueRange::ValueRange(llvm::iterator_range<OperandIterator> values)
-    : count(llvm::size(values)) {
-  if (count != 0) {
+    : ValueRange(nullptr, llvm::size(values)) {
+  if (!empty()) {
     auto begin = values.begin();
-    owner = &begin.getObject()->getOpOperand(begin.getIndex());
+    base = &begin.getBase()->getOpOperand(begin.getIndex());
   }
 }
 ValueRange::ValueRange(llvm::iterator_range<ResultIterator> values)
-    : count(llvm::size(values)) {
-  if (count != 0) {
+    : ValueRange(nullptr, llvm::size(values)) {
+  if (!empty()) {
     auto begin = values.begin();
-    owner = &begin.getObject()->getOpResult(begin.getIndex());
+    base = &begin.getBase()->getOpResult(begin.getIndex());
   }
 }
 
-/// Drop the first N elements, and keep M elements.
-ValueRange ValueRange::slice(unsigned n, unsigned m) const {
-  assert(n + m <= size() && "Invalid specifier");
-  OwnerT newOwner;
+/// See `detail::indexed_accessor_range_base` for details.
+ValueRange::OwnerT ValueRange::offset_base(const OwnerT &owner,
+                                           ptrdiff_t index) {
   if (OpOperand *operand = owner.dyn_cast<OpOperand *>())
-    newOwner = operand + n;
-  else if (OpResult *result = owner.dyn_cast<OpResult *>())
-    newOwner = result + n;
-  else
-    newOwner = owner.get<Value *const *>() + n;
-  return ValueRange(newOwner, m);
-}
-
-/// Drop the first n elements.
-ValueRange ValueRange::drop_front(unsigned n) const {
-  assert(size() >= n && "Dropping more elements than exist");
-  return slice(n, size() - n);
-}
-
-/// Drop the last n elements.
-ValueRange ValueRange::drop_back(unsigned n) const {
-  assert(size() >= n && "Dropping more elements than exist");
-  return ValueRange(owner, size() - n);
+    return operand + index;
+  if (OpResult *result = owner.dyn_cast<OpResult *>())
+    return result + index;
+  return owner.get<Value *const *>() + index;
 }
-
-ValueRange::Iterator::Iterator(OwnerT owner, unsigned curIndex)
-    : indexed_accessor_iterator<Iterator, OwnerT, Value *, Value *, Value *>(
-          owner, curIndex) {}
-
-Value *ValueRange::Iterator::operator*() const {
+/// See `detail::indexed_accessor_range_base` for details.
+Value *ValueRange::dereference_iterator(const OwnerT &owner, ptrdiff_t index) {
   // Operands access the held value via 'get'.
-  if (OpOperand *operand = object.dyn_cast<OpOperand *>())
+  if (OpOperand *operand = owner.dyn_cast<OpOperand *>())
     return operand[index].get();
   // An OpResult is a value, so we can return it directly.
-  if (OpResult *result = object.dyn_cast<OpResult *>())
+  if (OpResult *result = owner.dyn_cast<OpResult *>())
     return &result[index];
   // Otherwise, this is a raw value array so just index directly.
-  return object.get<Value *const *>()[index];
+  return owner.get<Value *const *>()[index];
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/Region.cpp b/mlir/lib/IR/Region.cpp
index a5a19cbcbf2..c588e567bc3 100644
--- a/mlir/lib/IR/Region.cpp
+++ b/mlir/lib/IR/Region.cpp
@@ -217,17 +217,23 @@ void llvm::ilist_traits<::mlir::Block>::transferNodesFromList(
 //===----------------------------------------------------------------------===//
 // RegionRange
 //===----------------------------------------------------------------------===//
+
 RegionRange::RegionRange(MutableArrayRef<Region> regions)
-    : owner(regions.data()), count(regions.size()) {}
+    : RegionRange(regions.data(), regions.size()) {}
 RegionRange::RegionRange(ArrayRef<std::unique_ptr<Region>> regions)
-    : owner(regions.data()), count(regions.size()) {}
-RegionRange::Iterator::Iterator(OwnerT owner, unsigned curIndex)
-    : indexed_accessor_iterator<Iterator, OwnerT, Region *, Region *, Region *>(
-          owner, curIndex) {}
-
-Region *RegionRange::Iterator::operator*() const {
-  if (const std::unique_ptr<Region> *operand =
-          object.dyn_cast<const std::unique_ptr<Region> *>())
+    : RegionRange(regions.data(), regions.size()) {}
+
+/// See `detail::indexed_accessor_range_base` for details.
+RegionRange::OwnerT RegionRange::offset_base(const OwnerT &owner,
+                                             ptrdiff_t index) {
+  if (auto *operand = owner.dyn_cast<const std::unique_ptr<Region> *>())
+    return operand + index;
+  return &owner.get<Region *>()[index];
+}
+/// See `detail::indexed_accessor_range_base` for details.
+Region *RegionRange::dereference_iterator(const OwnerT &owner,
+                                          ptrdiff_t index) {
+  if (auto *operand = owner.dyn_cast<const std::unique_ptr<Region> *>())
     return operand[index].get();
-  return &object.get<Region *>()[index];
+  return &owner.get<Region *>()[index];
 }