[llvm] 1320036 - [ADT] Add `at` method (assertive lookup) to DenseMap and StringMap

[Ryan Guo via llvm-commits]

Author: Ryan Guo
Date: 2023-02-17T08:32:29-08:00
New Revision: 132003603ae3453bc385ffd5ed53f5e8057ae1bc

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

LOG: [ADT] Add `at` method (assertive lookup) to DenseMap and StringMap

This patch makes it easier for users when they want to use validated
lookup on DenseMap/StringMap as a composable C++ expression. For
instance:

```
// instead of
if (auto val = map.lookup(key))
   return val;
assert("...");

// we can write
return map.at(key);
```

Differential Revision: https://reviews.llvm.org/D143976

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h
index 9da6f479c427..12a6673cc62c 100644
--- a/llvm/include/llvm/ADT/DenseMap.h
+++ b/llvm/include/llvm/ADT/DenseMap.h
@@ -201,6 +201,14 @@ class DenseMapBase : public DebugEpochBase {
     return ValueT();
   }
 
+  /// 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 {
+    auto Iter = this->find(std::move(Val));
+    assert(Iter != this->end() && "DenseMap::at failed due to a missing key");
+    return Iter->second;
+  }
+
   // Inserts key,value pair into the map if the key isn't already in the map.
   // If the key is already in the map, it returns false and doesn't update the
   // value.

diff  --git a/llvm/include/llvm/ADT/StringMap.h b/llvm/include/llvm/ADT/StringMap.h
index 0849bef53ba1..443714c02389 100644
--- a/llvm/include/llvm/ADT/StringMap.h
+++ b/llvm/include/llvm/ADT/StringMap.h
@@ -231,12 +231,20 @@ class StringMap : public StringMapImpl,
   /// lookup - Return the entry for the specified key, or a default
   /// constructed value if no such entry exists.
   ValueTy lookup(StringRef Key) const {
-    const_iterator it = find(Key);
-    if (it != end())
-      return it->second;
+    const_iterator Iter = find(Key);
+    if (Iter != end())
+      return Iter->second;
     return ValueTy();
   }
 
+  /// at - Return the entry for the specified key, or abort if no such
+  /// entry exists.
+  const ValueTy &at(StringRef Val) const {
+    auto Iter = this->find(std::move(Val));
+    assert(Iter != this->end() && "StringMap::at failed due to a missing key");
+    return Iter->second;
+  }
+
   /// Lookup the ValueTy for the \p Key, or create a default constructed value
   /// if the key is not in the map.
   ValueTy &operator[](StringRef Key) { return try_emplace(Key).first->second; }

diff  --git a/llvm/unittests/ADT/DenseMapTest.cpp b/llvm/unittests/ADT/DenseMapTest.cpp
index 2d01316e658b..ba4e76424f04 100644
--- a/llvm/unittests/ADT/DenseMapTest.cpp
+++ b/llvm/unittests/ADT/DenseMapTest.cpp
@@ -125,6 +125,10 @@ TYPED_TEST(DenseMapTest, EmptyIntMapTest) {
   EXPECT_TRUE(this->Map.find(this->getKey()) == this->Map.end());
   EXPECT_EQ(typename TypeParam::mapped_type(),
             this->Map.lookup(this->getKey()));
+
+  // LookupOrTrap tests
+  EXPECT_DEATH({ this->Map.at(this->getKey()); },
+               "DenseMap::at failed due to a missing key");
 }
 
 // Constant map tests
@@ -156,6 +160,10 @@ TYPED_TEST(DenseMapTest, SingleEntryMapTest) {
   EXPECT_TRUE(this->Map.find(this->getKey()) == this->Map.begin());
   EXPECT_EQ(this->getValue(), this->Map.lookup(this->getKey()));
   EXPECT_EQ(this->getValue(), this->Map[this->getKey()]);
+
+  // LookupOrTrap tests
+  EXPECT_DEATH({ this->Map.at(this->getKey(1)); },
+               "DenseMap::at failed due to a missing key");
 }
 
 // Test clear() method

diff  --git a/llvm/unittests/ADT/StringMapTest.cpp b/llvm/unittests/ADT/StringMapTest.cpp
index f40f22ad140c..25562d366c42 100644
--- a/llvm/unittests/ADT/StringMapTest.cpp
+++ b/llvm/unittests/ADT/StringMapTest.cpp
@@ -205,6 +205,22 @@ TEST_F(StringMapTest, CopyCtorTest) {
   EXPECT_EQ(5, Map2.lookup("funf"));
 }
 
+TEST_F(StringMapTest, LookupOrTrapTest) {
+  llvm::StringMap<int> Map;
+
+  // key not found on empty map
+  EXPECT_DEATH({ Map.at("a"); }, "StringMap::at failed due to a missing key");
+
+  // keys both found and not found on non-empty map
+  Map["a"] = 1;
+  Map["b"] = 2;
+  Map["c"] = 3;
+  EXPECT_EQ(1, Map.at("a"));
+  EXPECT_EQ(2, Map.at("b"));
+  EXPECT_EQ(3, Map.at("c"));
+  EXPECT_DEATH({ Map.at("d"); }, "StringMap::at failed due to a missing key");
+}
+
 // A more complex iteration test.
 TEST_F(StringMapTest, IterationTest) {
   bool visited[100];