Refactor the various operand/result/type iterators to use indexed_accessor_range.

This has several benefits:
* The implementation is much cleaner and more efficient.
* The ranges now have support for many useful operations: operator[], slice, drop_front, size, etc.
* Value ranges can now directly query a range for their types via 'getTypes()': e.g:
   void foo(Operation::operand_range operands) {
     auto operandTypes = operands.getTypes();
   }

PiperOrigin-RevId: 284834912

3 files changed, 64 insertions, 54 deletions

diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index 0483c27e968..fd747a98a40 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -597,9 +597,8 @@ void Operation::setSuccessor(Block *block, unsigned index) {
 }
 
 auto Operation::getNonSuccessorOperands() -> operand_range {
-  return {operand_iterator(this, 0),
-          operand_iterator(this, hasSuccessors() ? getSuccessorOperandIndex(0)
-                                                 : getNumOperands())};
+  return getOperands().take_front(hasSuccessors() ? getSuccessorOperandIndex(0)
+                                                  : getNumOperands());
 }
 
 /// Get the index of the first operand of the successor at the provided
@@ -635,9 +634,7 @@ Operation::decomposeSuccessorOperandIndex(unsigned operandIndex) {
 
 auto Operation::getSuccessorOperands(unsigned index) -> operand_range {
   unsigned succOperandIndex = getSuccessorOperandIndex(index);
-  return {operand_iterator(this, succOperandIndex),
-          operand_iterator(this,
-                           succOperandIndex + getNumSuccessorOperands(index))};
+  return getOperands().slice(succOperandIndex, getNumSuccessorOperands(index));
 }
 
 /// Attempt to fold this operation using the Op's registered foldHook.
@@ -746,48 +743,6 @@ Operation *Operation::clone() {
 }
 
 //===----------------------------------------------------------------------===//
-// ValueRange
-//===----------------------------------------------------------------------===//
-
-ValueRange::ValueRange(ArrayRef<Value *> values)
-    : ValueRange(values.data(), values.size()) {}
-ValueRange::ValueRange(llvm::iterator_range<OperandIterator> values)
-    : ValueRange(nullptr, llvm::size(values)) {
-  if (!empty()) {
-    auto begin = values.begin();
-    base = &begin.getBase()->getOpOperand(begin.getIndex());
-  }
-}
-ValueRange::ValueRange(llvm::iterator_range<ResultIterator> values)
-    : ValueRange(nullptr, llvm::size(values)) {
-  if (!empty()) {
-    auto begin = values.begin();
-    base = &begin.getBase()->getOpResult(begin.getIndex());
-  }
-}
-
-/// 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 *>())
-    return operand + index;
-  if (OpResult *result = owner.dyn_cast<OpResult *>())
-    return result + index;
-  return owner.get<Value *const *>() + index;
-}
-/// 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 = owner.dyn_cast<OpOperand *>())
-    return operand[index].get();
-  // An OpResult is a value, so we can return it directly.
-  if (OpResult *result = owner.dyn_cast<OpResult *>())
-    return &result[index];
-  // Otherwise, this is a raw value array so just index directly.
-  return owner.get<Value *const *>()[index];
-}
-
-//===----------------------------------------------------------------------===//
 // OpState trait class.
 //===----------------------------------------------------------------------===//
 
