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 @@ -16,8 +16,10 @@ #include "llvm/ADT/DenseMapInfo.h" #include "llvm/ADT/EpochTracker.h" +#include "llvm/ADT/Twine.h" #include "llvm/Support/AlignOf.h" #include "llvm/Support/Compiler.h" +#include "llvm/Support/Debug.h" #include "llvm/Support/MathExtras.h" #include "llvm/Support/MemAlloc.h" #include "llvm/Support/ReverseIteration.h" @@ -201,6 +203,16 @@ return ValueT(); } + /// lookupOrTrap - Return the entry for the specified key, or abort with given + /// message if no such entry exists. + ValueT lookupOrTrap(const_arg_type_t Val, const Twine &Msg = "") const { + auto Iter = this->find(std::move(Val)); + if (Iter != this->end()) + return Iter->second; + dbgs() << "[DenseMap::lookupOrTrap] key not found; " << Msg << "\n"; + std::abort(); + } + // 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/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,13 @@ 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.lookupOrTrap(this->getKey()); }, + "\\[DenseMap::lookupOrTrap\\] key not found;"); + EXPECT_DEATH( + { this->Map.lookupOrTrap(this->getKey(), "clarification message 1"); }, + "\\[DenseMap::lookupOrTrap\\] key not found; clarification message 1"); } // Constant map tests @@ -156,6 +163,13 @@ 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.lookupOrTrap(this->getKey(1)); }, + "\\[DenseMap::lookupOrTrap\\] key not found;"); + EXPECT_DEATH( + { this->Map.lookupOrTrap(this->getKey(1), "clarification message 2"); }, + "\\[DenseMap::lookupOrTrap\\] key not found; clarification message 2"); } // Test clear() method