[MapVector] Introduce {keys,values} iterator, use in VPlan

Similar to DenseMap::{keys,values}, introduce MapVector::{keys,values},
and use it to simplify some code in VPlan.

Author: artagnon

llvm/include/llvm/ADT/MapVector.h

@@ -99,6 +99,26 @@ class MapVector {
     return try_emplace_impl(Key).first->second;
   }
 
+  [[nodiscard]] inline auto keys() {
+    return map_range(
+        Vector, [](const std::pair<KeyT, ValueT> &KV) { return KV.first; });
+  }
+
+  [[nodiscard]] inline auto values() {
+    return map_range(
+        Vector, [](const std::pair<KeyT, ValueT> &KV) { return KV.second; });
+  }
+
+  [[nodiscard]] inline auto keys() const {
+    return map_range(
+        Vector, [](const std::pair<KeyT, ValueT> &KV) { return KV.first; });
+  }
+
+  [[nodiscard]] inline auto values() const {
+    return map_range(
+        Vector, [](const std::pair<KeyT, ValueT> &KV) { return KV.second; });
+  }
+
   // Returns a copy of the value.  Only allowed if ValueT is copyable.
   [[nodiscard]] ValueT lookup(const KeyT &Key) const {
     static_assert(std::is_copy_constructible_v<ValueT>,

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -25,7 +25,7 @@
 #define LLVM_TRANSFORMS_VECTORIZE_VPLAN_H
 
 #include "VPlanValue.h"
-#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Twine.h"
@@ -4336,13 +4336,9 @@ class VPlan {
   /// Represents the loop-invariant VF * UF of the vector loop region.
   VPValue VFxUF;
 
-  /// Holds a mapping between Values and their corresponding VPValue inside
-  /// VPlan.
-  Value2VPValueTy Value2VPValue;
-
-  /// Contains all the external definitions created for this VPlan. External
-  /// definitions are VPValues that hold a pointer to their underlying IR.
-  SmallVector<VPValue *, 16> VPLiveIns;
+  /// Contains all the external definitions created for this VPlan, as a mapping
+  /// from IR Values to VPValues.
+  SmallMapVector<Value *, VPValue *, 16> LiveIns;
 
   /// Blocks allocated and owned by the VPlan. They will be deleted once the
   /// VPlan is destroyed.
@@ -4539,10 +4535,9 @@ class VPlan {
   ///  yet) for \p V.
   VPValue *getOrAddLiveIn(Value *V) {
     assert(V && "Trying to get or add the VPValue of a null Value");
-    auto [It, Inserted] = Value2VPValue.try_emplace(V);
+    auto [It, Inserted] = LiveIns.try_emplace(V);
     if (Inserted) {
       VPValue *VPV = new VPValue(V);
-      VPLiveIns.push_back(VPV);
       assert(VPV->isLiveIn() && "VPV must be a live-in.");
       It->second = VPV;
     }
@@ -4574,17 +4569,10 @@ class VPlan {
   }
 
   /// Return the live-in VPValue for \p V, if there is one or nullptr otherwise.
-  VPValue *getLiveIn(Value *V) const { return Value2VPValue.lookup(V); }
+  VPValue *getLiveIn(Value *V) const { return LiveIns.lookup(V); }
 
   /// Return the list of live-in VPValues available in the VPlan.
-  ArrayRef<VPValue *> getLiveIns() const {
-    assert(all_of(Value2VPValue,
-                  [this](const auto &P) {
-                    return is_contained(VPLiveIns, P.second);
-                  }) &&
-           "all VPValues in Value2VPValue must also be in VPLiveIns");
-    return VPLiveIns;
-  }
+  auto getLiveIns() const { return LiveIns.values(); }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the live-ins of this VPlan to \p O.

llvm/lib/Transforms/Vectorize/VPlanValue.h

@@ -20,10 +20,8 @@
 #ifndef LLVM_TRANSFORMS_VECTORIZE_VPLAN_VALUE_H
 #define LLVM_TRANSFORMS_VECTORIZE_VPLAN_VALUE_H
 
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/Compiler.h"
@@ -196,9 +194,6 @@ class LLVM_ABI_FOR_TEST VPValue {
   }
 };
 
-typedef DenseMap<Value *, VPValue *> Value2VPValueTy;
-typedef DenseMap<VPValue *, Value *> VPValue2ValueTy;
-
 LLVM_ABI_FOR_TEST raw_ostream &operator<<(raw_ostream &OS,
                                           const VPRecipeBase &R);
 

llvm/unittests/ADT/MapVectorTest.cpp

@@ -307,6 +307,24 @@ TEST(MapVectorTest, AtTest) {
   EXPECT_EQ(ConstMV.at(1), 12);
 }
 
+TEST(MapVectorTest, KeysValuesIterator) {
+  MapVector<int, int> MV;
+
+  MV.insert(std::make_pair(1, 11));
+  MV.insert(std::make_pair(2, 12));
+  MV.insert(std::make_pair(3, 13));
+  MV.insert(std::make_pair(4, 14));
+  MV.insert(std::make_pair(5, 15));
+  MV.insert(std::make_pair(6, 16));
+
+  EXPECT_THAT(MV.keys(), testing::ElementsAre(1, 2, 3, 4, 5, 6));
+  EXPECT_THAT(MV.values(), testing::ElementsAre(11, 12, 13, 14, 15, 16));
+
+  const MapVector<int, int> &ConstMV = MV;
+  EXPECT_THAT(ConstMV.keys(), testing::ElementsAre(1, 2, 3, 4, 5, 6));
+  EXPECT_THAT(ConstMV.values(), testing::ElementsAre(11, 12, 13, 14, 15, 16));
+}
+
 template <class IntType> struct MapVectorMappedTypeTest : ::testing::Test {
   using int_type = IntType;
 };