[DenseMap] Introduce lookup_or

[artagnon]

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

[llvmbot]

@llvm/pr-subscribers-llvm-adt

Author: Ramkumar Ramachandra (artagnon)

Changes

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

---

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

2 Files Affected:

• (modified) llvm/include/llvm/ADT/DenseMap.h (+10)
• (modified) llvm/unittests/ADT/DenseMapTest.cpp (+20)

diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index bb99a41646b08..517d9f4dd8c41 100644
--- a/llvm/include/llvm/ADT/DenseMap.h
+++ b/llvm/include/llvm/ADT/DenseMap.h
@@ -198,6 +198,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 a4c045585fc28..81b86cff73079 100644
--- a/llvm/unittests/ADT/DenseMapTest.cpp
+++ b/llvm/unittests/ADT/DenseMapTest.cpp
@@ -570,6 +570,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 {

[fhahn]

LGTM, thanks