diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h --- a/llvm/include/llvm/ADT/DenseMap.h +++ b/llvm/include/llvm/ADT/DenseMap.h @@ -201,6 +201,14 @@ return ValueT(); } + /// at - Return the entry for the specified key, or abort if no such + /// entry exists. + ValueT at(const_arg_type_t 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 --- a/llvm/include/llvm/ADT/StringMap.h +++ b/llvm/include/llvm/ADT/StringMap.h @@ -231,12 +231,20 @@ /// 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. + 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 --- a/llvm/unittests/ADT/DenseMapTest.cpp +++ b/llvm/unittests/ADT/DenseMapTest.cpp @@ -125,6 +125,10 @@ 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 @@ 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 --- a/llvm/unittests/ADT/StringMapTest.cpp +++ b/llvm/unittests/ADT/StringMapTest.cpp @@ -205,6 +205,22 @@ EXPECT_EQ(5, Map2.lookup("funf")); } +TEST_F(StringMapTest, LookupOrTrapTest) { + llvm::StringMap 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];