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

[StrongerXi]

Summary

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);

[lattner]

Very cool, I'm excited to get this. This will simplify a lot of code! Could you also add this to StringMap? It has a similar lookup method but it can't be used for non-default-constructible values, and has the same usage pattern as this container.

llvm/include/llvm/ADT/DenseMap.h

211

Unfortunately, I don't think we can afford to include this level of debug tracing in here. This adds a bunch of #include dependencies, and (more problematically) includes a bunch of code in the non assert build.

The standard LLVM way to do this (see in things like ArrayRef's subscript) is to just assert with a fixed message "lookupOrTrap failed due to a missing key", using the standard assert macro.

[StrongerXi]

1. Use assert and fixed error message
2. Add the same method and unittests to StringMap

[StrongerXi]

llvm/include/llvm/ADT/DenseMap.h

211

Ah good points, updated.

llvm/include/llvm/ADT/StringMap.h

234–237

Updated because I noticed they didn't follow the convention of capitalized variable names.

[foad]

With apologies for bikeshedding:

I don't love the "OrTrap" name. It might trap in a Debug build, but it surely won't in a Release build. Anyway I don't think the name needs to say what the consequences might be if the key does not exist.

How about calling it just "lookup", and using a longer name like "lookupOrDefaultConstruct" for the less-obvious behaviour of the current lookup method? (Is there really much code that relies on its "or default construct" behaviour? And anyway aren't there more STL-like ways of getting that behaviour, like using operator[]?)

Or why not call it "at"?

[lattner]

Bikeshedding is good @foad, these are important core data structures :)

I agree with you that at seems like the right name for this, given the STL precedent. I wasn't aware of that existing for maps.

Renaming lookup -> lookupOrDefaultConstruct() or getValueIfPresent() or something like that would make sense as well, but seems like a separate issue.

[StrongerXi]

I agree that at has the advantage of being similar (in spelling & semantics) to std::map::at. Possible counterarguments:

1. at too implicit about the "assertive/validated" semantics.
2. having something similar to lookup is a easier to adapt for existing users, also makes the API more consistent/intuitive.

I can't seem to think of many alternatives though, lookupChecked? What do you all think?

I also agree that renaming the old lookup should be a separate issue, and I think it has backward compatibility concerns too.

[lattner]

We generally don't worry to much about compatibility issues. It is good to deprecate the old name for a few weeks before removing it, but we don't do long cycles here.

I also don't like the 'at' name particularly, but it is highly precedented by the STL. I think any alternate name would have to be really good to overcome that.

[StrongerXi]

Rename lookupOrTrap to at

[StrongerXi]

In D143976#4127948, @lattner wrote:

We generally don't worry to much about compatibility issues. It is good to deprecate the old name for a few weeks before removing it, but we don't do long cycles here.

I also don't like the 'at' name particularly, but it is highly precedented by the STL. I think any alternate name would have to be really good to overcome that.

Sounds good; renaming to at:).

[foad]

Could you also add this to StringMap?

Maybe IntervalMap and MapVector too?

I agree with you that at seems like the right name for this, given the STL precedent. I wasn't aware of that existing for maps.

Agreed.

Renaming lookup -> lookupOrDefaultConstruct() or getValueIfPresent() or something like that would make sense as well, but seems like a separate issue.

I only suggested that to free up the name "lookup". I don't see any pressing need to change it if we are going with "at" for the new function.

llvm/include/llvm/ADT/DenseMap.h

206

Just a question: would it be better to return const ValueT & like std::map::at would? Do we also want a non-const version of this?

[StrongerXi]

Could you also add this to StringMap?

Maybe IntervalMap and MapVector too?

I was thinking about this, didn't want to go too aggressive. But it should be easy to add once we reach consensus on the API here:).

llvm/include/llvm/ADT/DenseMap.h

206

Good question, I'm following the signature of existing lookup methods, I _guess_ the reason is more flexible API?

[lattner]

Awesome, thank you for driving this!

llvm/include/llvm/ADT/DenseMap.h

206

+1, no need to copy non-trivial values.

llvm/include/llvm/ADT/StringMap.h

242

Similarly should return const&

[StrongerXi]

Address all feedbacks:

1. Return const reference

[StrongerXi]

git clang-format

llvm/include/llvm/ADT/DenseMap.h

public:
   /// constructed value if no such entry exists.
   ValueT lookup(const_arg_type_t<KeyT> Val) const {
     const BucketT *TheBucket;
     if (LookupBucketFor(Val, TheBucket))
       return TheBucket->getSecond();
     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.
   std::pair<iterator, bool> insert(const std::pair<KeyT, ValueT> &KV) {
     return try_emplace(KV.first, KV.second);
   }
 
   // Inserts key,value pair into the map if the key isn't already in the map.

llvm/include/llvm/ADT/StringMap.h

const_iterator find(StringRef Key) const {
     if (Bucket == -1)
       return end();
     return const_iterator(TheTable + Bucket, true);
   }
 
   /// 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; }
 
   /// count - Return 1 if the element is in the map, 0 otherwise.
   size_type count(StringRef Key) const { return find(Key) == end() ? 0 : 1; }
 
   template <typename InputTy>

llvm/unittests/ADT/DenseMapTest.cpp

TYPED_TEST(DenseMapTest, EmptyIntMapTest) {
   // Iterator tests
   EXPECT_TRUE(this->Map.begin() == this->Map.end());
 
   // Lookup tests
   EXPECT_FALSE(this->Map.count(this->getKey()));
   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
 TYPED_TEST(DenseMapTest, ConstEmptyMapTest) {
   const TypeParam &ConstMap = this->Map;
   EXPECT_EQ(0u, ConstMap.size());
   EXPECT_TRUE(ConstMap.empty());
   EXPECT_TRUE(ConstMap.begin() == ConstMap.end());
TYPED_TEST(DenseMapTest, SingleEntryMapTest) {
   ++it;
   EXPECT_TRUE(it == this->Map.end());
 
   // Lookup tests
   EXPECT_TRUE(this->Map.count(this->getKey()));
   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
 TYPED_TEST(DenseMapTest, ClearTest) {
   this->Map[this->getKey()] = this->getValue();
   this->Map.clear();
 
   EXPECT_EQ(0u, this->Map.size());

llvm/unittests/ADT/StringMapTest.cpp

TEST_F(StringMapTest, CopyCtorTest) {
   EXPECT_EQ(3u, Map2.size());
   EXPECT_EQ(0, Map2.lookup("eins"));
   EXPECT_EQ(2, Map2.lookup("zwei"));
   EXPECT_EQ(0, Map2.lookup("drei"));
   EXPECT_EQ(4, Map2.lookup("veir"));
   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];
 
   // Insert 100 numbers into the map
   for (int i = 0; i < 100; ++i) {
     std::stringstream ss;
     ss << "key_" << i;