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

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.

Differential Revision: https://reviews.llvm.org/D74646

Author: River707

mlir/include/mlir/IR/OpImplementation.h

@@ -188,10 +188,9 @@ 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;
 }

mlir/include/mlir/IR/Operation.h

@@ -232,7 +232,7 @@ class Operation final
 
   // 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(); }
@@ -260,7 +260,7 @@ class Operation final
 
   /// 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();

mlir/include/mlir/IR/OperationSupport.h

@@ -32,14 +32,17 @@ 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
@@ -535,6 +538,46 @@ class OpPrintingFlags {
 // 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
 
@@ -555,6 +598,18 @@ class ValueTypeIterator final
       : 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
 
@@ -568,7 +623,8 @@ class OperandRange final
 
   /// 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.
@@ -598,7 +654,8 @@ class ResultRange final
 
   /// 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.
@@ -666,7 +723,8 @@ class ValueRange final
 
   /// 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;

mlir/lib/IR/OperationSupport.cpp

@@ -140,6 +140,43 @@ void detail::OperandStorage::grow(ResizableStorage &resizeUtil,
 // 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