[llvm][Support] Add DenseMapInfo for std::variant

Differential Revision: https://reviews.llvm.org/D133200

2 files changed, 51 insertions, 0 deletions

diff --git a/llvm/include/llvm/ADT/DenseMapInfo.h b/llvm/include/llvm/ADT/DenseMapInfo.h
index afd478f..01e96a5 100644
--- a/llvm/include/llvm/ADT/DenseMapInfo.h
+++ b/llvm/include/llvm/ADT/DenseMapInfo.h
@@ -18,7 +18,9 @@
 #include <cstddef>
 #include <cstdint>
 #include <tuple>
+#include <type_traits>
 #include <utility>
+#include <variant>
 
 namespace llvm {
 
@@ -288,6 +290,37 @@ template <typename... Ts> struct DenseMapInfo<std::tuple<Ts...>> {
   }
 };
 
+// Provide DenseMapInfo for variants whose all alternatives have DenseMapInfo.
+template <typename... Ts> struct DenseMapInfo<std::variant<Ts...>> {
+  using Variant = std::variant<Ts...>;
+  using FirstT = std::variant_alternative_t<0, Variant>;
+
+  static inline Variant getEmptyKey() {
+    return Variant(std::in_place_index<0>, DenseMapInfo<FirstT>::getEmptyKey());
+  }
+
+  static inline Variant getTombstoneKey() {
+    return Variant(std::in_place_index<0>,
+                   DenseMapInfo<FirstT>::getTombstoneKey());
+  }
+
+  static unsigned getHashValue(const Variant &Val) {
+    return std::visit(
+        [&Val](auto &&Alternative) {
+          using T = std::decay_t<decltype(Alternative)>;
+          // Include index in hash to make sure same value as different
+          // alternatives don't collide.
+          return detail::combineHashValue(
+              DenseMapInfo<size_t>::getHashValue(Val.index()),
+              DenseMapInfo<T>::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
index 619ddb0..cb77ad8 100644
--- 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 <map>
 #include <set>
+#include <utility>
+#include <variant>
 
 using namespace llvm;
 
@@ -707,4 +711,18 @@ TEST(DenseMapCustomTest, SFINAEMapInfo) {
   EXPECT_EQ(Map.find(Keys[1]), Map.end());
   EXPECT_EQ(Map.find(Keys[2]), Map.end());
 }
+
+TEST(DenseMapCustomTest, VariantSupport) {
+  using variant = std::variant<int, int>;
+  DenseMap<variant, int> 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<variant>::getHashValue(Keys[0]),
+            DenseMapInfo<variant>::getHashValue(Keys[1]));
+}
 } // namespace