From 979ffe97b9bf925a9d1448d13aa4737cfda8fc35 Mon Sep 17 00:00:00 2001
From: Mogball <jeffniu22@gmail.com>
Date: Fri, 25 Mar 2022 22:44:34 +0000
Subject: [PATCH] [mlir] Allow RegionRange to accept ArrayRef<Region *>

Adds another pointer to the union in RegionRange to allow RegionRange to work on ArrayRef<Region *> (i.e. vectors of Region *).

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D122514
---
 mlir/include/mlir/IR/Region.h | 12 ++++++++----
 mlir/lib/IR/Region.cpp        | 14 ++++++++++----
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/mlir/include/mlir/IR/Region.h b/mlir/include/mlir/IR/Region.h
index c2694ad..3a1e57b 100644
--- a/mlir/include/mlir/IR/Region.h
+++ b/mlir/include/mlir/IR/Region.h
@@ -321,11 +321,14 @@ private:
 /// parameter.
 class RegionRange
     : public llvm::detail::indexed_accessor_range_base<
-          RegionRange, PointerUnion<Region *, const std::unique_ptr<Region> *>,
+          RegionRange,
+          PointerUnion<Region *, const std::unique_ptr<Region> *, Region **>,
           Region *, Region *, Region *> {
-  /// The type representing the owner of this range. This is either a list of
-  /// values, operands, or results.
-  using OwnerT = PointerUnion<Region *, const std::unique_ptr<Region> *>;
+  /// The type representing the owner of this range. This is either an owning
+  /// list of regions, a list of region unique pointers, or a list of region
+  /// pointers.
+  using OwnerT =
+      PointerUnion<Region *, const std::unique_ptr<Region> *, Region **>;
 
 public:
   using RangeBaseT::RangeBaseT;
@@ -339,6 +342,7 @@ public:
       : RegionRange(ArrayRef<std::unique_ptr<Region>>(std::forward<Arg>(arg))) {
   }
   RegionRange(ArrayRef<std::unique_ptr<Region>> regions);
+  RegionRange(ArrayRef<Region *> regions);
 
 private:
   /// See `llvm::detail::indexed_accessor_range_base` for details.
diff --git a/mlir/lib/IR/Region.cpp b/mlir/lib/IR/Region.cpp
index 8739b98..a5bca20 100644
--- a/mlir/lib/IR/Region.cpp
+++ b/mlir/lib/IR/Region.cpp
@@ -228,18 +228,24 @@ RegionRange::RegionRange(MutableArrayRef<Region> regions)
     : RegionRange(regions.data(), regions.size()) {}
 RegionRange::RegionRange(ArrayRef<std::unique_ptr<Region>> regions)
     : RegionRange(regions.data(), regions.size()) {}
+RegionRange::RegionRange(ArrayRef<Region *> regions)
+    : RegionRange(const_cast<Region **>(regions.data()), regions.size()) {}
 
 /// See `llvm::detail::indexed_accessor_range_base` for details.
 RegionRange::OwnerT RegionRange::offset_base(const OwnerT &owner,
                                              ptrdiff_t index) {
-  if (auto *operand = owner.dyn_cast<const std::unique_ptr<Region> *>())
-    return operand + index;
+  if (auto *region = owner.dyn_cast<const std::unique_ptr<Region> *>())
+    return region + index;
+  if (auto **region = owner.dyn_cast<Region **>())
+    return region + index;
   return &owner.get<Region *>()[index];
 }
 /// See `llvm::detail::indexed_accessor_range_base` for details.
 Region *RegionRange::dereference_iterator(const OwnerT &owner,
                                           ptrdiff_t index) {
-  if (auto *operand = owner.dyn_cast<const std::unique_ptr<Region> *>())
-    return operand[index].get();
+  if (auto *region = owner.dyn_cast<const std::unique_ptr<Region> *>())
+    return region[index].get();
+  if (auto **region = owner.dyn_cast<Region **>())
+    return region[index];
   return &owner.get<Region *>()[index];
 }
-- 
2.7.4