@@ -979,7 +934,7 @@ OpTrait::impl::verifySameOperandsAndResultElementType(Operation *op) {
   auto elementType = getElementTypeOrSelf(op->getResult(0));
 
   // Verify result element type matches first result's element type.
-  for (auto result : drop_begin(op->getResults(), 1)) {
+  for (auto result : llvm::drop_begin(op->getResults(), 1)) {
     if (getElementTypeOrSelf(result) != elementType)
       return op->emitOpError(
           "requires the same element type for all operands and results");
@@ -1210,7 +1165,7 @@ Value *impl::foldCastOp(Operation *op) {
 }
 
 //===----------------------------------------------------------------------===//
-// CastOp implementation
+// Misc. utils
 //===----------------------------------------------------------------------===//
 
 /// Insert an operation, generated by `buildTerminatorOp`, at the end of the
@@ -1230,6 +1185,10 @@ void impl::ensureRegionTerminator(
   block.push_back(buildTerminatorOp());
 }
 
+//===----------------------------------------------------------------------===//
+// UseIterator
+//===----------------------------------------------------------------------===//
+
 UseIterator::UseIterator(Operation *op, bool end)
     : op(op), res(end ? op->result_end() : op->result_begin()) {
   // Only initialize current use if there are results/can be uses.
diff --git a/mlir/lib/IR/OperationSupport.cpp b/mlir/lib/IR/OperationSupport.cpp
index e4ff889e4dc..256a261acd8 100644
--- a/mlir/lib/IR/OperationSupport.cpp
+++ b/mlir/lib/IR/OperationSupport.cpp
@@ -144,3 +144,50 @@ void detail::OperandStorage::grow(ResizableStorage &resizeUtil,
     operand.~OpOperand();
   resizeUtil.setDynamicStorage(newStorage);
 }
+
+//===----------------------------------------------------------------------===//
+// Operation Value-Iterators
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// OperandRange
+
+OperandRange::OperandRange(Operation *op)
+    : OperandRange(op->getOpOperands().data(), op->getNumOperands()) {}
+
+//===----------------------------------------------------------------------===//
+// ResultRange
+
+ResultRange::ResultRange(Operation *op)
+    : ResultRange(op->getOpResults().data(), op->getNumResults()) {}
+
+//===----------------------------------------------------------------------===//
+// ValueRange
+
+ValueRange::ValueRange(ArrayRef<Value *> values)
+    : ValueRange(values.data(), values.size()) {}
+ValueRange::ValueRange(OperandRange values)
+    : ValueRange(values.begin().getBase(), values.size()) {}
+ValueRange::ValueRange(ResultRange values)
+    : ValueRange(values.begin().getBase(), values.size()) {}
+
+/// 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 *>())
+    return operand + index;
+  if (OpResult *result = owner.dyn_cast<OpResult *>())
+    return result + index;
+  return owner.get<Value *const *>() + index;
+}
+/// 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 = owner.dyn_cast<OpOperand *>())
+    return operand[index].get();
+  // An OpResult is a value, so we can return it directly.
+  if (OpResult *result = owner.dyn_cast<OpResult *>())
+    return &result[index];
+  // Otherwise, this is a raw value array so just index directly.
+  return owner.get<Value *const *>()[index];
+}
diff --git a/mlir/lib/IR/TypeUtilities.cpp b/mlir/lib/IR/TypeUtilities.cpp
index a963a8dd459..01721418340 100644
--- a/mlir/lib/IR/TypeUtilities.cpp
+++ b/mlir/lib/IR/TypeUtilities.cpp
@@ -92,15 +92,19 @@ LogicalResult mlir::verifyCompatibleShape(Type type1, Type type2) {
   return success();
 }
 
-OperandElementTypeIterator::OperandElementTypeIterator(OperandIterator it)
-    : llvm::mapped_iterator<OperandIterator, Type (*)(Value *)>(it, &unwrap) {}
+OperandElementTypeIterator::OperandElementTypeIterator(
+    Operation::operand_iterator it)
+    : llvm::mapped_iterator<Operation::operand_iterator, Type (*)(Value *)>(
+          it, &unwrap) {}
 
 Type OperandElementTypeIterator::unwrap(Value *value) {
   return value->getType().cast<ShapedType>().getElementType();
 }
 
-ResultElementTypeIterator::ResultElementTypeIterator(ResultIterator it)
-    : llvm::mapped_iterator<ResultIterator, Type (*)(Value *)>(it, &unwrap) {}
+ResultElementTypeIterator::ResultElementTypeIterator(
+    Operation::result_iterator it)
+    : llvm::mapped_iterator<Operation::result_iterator, Type (*)(Value *)>(
+          it, &unwrap) {}
 
 Type ResultElementTypeIterator::unwrap(Value *value) {
   return value->getType().cast<ShapedType>().getElementType();