[llvm] [ADT] Avoid map storage for small SmallMapVector (PR #196473)

via llvm-commits llvm-commits at lists.llvm.org
Fri May 8 00:09:02 PDT 2026


llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-adt

Author: Mingjie Xu (Enna1)

<details>
<summary>Changes</summary>

SmallMapVector previously used SmallDenseMap for its index, which still
initializes and maintains map storage even when the number of entries is tiny.

Teach MapVector to support a vector-only small mode. While the entry count stays
within the configured small size, operations use the underlying vector directly.
When the size grows past the threshold, the map index is built and subsequent
operations use the regular MapVector path.

This mirrors the small-size strategy used by SmallSetVector.

---
Full diff: https://github.com/llvm/llvm-project/pull/196473.diff


2 Files Affected:

- (modified) llvm/include/llvm/ADT/MapVector.h (+125-6) 
- (modified) llvm/unittests/ADT/MapVectorTest.cpp (+20) 


``````````diff
diff --git a/llvm/include/llvm/ADT/MapVector.h b/llvm/include/llvm/ADT/MapVector.h
index 2b2f098dd3abf..423e2a0fd6145 100644
--- a/llvm/include/llvm/ADT/MapVector.h
+++ b/llvm/include/llvm/ADT/MapVector.h
@@ -18,6 +18,7 @@
 #define LLVM_ADT_MAPVECTOR_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include <cassert>
 #include <cstddef>
@@ -27,12 +28,26 @@
 
 namespace llvm {
 
+namespace detail {
+template <typename MapT> struct MapVectorKeyInfo {};
+
+template <typename KeyT, typename ValueT, typename KeyInfoT, typename BucketT>
+struct MapVectorKeyInfo<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>> {
+  using type = KeyInfoT;
+};
+
+template <typename MapT>
+using map_vector_key_info_t = typename MapVectorKeyInfo<MapT>::type;
+
+} // namespace detail
+
 /// This class implements a map that also provides access to all stored values
 /// in a deterministic order. The values are kept in a SmallVector<*, 0> and the
 /// mapping is done with DenseMap from Keys to indexes in that vector.
 template <typename KeyT, typename ValueT,
           typename MapType = DenseMap<KeyT, unsigned>,
-          typename VectorType = SmallVector<std::pair<KeyT, ValueT>, 0>>
+          typename VectorType = SmallVector<std::pair<KeyT, ValueT>, 0>,
+          unsigned N = 0>
 class MapVector {
 public:
   using key_type = KeyT;
@@ -108,6 +123,14 @@ class MapVector {
   [[nodiscard]] ValueT lookup(const KeyT &Key) const {
     static_assert(std::is_copy_constructible_v<ValueT>,
                   "Cannot call lookup() if ValueT is not copyable.");
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        auto I = findInVector(Key);
+        if (I != Vector.end())
+          return I->second;
+        return ValueT();
+      }
+
     typename MapType::const_iterator Pos = Map.find(Key);
     return Pos == Map.end()? ValueT() : Vector[Pos->second].second;
   }
@@ -144,6 +167,10 @@ class MapVector {
   }
 
   [[nodiscard]] bool contains(const KeyT &Key) const {
+    if constexpr (canBeSmall())
+      if (isSmall())
+        return findInVector(Key) != Vector.end();
+
     return Map.find(Key) != Map.end();
   }
 
@@ -152,11 +179,19 @@ class MapVector {
   }
 
   [[nodiscard]] iterator find(const KeyT &Key) {
+    if constexpr (canBeSmall())
+      if (isSmall())
+        return findInVector(Key);
+
     typename MapType::const_iterator Pos = Map.find(Key);
     return Pos == Map.end() ? Vector.end() : (Vector.begin() + Pos->second);
   }
 
   [[nodiscard]] const_iterator find(const KeyT &Key) const {
+    if constexpr (canBeSmall())
+      if (isSmall())
+        return findInVector(Key);
+
     typename MapType::const_iterator Pos = Map.find(Key);
     return Pos == Map.end() ? Vector.end() : (Vector.begin() + Pos->second);
   }
@@ -164,6 +199,14 @@ class MapVector {
   /// at - Return the entry for the specified key, or abort if no such
   /// entry exists.
   [[nodiscard]] ValueT &at(const KeyT &Key) {
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        auto I = findInVector(Key);
+        assert(I != Vector.end() &&
+               "MapVector::at failed due to a missing key");
+        return I->second;
+      }
+
     auto Pos = Map.find(Key);
     assert(Pos != Map.end() && "MapVector::at failed due to a missing key");
     return Vector[Pos->second].second;
@@ -172,6 +215,14 @@ class MapVector {
   /// at - Return the entry for the specified key, or abort if no such
   /// entry exists.
   [[nodiscard]] const ValueT &at(const KeyT &Key) const {
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        auto I = findInVector(Key);
+        assert(I != Vector.end() &&
+               "MapVector::at failed due to a missing key");
+        return I->second;
+      }
+
     auto Pos = Map.find(Key);
     assert(Pos != Map.end() && "MapVector::at failed due to a missing key");
     return Vector[Pos->second].second;
@@ -179,6 +230,12 @@ class MapVector {
 
   /// Remove the last element from the vector.
   void pop_back() {
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        Vector.pop_back();
+        return;
+      }
+
     typename MapType::iterator Pos = Map.find(Vector.back().first);
     Map.erase(Pos);
     Vector.pop_back();
@@ -192,6 +249,10 @@ class MapVector {
   /// \note This is a deceivingly expensive operation (linear time).  It's
   /// usually better to use \a remove_if() if possible.
   typename VectorType::iterator erase(typename VectorType::iterator Iterator) {
+    if constexpr (canBeSmall())
+      if (isSmall())
+        return Vector.erase(Iterator);
+
     Map.erase(Iterator->first);
     auto Next = Vector.erase(Iterator);
     if (Next == Vector.end())
@@ -225,15 +286,67 @@ class MapVector {
   template <class Predicate> void remove_if(Predicate Pred);
 
 private:
+  template <typename LookupKeyT, typename StoredKeyT>
+  [[nodiscard]] static bool keyEqual(const LookupKeyT &LookupKey,
+                                     const StoredKeyT &StoredKey) {
+    if constexpr (is_detected<detail::map_vector_key_info_t, MapType>::value) {
+      using KeyInfo = detail::map_vector_key_info_t<MapType>;
+      return KeyInfo::isEqual(LookupKey, StoredKey);
+    } else {
+      static_assert(has_equality_comparison_v<LookupKeyT, StoredKeyT>,
+                    "MapVector small mode requires MapType key info or "
+                    "operator== for key comparison");
+      return LookupKey == StoredKey;
+    }
+  }
+
+  template <typename T> [[nodiscard]] iterator findInVector(const T &Key) {
+    return find_if(Vector,
+                   [&Key](const auto &P) { return keyEqual(Key, P.first); });
+  }
+
+  template <typename T>
+  [[nodiscard]] const_iterator findInVector(const T &Key) const {
+    return find_if(Vector,
+                   [&Key](const auto &P) { return keyEqual(Key, P.first); });
+  }
+
+  [[nodiscard]] static constexpr bool canBeSmall() { return N != 0; }
+
+  [[nodiscard]] bool isSmall() const { return Map.empty(); }
+
+  void makeBig() {
+    if constexpr (canBeSmall()) {
+      unsigned Index = 0;
+      for (const auto &entry : Vector)
+        Map[entry.first] = Index++;
+    }
+  }
+
   MapType Map;
   VectorType Vector;
 
+  static_assert(N <= 32, "Small size should be less than or equal to 32!");
+
   static_assert(
       std::is_integral_v<typename MapType::mapped_type>,
       "The mapped_type of the specified Map must be an integral type");
 
   template <typename KeyArgT, typename... Ts>
   std::pair<iterator, bool> try_emplace_impl(KeyArgT &&Key, Ts &&...Args) {
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        auto I = findInVector(Key);
+        if (I != Vector.end())
+          return {I, false};
+        Vector.emplace_back(std::piecewise_construct,
+                            std::forward_as_tuple(std::forward<KeyArgT>(Key)),
+                            std::forward_as_tuple(std::forward<Ts>(Args)...));
+        if (Vector.size() > N)
+          makeBig();
+        return {std::prev(end()), true};
+      }
+
     auto [It, Inserted] = Map.try_emplace(Key);
     if (Inserted) {
       It->second = Vector.size();
@@ -246,9 +359,16 @@ class MapVector {
   }
 };
 
-template <typename KeyT, typename ValueT, typename MapType, typename VectorType>
+template <typename KeyT, typename ValueT, typename MapType, typename VectorType,
+          unsigned N>
 template <class Function>
-void MapVector<KeyT, ValueT, MapType, VectorType>::remove_if(Function Pred) {
+void MapVector<KeyT, ValueT, MapType, VectorType, N>::remove_if(Function Pred) {
+  if constexpr (canBeSmall())
+    if (isSmall()) {
+      Vector.erase(llvm::remove_if(Vector, Pred), Vector.end());
+      return;
+    }
+
   auto O = Vector.begin();
   for (auto I = O, E = Vector.end(); I != E; ++I) {
     if (Pred(*I)) {
@@ -271,9 +391,8 @@ void MapVector<KeyT, ValueT, MapType, VectorType>::remove_if(Function Pred) {
 /// A MapVector that performs no allocations if smaller than a certain
 /// size.
 template <typename KeyT, typename ValueT, unsigned N>
-struct SmallMapVector
-    : MapVector<KeyT, ValueT, SmallDenseMap<KeyT, unsigned, N>,
-                SmallVector<std::pair<KeyT, ValueT>, N>> {
+struct SmallMapVector : MapVector<KeyT, ValueT, DenseMap<KeyT, unsigned>,
+                                  SmallVector<std::pair<KeyT, ValueT>, N>, N> {
 };
 
 } // end namespace llvm
diff --git a/llvm/unittests/ADT/MapVectorTest.cpp b/llvm/unittests/ADT/MapVectorTest.cpp
index b11d4603b90b7..b1d7b9bb09e29 100644
--- a/llvm/unittests/ADT/MapVectorTest.cpp
+++ b/llvm/unittests/ADT/MapVectorTest.cpp
@@ -482,6 +482,26 @@ TEST(SmallMapVectorSmallTest, NonCopyable) {
   ASSERT_EQ(*MV.find(2)->second, 2);
 }
 
+TEST(SmallMapVectorSmallTest, UsesDenseMapInfoEquality) {
+  SmallMapVector<A, int, 4> MV;
+
+  auto R0 = MV.try_emplace(A(0), 1);
+  EXPECT_TRUE(R0.second);
+
+  auto R1 = MV.try_emplace(A(0), 2);
+  EXPECT_FALSE(R1.second);
+  EXPECT_EQ(R1.first, R0.first);
+  EXPECT_EQ(R1.first->second, 1);
+  EXPECT_EQ(MV.size(), 1u);
+
+  EXPECT_TRUE(MV.contains(A(0)));
+  EXPECT_EQ(MV.find(A(0)), MV.begin());
+  EXPECT_EQ(MV.lookup(A(0)), 1);
+  EXPECT_EQ(MV.at(A(0)), 1);
+  EXPECT_EQ(MV.erase(A(0)), 1u);
+  EXPECT_TRUE(MV.empty());
+}
+
 TEST(SmallMapVectorLargeTest, insert_pop) {
   SmallMapVector<int, int, 1> MV;
   std::pair<SmallMapVector<int, int, 1>::iterator, bool> R;

``````````

</details>


https://github.com/llvm/llvm-project/pull/196473


More information about the llvm-commits mailing list