[llvm] 7302fe4 - [VPlan] Make blocksOnly work properly with ranges over const pointers.

[Florian Hahn via llvm-commits]

Author: Florian Hahn
Date: 2021-04-26T10:52:35+01:00
New Revision: 7302fe432843bd7e9b41b50ddaa37781f62b08cf

URL: https://github.com/llvm/llvm-project/commit/7302fe432843bd7e9b41b50ddaa37781f62b08cf
DIFF: https://github.com/llvm/llvm-project/commit/7302fe432843bd7e9b41b50ddaa37781f62b08cf.diff

LOG: [VPlan] Make blocksOnly work properly with ranges over const pointers.

When iterating over const blocks, the base type in the lambdas needs
to use const VPBlockBase *, otherwise it cannot be used with input
iterators over const VPBlockBase.

Also adjust the type of the input iterator range to const &, as it
does not take ownership of the input range.

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/VPlan.h
    llvm/unittests/Transforms/Vectorize/VPlanTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 91a3d9153c68..d0b20c20ad79 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -2202,16 +2202,22 @@ class VPBlockUtils {
     return Count;
   }
 
-  /// Return an iterator range over \p Iter which only includes \p BlockTy
+  /// Return an iterator range over \p Range which only includes \p BlockTy
   /// blocks. The accesses are casted to \p BlockTy.
-  template <typename BlockTy, typename T> static auto blocksOnly(T Iter) {
-    // We need to first create an iterator range over VPBlockBase & instead of
-    // VPBlockBase * for filter_range to work properly.
-    auto Mapped = map_range(
-        Iter, [](VPBlockBase *Block) -> VPBlockBase & { return *Block; });
+  template <typename BlockTy, typename T>
+  static auto blocksOnly(const T &Range) {
+    // Create BaseTy with correct const-ness based on BlockTy.
+    using BaseTy =
+        typename std::conditional<std::is_const<BlockTy>::value,
+                                  const VPBlockBase, VPBlockBase>::type;
+
+    // We need to first create an iterator range over (const) BlocktTy & instead
+    // of (const) BlockTy * for filter_range to work properly.
+    auto Mapped =
+        map_range(Range, [](BaseTy *Block) -> BaseTy & { return *Block; });
     auto Filter = make_filter_range(
-        Mapped, [](VPBlockBase &Block) { return isa<BlockTy>(&Block); });
-    return map_range(Filter, [](VPBlockBase &Block) -> BlockTy * {
+        Mapped, [](BaseTy &Block) { return isa<BlockTy>(&Block); });
+    return map_range(Filter, [](BaseTy &Block) -> BlockTy * {
       return cast<BlockTy>(&Block);
     });
   }

diff  --git a/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp b/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp
index 7c1843b6abb9..41bf33bc8db9 100644
--- a/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/VPlanTest.cpp
@@ -639,8 +639,9 @@ TEST(VPBasicBlockTest, TraversingIteratorTest) {
     EXPECT_EQ(VPBB1, FromIterator[6]);
 
     // Post-order, const VPRegionBlocks only.
+    VPBlockRecursiveTraversalWrapper<const VPBlockBase *> StartConst(VPBB1);
     SmallVector<const VPRegionBlock *> FromIteratorVPRegion(
-        VPBlockUtils::blocksOnly<const VPRegionBlock>(post_order(Start)));
+        VPBlockUtils::blocksOnly<const VPRegionBlock>(post_order(StartConst)));
     EXPECT_EQ(3u, FromIteratorVPRegion.size());
     EXPECT_EQ(R3, FromIteratorVPRegion[0]);
     EXPECT_EQ(R2, FromIteratorVPRegion[1]);