[llvm] 358ebdd - [DenseMap] Introduce lookup_or (#138887)

[via llvm-commits]

Author: Ramkumar Ramachandra
Date: 2025-05-08T11:22:10+01:00
New Revision: 358ebddeb836d1c0ac665a8a2faa2e07fd89da63

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

LOG: [DenseMap] Introduce lookup_or (#138887)

Introduce lookup_or, a variant of lookup, for non-default-constructible
values.

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index 3175b3ece467c..9f1dd8e51339a 100644
--- a/llvm/include/llvm/ADT/DenseMap.h
+++ b/llvm/include/llvm/ADT/DenseMap.h
@@ -217,6 +217,16 @@ class DenseMapBase : public DebugEpochBase {
     return ValueT();
   }
 
+  // Return the entry with the specified key, or \p Default. This variant is
+  // useful, because `lookup` cannot be used with non-default-constructible
+  // values.
+  ValueT lookup_or(const_arg_type_t<KeyT> Val,
+                   const_arg_type_t<ValueT> Default) const {
+    if (const BucketT *Bucket = doFind(Val))
+      return Bucket->getSecond();
+    return Default;
+  }
+
   /// at - Return the entry for the specified key, or abort if no such
   /// entry exists.
   const ValueT &at(const_arg_type_t<KeyT> Val) const {

diff  --git a/llvm/unittests/ADT/DenseMapTest.cpp b/llvm/unittests/ADT/DenseMapTest.cpp
index b9d519a23c9be..2bf0f90e6b070 100644
--- a/llvm/unittests/ADT/DenseMapTest.cpp
+++ b/llvm/unittests/ADT/DenseMapTest.cpp
@@ -616,6 +616,26 @@ TEST(DenseMapCustomTest, StringRefTest) {
   EXPECT_EQ(42, M.lookup(StringRef("a", 0)));
 }
 
+struct NonDefaultConstructible {
+  unsigned V;
+  NonDefaultConstructible(unsigned V) : V(V) {};
+  bool operator==(const NonDefaultConstructible &Other) const {
+    return V == Other.V;
+  }
+};
+
+TEST(DenseMapCustomTest, LookupOr) {
+  DenseMap<int, NonDefaultConstructible> M;
+
+  M.insert_or_assign(0, 3u);
+  M.insert_or_assign(1, 2u);
+  M.insert_or_assign(1, 0u);
+
+  EXPECT_EQ(M.lookup_or(0, 4u), 3u);
+  EXPECT_EQ(M.lookup_or(1, 4u), 0u);
+  EXPECT_EQ(M.lookup_or(2, 4u), 4u);
+}
+
 // Key traits that allows lookup with either an unsigned or char* key;
 // In the latter case, "a" == 0, "b" == 1 and so on.
 struct TestDenseMapInfo {