diff --git a/llvm/include/llvm/ADT/DenseMapInfo.h b/llvm/include/llvm/ADT/DenseMapInfo.h --- a/llvm/include/llvm/ADT/DenseMapInfo.h +++ b/llvm/include/llvm/ADT/DenseMapInfo.h @@ -18,7 +18,9 @@ #include #include #include +#include #include +#include namespace llvm { @@ -288,6 +290,37 @@ } }; +// Provide DenseMapInfo for variants whose all alternatives have DenseMapInfo. +template struct DenseMapInfo> { + using Variant = std::variant; + using FirstT = std::variant_alternative_t<0, Variant>; + + static inline Variant getEmptyKey() { + return Variant(std::in_place_index<0>, DenseMapInfo::getEmptyKey()); + } + + static inline Variant getTombstoneKey() { + return Variant(std::in_place_index<0>, + DenseMapInfo::getTombstoneKey()); + } + + static unsigned getHashValue(const Variant &Val) { + return std::visit( + [&Val](auto &&Alternative) { + using T = std::decay_t; + // Include index in hash to make sure same value as different + // alternatives don't collide. + return detail::combineHashValue( + DenseMapInfo::getHashValue(Val.index()), + DenseMapInfo::getHashValue(Alternative)); + }, + Val); + } + + static bool isEqual(const Variant &LHS, const Variant &RHS) { + return LHS == RHS; + } +}; } // end namespace llvm #endif // LLVM_ADT_DENSEMAPINFO_H diff --git a/llvm/unittests/ADT/DenseMapTest.cpp b/llvm/unittests/ADT/DenseMapTest.cpp --- a/llvm/unittests/ADT/DenseMapTest.cpp +++ b/llvm/unittests/ADT/DenseMapTest.cpp @@ -7,9 +7,13 @@ //===----------------------------------------------------------------------===// #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseMapInfo.h" +#include "gmock/gmock.h" #include "gtest/gtest.h" #include #include +#include +#include using namespace llvm; @@ -707,4 +711,18 @@ EXPECT_EQ(Map.find(Keys[1]), Map.end()); EXPECT_EQ(Map.find(Keys[2]), Map.end()); } + +TEST(DenseMapCustomTest, VariantSupport) { + using variant = std::variant; + DenseMap Map; + variant Keys[] = { + variant(std::in_place_index<0>, 1), + variant(std::in_place_index<1>, 1), + }; + Map.try_emplace(Keys[0], 0); + Map.try_emplace(Keys[1], 1); + EXPECT_THAT(Map, testing::SizeIs(2)); + EXPECT_NE(DenseMapInfo::getHashValue(Keys[0]), + DenseMapInfo::getHashValue(Keys[1])); +} } // namespace