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

llvm/include/llvm/ADT/DenseMapInfo.h

 
 #ifndef LLVM_ADT_DENSEMAPINFO_H
 #define LLVM_ADT_DENSEMAPINFO_H
 
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <tuple>
+#include <type_traits>
 #include <utility>
+#include <variant>
 
 namespace llvm {
 
 namespace detail {
 
 /// Simplistic combination of 32-bit hash values into 32-bit hash values.
 static inline unsigned combineHashValue(unsigned a, unsigned b) {
   uint64_t key = (uint64_t)a << 32 | (uint64_t)b;
template <typename... Ts> struct DenseMapInfo<std::tuple<Ts...>> {
   }
 
   static bool isEqual(const Tuple &lhs, const Tuple &rhs) {
     std::integral_constant<bool, 0 == sizeof...(Ts)> atEnd;
     return isEqualImpl<0>(lhs, rhs, atEnd);
   }
 };
 
+// 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

llvm/unittests/ADT/DenseMapTest.cpp

 //===- llvm/unittest/ADT/DenseMapMap.cpp - DenseMap unit tests --*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #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;
 
 namespace {
 
 uint32_t getTestKey(int i, uint32_t *) { return i; }
 uint32_t getTestValue(int i, uint32_t *) { return 42 + i; }
 
TEST(DenseMapCustomTest, SFINAEMapInfo) {
   EXPECT_EQ(Map[Keys[0]], 1);
   EXPECT_EQ(Map[Keys[1]], 2);
   EXPECT_EQ(Map[Keys[2]], 3);
   Map.clear();
   EXPECT_EQ(Map.find(Keys[0]), Map.end());
   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