[llvm] d72b57b - DenseMap: support enum class keys (#95972)

[via llvm-commits]

Author: Ramkumar Ramachandra
Date: 2024-06-19T10:29:37+01:00
New Revision: d72b57bbf7c3448077f82147e3ce1a72b2c86a34

URL: https://github.com/llvm/llvm-project/commit/d72b57bbf7c3448077f82147e3ce1a72b2c86a34
DIFF: https://github.com/llvm/llvm-project/commit/d72b57bbf7c3448077f82147e3ce1a72b2c86a34.diff

LOG: DenseMap: support enum class keys (#95972)

Implemented using std::underlying_type.

Added: 
    

Modified: 
    llvm/include/llvm/ADT/DenseMapInfo.h
    llvm/unittests/ADT/DenseMapTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/DenseMapInfo.h b/llvm/include/llvm/ADT/DenseMapInfo.h
index 5b7dce7b53c62..c405f7446d5a9 100644
--- a/llvm/include/llvm/ADT/DenseMapInfo.h
+++ b/llvm/include/llvm/ADT/DenseMapInfo.h
@@ -297,6 +297,24 @@ template <typename... Ts> struct DenseMapInfo<std::tuple<Ts...>> {
   }
 };
 
+// Provide DenseMapInfo for enum classes.
+template <typename Enum>
+struct DenseMapInfo<Enum, std::enable_if_t<std::is_enum_v<Enum>>> {
+  using UnderlyingType = std::underlying_type_t<Enum>;
+  using Info = DenseMapInfo<UnderlyingType>;
+
+  static Enum getEmptyKey() { return static_cast<Enum>(Info::getEmptyKey()); }
+
+  static Enum getTombstoneKey() {
+    return static_cast<Enum>(Info::getTombstoneKey());
+  }
+
+  static unsigned getHashValue(const Enum &Val) {
+    return Info::getHashValue(static_cast<UnderlyingType>(Val));
+  }
+
+  static bool isEqual(const Enum &LHS, const Enum &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 cb45a6e681ce4..d1bbdde8dfc26 100644
--- a/llvm/unittests/ADT/DenseMapTest.cpp
+++ b/llvm/unittests/ADT/DenseMapTest.cpp
@@ -21,7 +21,6 @@
 using namespace llvm;
 
 namespace {
-
 uint32_t getTestKey(int i, uint32_t *) { return i; }
 uint32_t getTestValue(int i, uint32_t *) { return 42 + i; }
 
@@ -36,6 +35,14 @@ uint32_t *getTestValue(int i, uint32_t **) {
   return &dummy_arr1[i];
 }
 
+enum class EnumClass { Val };
+
+EnumClass getTestKey(int i, EnumClass *) {
+  // We can't possibly support 100 values for the swap test, so just return an
+  // invalid EnumClass for testing.
+  return static_cast<EnumClass>(i);
+}
+
 /// A test class that tries to check that construction and destruction
 /// occur correctly.
 class CtorTester {
@@ -104,14 +111,19 @@ template <typename T>
 typename T::mapped_type *const DenseMapTest<T>::dummy_value_ptr = nullptr;
 
 // Register these types for testing.
+// clang-format off
 typedef ::testing::Types<DenseMap<uint32_t, uint32_t>,
                          DenseMap<uint32_t *, uint32_t *>,
                          DenseMap<CtorTester, CtorTester, CtorTesterMapInfo>,
+                         DenseMap<EnumClass, uint32_t>,
                          SmallDenseMap<uint32_t, uint32_t>,
                          SmallDenseMap<uint32_t *, uint32_t *>,
                          SmallDenseMap<CtorTester, CtorTester, 4,
-                                       CtorTesterMapInfo>
+                                       CtorTesterMapInfo>,
+                         SmallDenseMap<EnumClass, uint32_t>
                          > DenseMapTestTypes;
+// clang-format on
+
 TYPED_TEST_SUITE(DenseMapTest, DenseMapTestTypes, );
 
 // Empty map tests