Add a new ValueRange class.

This class represents a generic abstraction over the different ways to represent a range of Values: ArrayRef<Value *>, operand_range, result_range. This class will allow for removing the many instances of explicit SmallVector<Value *, N> construction. It has the same memory cost as ArrayRef, and only suffers cost from indexing(if+elsing the different underlying representations).

This change only updates a few of the existing usages, with more to be changed in followups; e.g. 'build' API.

PiperOrigin-RevId: 284307996

3 files changed, 46 insertions, 14 deletions

diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index ed97b8b5940..0ea447ed324 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -1486,10 +1486,10 @@ public:
   /// SSA values in namesToUse.  This may only be used for IsolatedFromAbove
   /// operations.  If any entry in namesToUse is null, the corresponding
   /// argument name is left alone.
-  void shadowRegionArgs(Region &region, ArrayRef<Value *> namesToUse) override;
+  void shadowRegionArgs(Region &region, ValueRange namesToUse) override;
 
   void printAffineMapOfSSAIds(AffineMapAttr mapAttr,
-                              ArrayRef<Value *> operands) override {
+                              ValueRange operands) override {
     AffineMap map = mapAttr.getValue();
     unsigned numDims = map.getNumDims();
     auto printValueName = [&](unsigned pos, bool isSymbol) {
@@ -1851,8 +1851,7 @@ void OperationPrinter::printValueIDImpl(Value *value, bool printResultNo,
 /// SSA values in namesToUse.  This may only be used for IsolatedFromAbove
 /// operations.  If any entry in namesToUse is null, the corresponding
 /// argument name is left alone.
-void OperationPrinter::shadowRegionArgs(Region &region,
-                                        ArrayRef<Value *> namesToUse) {
+void OperationPrinter::shadowRegionArgs(Region &region, ValueRange namesToUse) {
   assert(!region.empty() && "cannot shadow arguments of an empty region");
   assert(region.front().getNumArguments() == namesToUse.size() &&
          "incorrect number of names passed in");
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index 1d213f45dd5..3ebc7129305 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -740,6 +740,42 @@ Operation *Operation::clone() {
 }
 
 //===----------------------------------------------------------------------===//
+// ValueRange
+//===----------------------------------------------------------------------===//
+
+ValueRange::ValueRange(ArrayRef<Value *> values)
+    : owner(values.data()), count(values.size()) {}
+ValueRange::ValueRange(llvm::iterator_range<OperandIterator> values)
+    : count(llvm::size(values)) {
+  if (count != 0) {
+    auto begin = values.begin();
+    owner = &begin.getObject()->getOpOperand(begin.getIndex());
+  }
+}
+ValueRange::ValueRange(llvm::iterator_range<ResultIterator> values)
+    : count(llvm::size(values)) {
+  if (count != 0) {
+    auto begin = values.begin();
+    owner = &begin.getObject()->getOpResult(begin.getIndex());
+  }
+}
+
+ValueRange::Iterator::Iterator(OwnerT owner, unsigned curIndex)
+    : indexed_accessor_iterator<Iterator, OwnerT, Value *, Value *, Value *>(
+          owner, curIndex) {}
+
+Value *ValueRange::Iterator::operator*() const {
+  // Operands access the held value via 'get'.
+  if (OpOperand *operand = object.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 *>())
+    return &result[index];
+  // Otherwise, this is a raw value array so just index directly.
+  return object.get<Value *const *>()[index];
+}
+
+//===----------------------------------------------------------------------===//
 // OpState trait class.
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index b8ecab97552..3887a0308b0 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -85,8 +85,8 @@ PatternRewriter::~PatternRewriter() {
 /// clients can specify a list of other nodes that this replacement may make
 /// (perhaps transitively) dead.  If any of those ops are dead, this will
 /// remove them as well.
-void PatternRewriter::replaceOp(Operation *op, ArrayRef<Value *> newValues,
-                                ArrayRef<Value *> valuesToRemoveIfDead) {
+void PatternRewriter::replaceOp(Operation *op, ValueRange newValues,
+                                ValueRange valuesToRemoveIfDead) {
   // Notify the rewriter subclass that we're about to replace this root.
   notifyRootReplaced(op);
 
@@ -114,7 +114,7 @@ void PatternRewriter::eraseOp(Operation *op) {
 /// 'argValues' is used to replace the block arguments of 'source' after
 /// merging.
 void PatternRewriter::mergeBlocks(Block *source, Block *dest,
-                                  ArrayRef<Value *> argValues) {
+                                  ValueRange argValues) {
   assert(llvm::all_of(source->getPredecessors(),
                       [dest](Block *succ) { return succ == dest; }) &&
          "expected 'source' to have no predecessors or only 'dest'");
@@ -141,15 +141,12 @@ Block *PatternRewriter::splitBlock(Block *block, Block::iterator before) {
 /// op and newOp are known to have the same number of results, replace the
 /// uses of op with uses of newOp
 void PatternRewriter::replaceOpWithResultsOfAnotherOp(
-    Operation *op, Operation *newOp, ArrayRef<Value *> valuesToRemoveIfDead) {
+    Operation *op, Operation *newOp, ValueRange valuesToRemoveIfDead) {
   assert(op->getNumResults() == newOp->getNumResults() &&
          "replacement op doesn't match results of original op");
   if (op->getNumResults() == 1)
     return replaceOp(op, newOp->getResult(0), valuesToRemoveIfDead);
-
-  SmallVector<Value *, 8> newResults(newOp->getResults().begin(),
-                                     newOp->getResults().end());
-  return replaceOp(op, newResults, valuesToRemoveIfDead);
+  return replaceOp(op, newOp->getResults(), valuesToRemoveIfDead);
 }
 
 /// Move the blocks that belong to "region" before the given position in
@@ -190,8 +187,8 @@ void PatternRewriter::cloneRegionBefore(Region &region, Block *before) {
 /// The opsToRemoveIfDead list is an optional list of nodes that the rewriter
 /// should remove if they are dead at this point.
 ///
-void PatternRewriter::updatedRootInPlace(
-    Operation *op, ArrayRef<Value *> valuesToRemoveIfDead) {
+void PatternRewriter::updatedRootInPlace(Operation *op,
+                                         ValueRange valuesToRemoveIfDead) {
   // Notify the rewriter subclass that we're about to replace this root.
   notifyRootUpdated(op);