DenseMap: support enum class keys

[artagnon]

Implemented using std::underlying_type.

[llvmbot]

@llvm/pr-subscribers-llvm-adt

Author: Ramkumar Ramachandra (artagnon)

Changes

Implemented using std::underlying_type.

---

Full diff: https://github.com/llvm/llvm-project/pull/95972.diff

2 Files Affected:

• (modified) llvm/include/llvm/ADT/DenseMapInfo.h (+21)
• (modified) llvm/unittests/ADT/MapVectorTest.cpp (+24)

diff --git a/llvm/include/llvm/ADT/DenseMapInfo.h b/llvm/include/llvm/ADT/DenseMapInfo.h
index 5b7dce7b53c62..14440d123289e 100644
--- a/llvm/include/llvm/ADT/DenseMapInfo.h
+++ b/llvm/include/llvm/ADT/DenseMapInfo.h
@@ -297,6 +297,27 @@ 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 Info::isEqual(static_cast<UnderlyingType>(LHS),
+                         static_cast<UnderlyingType>(RHS));
+  }
+};
 } // end namespace llvm
 
 #endif // LLVM_ADT_DENSEMAPINFO_H
diff --git a/llvm/unittests/ADT/MapVectorTest.cpp b/llvm/unittests/ADT/MapVectorTest.cpp
index e0f11b60a0223..2c61acb29ce20 100644
--- a/llvm/unittests/ADT/MapVectorTest.cpp
+++ b/llvm/unittests/ADT/MapVectorTest.cpp
@@ -267,6 +267,30 @@ TEST(MapVectorTest, NonCopyable) {
   ASSERT_EQ(*MV.find(2)->second, 2);
 }
 
+TEST(MapVectorTest, EnumClassKey) {
+  enum class EC1 { ValA, ValB };
+  enum class EC2 { ValA, ValB };
+  MapVector<EC1, int> MV1;
+  MapVector<EC2, int> MV2;
+
+  ASSERT_TRUE(MV1.empty());
+  ASSERT_TRUE(MV2.empty());
+  MV1.insert({EC1::ValA, 13});
+  MV1.insert({EC1::ValB, 7});
+  MV2.insert({EC2::ValA, 42});
+  MV2.insert({EC2::ValB, 43});
+
+  ASSERT_EQ(MV1.count(EC1::ValA), 1u);
+  ASSERT_EQ(MV2.count(EC2::ValA), 1u);
+  ASSERT_EQ(MV2[EC2::ValB], 43);
+  ASSERT_NE(DenseMapInfo<EC1>::getHashValue(EC1::ValA),
+            DenseMapInfo<EC1>::getHashValue(EC1::ValB));
+  ASSERT_NE(EC2::ValA, DenseMapInfo<EC2>::getTombstoneKey());
+  ASSERT_NE(EC2::ValB, DenseMapInfo<EC2>::getTombstoneKey());
+  ASSERT_NE(EC2::ValA, DenseMapInfo<EC2>::getEmptyKey());
+  ASSERT_NE(EC2::ValB, DenseMapInfo<EC2>::getEmptyKey());
+}
+
 template <class IntType> struct MapVectorMappedTypeTest : ::testing::Test {
   using int_type = IntType;
 };

[nikic]

Looks reasonable.

[nikic]

Though I would like to hear where you plan to use this. Usually enums imply that there is a reasonably small, continuously numbered and fixed set of keys, in which case using a vector instead of map is often preferable.

[artagnon]

Though I would like to hear where you plan to use this. Usually enums imply that there is a reasonably small, continuously numbered and fixed set of keys, in which case using a vector instead of map is often preferable.

This is a requirement for my new Presburger parser, which keeps a map of {DimOrSymbol, position} -> multiplier. The parser isn't ready just yet, but will be soon'ish.

[nikic]

Though I would like to hear where you plan to use this. Usually enums imply that there is a reasonably small, continuously numbered and fixed set of keys, in which case using a vector instead of map is often preferable.

This is a requirement for my new Presburger parser, which keeps a map of {DimOrSymbol, position} -> multiplier. The parser isn't ready just yet, but will be soon'ish.

Ah, so this is mostly a building block for hashing std::pair? That makes sense.