[MLIR] Change Operation::create() methods to use Value/Type/Block ranges.

- Introduce a new BlockRange class to represent range of blocks (constructible from
  an ArrayRef<Block *> or a SuccessorRange);
- Change Operation::create() methods to use TypeRange for result types, ValueRange for
  operands and BlockRange for successors.

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

mlir/include/mlir/IR/BlockSupport.h

@@ -75,6 +75,47 @@ private:
   friend RangeBaseT;
 };
 
+//===----------------------------------------------------------------------===//
+// BlockRange
+//===----------------------------------------------------------------------===//
+
+/// This class provides an abstraction over the different types of ranges over
+/// Blocks. 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 BlockRange final
+    : public llvm::detail::indexed_accessor_range_base<
+          BlockRange, llvm::PointerUnion<BlockOperand *, Block *const *>,
+          Block *, Block *, Block *> {
+public:
+  using RangeBaseT::RangeBaseT;
+  BlockRange(ArrayRef<Block *> blocks = llvm::None);
+  BlockRange(SuccessorRange successors);
+  template <typename Arg,
+            typename = typename std::enable_if_t<
+                std::is_constructible<ArrayRef<Block *>, Arg>::value>>
+  BlockRange(Arg &&arg)
+      : BlockRange(ArrayRef<Block *>(std::forward<Arg>(arg))) {}
+  BlockRange(std::initializer_list<Block *> blocks)
+      : BlockRange(ArrayRef<Block *>(blocks)) {}
+
+private:
+  /// The owner of the range is either:
+  /// * A pointer to the first element of an array of block operands.
+  /// * A pointer to the first element of an array of Block *.
+  using OwnerT = llvm::PointerUnion<BlockOperand *, Block *const *>;
+
+  /// See `llvm::detail::indexed_accessor_range_base` for details.
+  static OwnerT offset_base(OwnerT object, ptrdiff_t index);
+
+  /// See `llvm::detail::indexed_accessor_range_base` for details.
+  static Block *dereference_iterator(OwnerT object, ptrdiff_t index);
+
+  /// Allow access to `offset_base` and `dereference_iterator`.
+  friend RangeBaseT;
+};
+
 //===----------------------------------------------------------------------===//
 // Operation Iterators
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/Operation.h

@@ -32,25 +32,25 @@ class Operation final
 public:
   /// Create a new Operation with the specific fields.
   static Operation *create(Location location, OperationName name,
-                           ArrayRef<Type> resultTypes, ArrayRef<Value> operands,
+                           TypeRange resultTypes, ValueRange operands,
                            ArrayRef<NamedAttribute> attributes,
-                           ArrayRef<Block *> successors, unsigned numRegions);
+                           BlockRange successors, unsigned numRegions);
 
   /// Overload of create that takes an existing MutableDictionaryAttr to avoid
   /// unnecessarily uniquing a list of attributes.
   static Operation *create(Location location, OperationName name,
-                           ArrayRef<Type> resultTypes, ArrayRef<Value> operands,
+                           TypeRange resultTypes, ValueRange operands,
                            MutableDictionaryAttr attributes,
-                           ArrayRef<Block *> successors, unsigned numRegions);
+                           BlockRange successors, unsigned numRegions);
 
   /// Create a new Operation from the fields stored in `state`.
   static Operation *create(const OperationState &state);
 
   /// Create a new Operation with the specific fields.
   static Operation *create(Location location, OperationName name,
-                           ArrayRef<Type> resultTypes, ArrayRef<Value> operands,
+                           TypeRange resultTypes, ValueRange operands,
                            MutableDictionaryAttr attributes,
-                           ArrayRef<Block *> successors = {},
+                           BlockRange successors = {},
                            RegionRange regions = {});
 
   /// The name of an operation is the key identifier for it.
@@ -633,7 +633,7 @@ private:
   bool hasValidOrder() { return orderIndex != kInvalidOrderIdx; }
 
 private:
-  Operation(Location location, OperationName name, ArrayRef<Type> resultTypes,
+  Operation(Location location, OperationName name, TypeRange resultTypes,
             unsigned numSuccessors, unsigned numRegions,
             const MutableDictionaryAttr &attributes, bool hasOperandStorage);
 

mlir/include/mlir/IR/OperationSupport.h

@@ -29,6 +29,7 @@
 
 namespace mlir {
 class Block;
+class BlockRange;
 class Dialect;
 class Operation;
 struct OperationState;
@@ -42,7 +43,6 @@ class Pattern;
 class Region;
 class ResultRange;
 class RewritePattern;
-class SuccessorRange;
 class Type;
 class Value;
 class ValueRange;
@@ -394,12 +394,8 @@ public:
     attributes.append(newAttributes);
   }
 
-  /// Add an array of successors.
-  void addSuccessors(ArrayRef<Block *> newSuccessors) {
-    successors.append(newSuccessors.begin(), newSuccessors.end());
-  }
   void addSuccessors(Block *successor) { successors.push_back(successor); }
-  void addSuccessors(SuccessorRange newSuccessors);
+  void addSuccessors(BlockRange newSuccessors);
 
   /// Create a region that should be attached to the operation.  These regions
   /// can be filled in immediately without waiting for Operation to be

mlir/lib/IR/Block.cpp

@@ -282,7 +282,7 @@ unsigned PredecessorIterator::getSuccessorIndex() const {
 }
 
 //===----------------------------------------------------------------------===//
-// Successors
+// SuccessorRange
 //===----------------------------------------------------------------------===//
 
 SuccessorRange::SuccessorRange(Block *block) : SuccessorRange(nullptr, 0) {
@@ -295,3 +295,29 @@ SuccessorRange::SuccessorRange(Operation *term) : SuccessorRange(nullptr, 0) {
   if ((count = term->getNumSuccessors()))
     base = term->getBlockOperands().data();
 }
+
+//===----------------------------------------------------------------------===//
+// BlockRange
+//===----------------------------------------------------------------------===//
+
+BlockRange::BlockRange(ArrayRef<Block *> blocks) : BlockRange(nullptr, 0) {
+  if ((count = blocks.size()))
+    base = blocks.data();
+}
+
+BlockRange::BlockRange(SuccessorRange successors)
+    : BlockRange(successors.begin().getBase(), successors.size()) {}
+
+/// See `llvm::detail::indexed_accessor_range_base` for details.
+BlockRange::OwnerT BlockRange::offset_base(OwnerT object, ptrdiff_t index) {
+  if (auto *operand = object.dyn_cast<BlockOperand *>())
+    return {operand + index};
+  return {object.dyn_cast<Block *const *>() + index};
+}
+
+/// See `llvm::detail::indexed_accessor_range_base` for details.
+Block *BlockRange::dereference_iterator(OwnerT object, ptrdiff_t index) {
+  if (const auto *operand = object.dyn_cast<BlockOperand *>())
+    return operand[index].get();
+  return object.dyn_cast<Block *const *>()[index];
+}

mlir/lib/IR/Operation.cpp

@@ -71,29 +71,24 @@ OperationName OperationName::getFromOpaquePointer(void *pointer) {
 
 /// Create a new Operation with the specific fields.
 Operation *Operation::create(Location location, OperationName name,
-                             ArrayRef<Type> resultTypes,
-                             ArrayRef<Value> operands,
+                             TypeRange resultTypes, ValueRange operands,
                              ArrayRef<NamedAttribute> attributes,
-                             ArrayRef<Block *> successors,
-                             unsigned numRegions) {
+                             BlockRange successors, unsigned numRegions) {
   return create(location, name, resultTypes, operands,
                 MutableDictionaryAttr(attributes), successors, numRegions);
 }
 
 /// Create a new Operation from operation state.
 Operation *Operation::create(const OperationState &state) {
-  return Operation::create(state.location, state.name, state.types,
-                           state.operands, state.attributes, state.successors,
-                           state.regions);
+  return create(state.location, state.name, state.types, state.operands,
+                state.attributes, state.successors, state.regions);
 }
 
 /// Create a new Operation with the specific fields.
 Operation *Operation::create(Location location, OperationName name,
-                             ArrayRef<Type> resultTypes,
-                             ArrayRef<Value> operands,
+                             TypeRange resultTypes, ValueRange operands,
                              MutableDictionaryAttr attributes,
-                             ArrayRef<Block *> successors,
-                             RegionRange regions) {
+                             BlockRange successors, RegionRange regions) {
   unsigned numRegions = regions.size();
   Operation *op = create(location, name, resultTypes, operands, attributes,
                          successors, numRegions);
@@ -106,11 +101,9 @@ Operation *Operation::create(Location location, OperationName name,
 /// Overload of create that takes an existing MutableDictionaryAttr to avoid
 /// unnecessarily uniquing a list of attributes.
 Operation *Operation::create(Location location, OperationName name,
-                             ArrayRef<Type> resultTypes,
-                             ArrayRef<Value> operands,
+                             TypeRange resultTypes, ValueRange operands,
                              MutableDictionaryAttr attributes,
-                             ArrayRef<Block *> successors,
-                             unsigned numRegions) {
+                             BlockRange successors, unsigned numRegions) {
   // We only need to allocate additional memory for a subset of results.
   unsigned numTrailingResults = OpResult::getNumTrailing(resultTypes.size());
   unsigned numInlineResults = OpResult::getNumInline(resultTypes.size());
@@ -167,7 +160,7 @@ Operation *Operation::create(Location location, OperationName name,
 }
 
 Operation::Operation(Location location, OperationName name,
-                     ArrayRef<Type> resultTypes, unsigned numSuccessors,
+                     TypeRange resultTypes, unsigned numSuccessors,
                      unsigned numRegions,
                      const MutableDictionaryAttr &attributes,
                      bool hasOperandStorage)
@@ -611,8 +604,8 @@ Operation *Operation::cloneWithoutRegions(BlockAndValueMapping &mapper) {
     successors.push_back(mapper.lookupOrDefault(successor));
 
   // Create the new operation.
-  auto *newOp = Operation::create(getLoc(), getName(), getResultTypes(),
-                                  operands, attrs, successors, getNumRegions());
+  auto *newOp = create(getLoc(), getName(), getResultTypes(), operands, attrs,
+                       successors, getNumRegions());
 
   // Remember the mapping of any results.
   for (unsigned i = 0, e = getNumResults(); i != e; ++i)

mlir/lib/IR/OperationSupport.cpp

@@ -186,7 +186,7 @@ void OperationState::addOperands(ValueRange newOperands) {
   operands.append(newOperands.begin(), newOperands.end());
 }
 
-void OperationState::addSuccessors(SuccessorRange newSuccessors) {
+void OperationState::addSuccessors(BlockRange newSuccessors) {
   successors.append(newSuccessors.begin(), newSuccessors.end());
 }