[llvm] [ADT] Avoid map storage for small SmallMapVector (PR #196473)
Mingjie Xu via llvm-commits
llvm-commits at lists.llvm.org
Thu May 7 23:14:26 PDT 2026
https://github.com/Enna1 created https://github.com/llvm/llvm-project/pull/196473
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.
>From a0ee51ce686455aa8bb8e24ac057e88b0bb9d9b9 Mon Sep 17 00:00:00 2001
From: Enna1 <xumingjie.enna1 at bytedance.com>
Date: Thu, 7 May 2026 11:45:53 +0800
Subject: [PATCH 1/2] [ADT] apply small-size optimization to SmallMapVector
---
llvm/include/llvm/ADT/MapVector.h | 99 +++++++++++++++++++++++++++++--
1 file changed, 93 insertions(+), 6 deletions(-)
diff --git a/llvm/include/llvm/ADT/MapVector.h b/llvm/include/llvm/ADT/MapVector.h
index 2b2f098dd3abf..fa52f806ca579 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>
@@ -32,7 +33,8 @@ namespace llvm {
/// 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 +110,15 @@ 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 =
+ find_if(Vector, [&Key](const auto &P) { return P.first == 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 +155,10 @@ class MapVector {
}
[[nodiscard]] bool contains(const KeyT &Key) const {
+ if constexpr (canBeSmall())
+ if (isSmall())
+ return any_of(Vector, [&Key](const auto &P) { return P.first == Key; });
+
return Map.find(Key) != Map.end();
}
@@ -152,11 +167,21 @@ class MapVector {
}
[[nodiscard]] iterator find(const KeyT &Key) {
+ if constexpr (canBeSmall())
+ if (isSmall())
+ return find_if(Vector,
+ [&Key](const auto &P) { return P.first == 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 find_if(Vector,
+ [&Key](const auto &P) { return P.first == Key; });
+
typename MapType::const_iterator Pos = Map.find(Key);
return Pos == Map.end() ? Vector.end() : (Vector.begin() + Pos->second);
}
@@ -164,6 +189,15 @@ 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 =
+ find_if(Vector, [&Key](const auto &P) { return P.first == 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 +206,15 @@ 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 =
+ find_if(Vector, [&Key](const auto &P) { return P.first == 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 +222,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 +241,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 +278,43 @@ class MapVector {
template <class Predicate> void remove_if(Predicate Pred);
private:
+ [[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 =
+ find_if(Vector, [&Key](const auto &P) { return P.first == 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 +327,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)) {
@@ -272,9 +360,8 @@ void MapVector<KeyT, ValueT, MapType, VectorType>::remove_if(Function Pred) {
/// size.
template <typename KeyT, typename ValueT, unsigned N>
struct SmallMapVector
- : MapVector<KeyT, ValueT, SmallDenseMap<KeyT, unsigned, N>,
- SmallVector<std::pair<KeyT, ValueT>, N>> {
-};
+ : MapVector<KeyT, ValueT, DenseMap<KeyT, unsigned>,
+ SmallVector<std::pair<KeyT, ValueT>, N>, N> {};
} // end namespace llvm
>From dad746a34d935753dd2b651353fd0f50071b8c7d Mon Sep 17 00:00:00 2001
From: Enna1 <xumingjie.enna1 at bytedance.com>
Date: Fri, 8 May 2026 11:44:28 +0800
Subject: [PATCH 2/2] use KeyInfo::isEqual and add test
---
llvm/include/llvm/ADT/MapVector.h | 64 +++++++++++++++++++++-------
llvm/unittests/ADT/MapVectorTest.cpp | 20 +++++++++
2 files changed, 68 insertions(+), 16 deletions(-)
diff --git a/llvm/include/llvm/ADT/MapVector.h b/llvm/include/llvm/ADT/MapVector.h
index fa52f806ca579..423e2a0fd6145 100644
--- a/llvm/include/llvm/ADT/MapVector.h
+++ b/llvm/include/llvm/ADT/MapVector.h
@@ -28,6 +28,19 @@
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.
@@ -112,8 +125,7 @@ class MapVector {
"Cannot call lookup() if ValueT is not copyable.");
if constexpr (canBeSmall())
if (isSmall()) {
- auto I =
- find_if(Vector, [&Key](const auto &P) { return P.first == Key; });
+ auto I = findInVector(Key);
if (I != Vector.end())
return I->second;
return ValueT();
@@ -157,7 +169,7 @@ class MapVector {
[[nodiscard]] bool contains(const KeyT &Key) const {
if constexpr (canBeSmall())
if (isSmall())
- return any_of(Vector, [&Key](const auto &P) { return P.first == Key; });
+ return findInVector(Key) != Vector.end();
return Map.find(Key) != Map.end();
}
@@ -169,8 +181,7 @@ class MapVector {
[[nodiscard]] iterator find(const KeyT &Key) {
if constexpr (canBeSmall())
if (isSmall())
- return find_if(Vector,
- [&Key](const auto &P) { return P.first == Key; });
+ return findInVector(Key);
typename MapType::const_iterator Pos = Map.find(Key);
return Pos == Map.end() ? Vector.end() : (Vector.begin() + Pos->second);
@@ -179,8 +190,7 @@ class MapVector {
[[nodiscard]] const_iterator find(const KeyT &Key) const {
if constexpr (canBeSmall())
if (isSmall())
- return find_if(Vector,
- [&Key](const auto &P) { return P.first == Key; });
+ return findInVector(Key);
typename MapType::const_iterator Pos = Map.find(Key);
return Pos == Map.end() ? Vector.end() : (Vector.begin() + Pos->second);
@@ -191,8 +201,7 @@ class MapVector {
[[nodiscard]] ValueT &at(const KeyT &Key) {
if constexpr (canBeSmall())
if (isSmall()) {
- auto I =
- find_if(Vector, [&Key](const auto &P) { return P.first == Key; });
+ auto I = findInVector(Key);
assert(I != Vector.end() &&
"MapVector::at failed due to a missing key");
return I->second;
@@ -208,8 +217,7 @@ class MapVector {
[[nodiscard]] const ValueT &at(const KeyT &Key) const {
if constexpr (canBeSmall())
if (isSmall()) {
- auto I =
- find_if(Vector, [&Key](const auto &P) { return P.first == Key; });
+ auto I = findInVector(Key);
assert(I != Vector.end() &&
"MapVector::at failed due to a missing key");
return I->second;
@@ -278,6 +286,31 @@ 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(); }
@@ -303,8 +336,7 @@ class MapVector {
std::pair<iterator, bool> try_emplace_impl(KeyArgT &&Key, Ts &&...Args) {
if constexpr (canBeSmall())
if (isSmall()) {
- auto I =
- find_if(Vector, [&Key](const auto &P) { return P.first == Key; });
+ auto I = findInVector(Key);
if (I != Vector.end())
return {I, false};
Vector.emplace_back(std::piecewise_construct,
@@ -359,9 +391,9 @@ void MapVector<KeyT, ValueT, MapType, VectorType, N>::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, DenseMap<KeyT, unsigned>,
- SmallVector<std::pair<KeyT, ValueT>, N>, 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;
More information about the llvm-commits
mailing list