[mlir] Add a TypeRange class that functions similar to ValueRange.

[rriddle]

Summary

This class wraps around the various different ways to construct a range of Type, without forcing the materialization of that range into a contiguous vector.

[benvanik]

Perfect!

[rriddle]

Rebase.

mlir/include/mlir/IR/OpImplementation.h

inline OpAsmPrinter &operator<<(OpAsmPrinter &p, const T &other) {
   p.getStream() << other;
   return p;
 }
 
 inline OpAsmPrinter &operator<<(OpAsmPrinter &p, bool value) {
   return p << (value ? StringRef("true") : "false");
 }
 
-template <typename IteratorT>
-inline OpAsmPrinter &
-operator<<(OpAsmPrinter &p,
-           const iterator_range<ValueTypeIterator<IteratorT>> &types) {
+template <typename ValueRangeT>
+inline OpAsmPrinter &operator<<(OpAsmPrinter &p,
+                                const ValueTypeRange<ValueRangeT> &types) {
   interleaveComma(types, p);
   return p;
 }
 inline OpAsmPrinter &operator<<(OpAsmPrinter &p, ArrayRef<Type> types) {
   interleaveComma(types, p);
   return p;
 }
 

mlir/include/mlir/IR/Operation.h

public:
   MutableArrayRef<OpOperand> getOpOperands() {
     return getOperandStorage().getOperands();
   }
 
   OpOperand &getOpOperand(unsigned idx) { return getOpOperands()[idx]; }
 
   // Support operand type iteration.
   using operand_type_iterator = operand_range::type_iterator;
-  using operand_type_range = iterator_range<operand_type_iterator>;
+  using operand_type_range = operand_range::type_range;
   operand_type_iterator operand_type_begin() { return operand_begin(); }
   operand_type_iterator operand_type_end() { return operand_end(); }
   operand_type_range getOperandTypes() { return getOperands().getTypes(); }
 
   //===--------------------------------------------------------------------===//
   // Results
   //===--------------------------------------------------------------------===//
 
public:
   result_iterator result_end() { return getResults().end(); }
   result_range getResults() { return result_range(this); }
 
   result_range getOpResults() { return getResults(); }
   OpResult getOpResult(unsigned idx) { return getResult(idx); }
 
   /// Support result type iteration.
   using result_type_iterator = result_range::type_iterator;
-  using result_type_range = ArrayRef<Type>;
+  using result_type_range = result_range::type_range;
   result_type_iterator result_type_begin() { return getResultTypes().begin(); }
   result_type_iterator result_type_end() { return getResultTypes().end(); }
   result_type_range getResultTypes();
 
   //===--------------------------------------------------------------------===//
   // Attributes
   //===--------------------------------------------------------------------===//
 

mlir/include/mlir/IR/OperationSupport.h

 namespace mlir {
 class Block;
 class Dialect;
 class Operation;
 struct OperationState;
 class OpAsmParser;
 class OpAsmParserResult;
 class OpAsmPrinter;
+class OperandRange;
 class OpFoldResult;
 class ParseResult;
 class Pattern;
 class Region;
+class ResultRange;
 class RewritePattern;
 class Type;
 class Value;
 class ValueRange;
+template <typename ValueRangeT> class ValueTypeRange;
 
 /// This is an adaptor from a list of values to named operands of OpTy.  In a
 /// generic operation context, e.g., in dialect conversions, an ordered array of
 /// `Value`s is treated as operands of `OpTy`.  This adaptor takes a reference
 /// to the array and provides accessors with the same names as `OpTy` for
 /// operands.  This makes possible to create function templates that operate on
 /// either OpTy or OperandAdaptor<OpTy> seamlessly.
 template <typename OpTy> using OperandAdaptor = typename OpTy::OperandAdaptor;
private:
   bool printLocalScope : 1;
 };
 
 //===----------------------------------------------------------------------===//
 // Operation Value-Iterators
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
+// TypeRange
+
+/// This class provides an abstraction over the various different ranges of
+/// value types. In many cases, this prevents the need to explicitly materialize
+/// a SmallVector/std::vector. This class should be used in places that are not
+/// suitable for a more derived type (e.g. ArrayRef) or a template range
+/// parameter.
+class TypeRange
+    : public detail::indexed_accessor_range_base<
+          TypeRange,
+          llvm::PointerUnion<const Value *, const Type *, OpOperand *>, Type,
+          Type, Type> {
+public:
+  using RangeBaseT::RangeBaseT;
+  TypeRange(ArrayRef<Type> types = llvm::None);
+  explicit TypeRange(OperandRange values);
+  explicit TypeRange(ResultRange values);
+  explicit TypeRange(ValueRange values);
+  template <typename ValueRangeT>
+  TypeRange(ValueTypeRange<ValueRangeT> values)
+      : TypeRange(ValueRangeT(values.begin().getCurrent(),
+                              values.end().getCurrent())) {}
+
+private:
+  /// The owner of the range is either:
+  /// * A pointer to the first element of an array of values.
+  /// * A pointer to the first element of an array of types.
+  /// * A pointer to the first element of an array of operands.
+  using OwnerT = llvm::PointerUnion<const Value *, const Type *, OpOperand *>;
+
+  /// See `detail::indexed_accessor_range_base` for details.
+  static OwnerT offset_base(OwnerT object, ptrdiff_t index);
+  /// See `detail::indexed_accessor_range_base` for details.
+  static Type dereference_iterator(OwnerT object, ptrdiff_t index);
+
+  /// Allow access to `offset_base` and `dereference_iterator`.
+  friend RangeBaseT;
+};
+
+//===----------------------------------------------------------------------===//
 // ValueTypeRange
 
 /// This class implements iteration on the types of a given range of values.
 template <typename ValueIteratorT>
 class ValueTypeIterator final
     : public llvm::mapped_iterator<ValueIteratorT, Type (*)(Value)> {
   static Type unwrap(Value value) { return value.getType(); }
 
 public:
   using reference = Type;
 
   /// Provide a const dereference method.
   Type operator*() const { return unwrap(*this->I); }
 
   /// Initializes the type iterator to the specified value iterator.
   ValueTypeIterator(ValueIteratorT it)
       : llvm::mapped_iterator<ValueIteratorT, Type (*)(Value)>(it, &unwrap) {}
 };
 
+/// This class implements iteration on the types of a given range of values.
+template <typename ValueRangeT>
+class ValueTypeRange final
+    : public llvm::iterator_range<
+          ValueTypeIterator<typename ValueRangeT::iterator>> {
+public:
+  using llvm::iterator_range<
+      ValueTypeIterator<typename ValueRangeT::iterator>>::iterator_range;
+  template <typename Container>
+  ValueTypeRange(Container &&c) : ValueTypeRange(c.begin(), c.end()) {}
+};
+
 //===----------------------------------------------------------------------===//
 // OperandRange
 
 /// This class implements the operand iterators for the Operation class.
 class OperandRange final
     : public detail::indexed_accessor_range_base<OperandRange, OpOperand *,
                                                  Value, Value, Value> {
 public:
   using RangeBaseT::RangeBaseT;
   OperandRange(Operation *op);
 
   /// Returns the types of the values within this range.
   using type_iterator = ValueTypeIterator<iterator>;
-  iterator_range<type_iterator> getTypes() const { return {begin(), end()}; }
+  using type_range = ValueTypeRange<OperandRange>;
+  type_range getTypes() const { return {begin(), end()}; }
 
 private:
   /// See `detail::indexed_accessor_range_base` for details.
   static OpOperand *offset_base(OpOperand *object, ptrdiff_t index) {
     return object + index;
   }
   /// See `detail::indexed_accessor_range_base` for details.
   static Value dereference_iterator(OpOperand *object, ptrdiff_t index) {
: public indexed_accessor_range<ResultRange, Operation *, OpResult,
                                     OpResult, OpResult> {
 public:
   using indexed_accessor_range<ResultRange, Operation *, OpResult, OpResult,
                                OpResult>::indexed_accessor_range;
   ResultRange(Operation *op);
 
   /// Returns the types of the values within this range.
   using type_iterator = ArrayRef<Type>::iterator;
-  ArrayRef<Type> getTypes() const;
+  using type_range = ArrayRef<Type>;
+  type_range getTypes() const;
 
 private:
   /// See `indexed_accessor_range` for details.
   static OpResult dereference(Operation *op, ptrdiff_t index);
 
   /// Allow access to `dereference_iterator`.
   friend indexed_accessor_range<ResultRange, Operation *, OpResult, OpResult,
                                 OpResult>;
public:
   ValueRange(ArrayRef<BlockArgument> values)
       : ValueRange(ArrayRef<Value>(values.data(), values.size())) {}
   ValueRange(ArrayRef<Value> values = llvm::None);
   ValueRange(OperandRange values);
   ValueRange(ResultRange values);
 
   /// Returns the types of the values within this range.
   using type_iterator = ValueTypeIterator<iterator>;
-  iterator_range<type_iterator> getTypes() const { return {begin(), end()}; }
+  using type_range = ValueTypeRange<ValueRange>;
+  type_range getTypes() const { return {begin(), end()}; }
 
 private:
   using OwnerT = detail::ValueRangeOwner;
 
   /// See `detail::indexed_accessor_range_base` for details.
   static OwnerT offset_base(const OwnerT &owner, ptrdiff_t index);
   /// See `detail::indexed_accessor_range_base` for details.
   static Value dereference_iterator(const OwnerT &owner, ptrdiff_t index);

mlir/lib/IR/OperationSupport.cpp

void detail::OperandStorage::grow(ResizableStorage &resizeUtil,
   resizeUtil.setDynamicStorage(newStorage);
 }
 
 //===----------------------------------------------------------------------===//
 // Operation Value-Iterators
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
+// TypeRange
+
+TypeRange::TypeRange(ArrayRef<Type> types)
+    : TypeRange(types.data(), types.size()) {}
+TypeRange::TypeRange(OperandRange values)
+    : TypeRange(values.begin().getBase(), values.size()) {}
+TypeRange::TypeRange(ResultRange values)
+    : TypeRange(values.getBase()->getResultTypes().slice(values.getStartIndex(),
+                                                         values.size())) {}
+TypeRange::TypeRange(ValueRange values) : TypeRange(OwnerT(), values.size()) {
+  detail::ValueRangeOwner owner = values.begin().getBase();
+  if (auto *op = reinterpret_cast<Operation *>(owner.ptr.dyn_cast<void *>()))
+    this->base = &op->getResultTypes()[owner.startIndex];
+  else if (auto *operand = owner.ptr.dyn_cast<OpOperand *>())
+    this->base = operand;
+  else
+    this->base = owner.ptr.get<const Value *>();
+}
+
+/// See `detail::indexed_accessor_range_base` for details.
+TypeRange::OwnerT TypeRange::offset_base(OwnerT object, ptrdiff_t index) {
+  if (auto *value = object.dyn_cast<const Value *>())
+    return {value + index};
+  if (auto *operand = object.dyn_cast<OpOperand *>())
+    return {operand + index};
+  return {object.dyn_cast<const Type *>() + index};
+}
+/// See `detail::indexed_accessor_range_base` for details.
+Type TypeRange::dereference_iterator(OwnerT object, ptrdiff_t index) {
+  if (auto *value = object.dyn_cast<const Value *>())
+    return (value + index)->getType();
+  if (auto *operand = object.dyn_cast<OpOperand *>())
+    return (operand + index)->get().getType();
+  return object.dyn_cast<const Type *>()[index];
+}
+
+//===----------------------------------------------------------------------===//
 // OperandRange
 
 OperandRange::OperandRange(Operation *op)
     : OperandRange(op->getOpOperands().data(), op->getNumOperands()) {}
 
 //===----------------------------------------------------------------------===//
 // ResultRange