ADT: Complete the at() methods for DenseMap and MapVector

[nhaehnle]

Make it easier to use these containers as drop-in replacements for
std::map.

[llvmbot]

@llvm/pr-subscribers-llvm-adt

Author: Nicolai Hähnle (nhaehnle)

Changes

Make it easier to use these containers as drop-in replacements for
std::map.

---

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

4 Files Affected:

• (modified) llvm/include/llvm/ADT/DenseMap.h (+8)
• (modified) llvm/include/llvm/ADT/MapVector.h (+18-4)
• (modified) llvm/unittests/ADT/DenseMapTest.cpp (+8)
• (modified) llvm/unittests/ADT/MapVectorTest.cpp (+15)

diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h index a706f68fab81b..fe8868619730e 100644 --- a/llvm/include/llvm/ADT/DenseMap.h +++ b/llvm/include/llvm/ADT/DenseMap.h @@ -219,6 +219,14 @@ class DenseMapBase : public DebugEpochBase { return Default; } + /// at - Return the entry for the specified key, or abort if no such + /// entry exists. + [[nodiscard]] ValueT &at(const_arg_type_t<KeyT> Val) { + auto Iter = this->find(std::move(Val)); + assert(Iter != this->end() && "DenseMap::at failed due to a missing key"); + return Iter->second; + } + /// at - Return the entry for the specified key, or abort if no such /// entry exists. [[nodiscard]] const ValueT &at(const_arg_type_t<KeyT> Val) const { diff --git a/llvm/include/llvm/ADT/MapVector.h b/llvm/include/llvm/ADT/MapVector.h index 82f2c4977e01d..80bcb7e0b7ba4 100644 --- a/llvm/include/llvm/ADT/MapVector.h +++ b/llvm/include/llvm/ADT/MapVector.h @@ -148,14 +148,28 @@ class MapVector { [[nodiscard]] iterator find(const KeyT &Key) { typename MapType::const_iterator Pos = Map.find(Key); - return Pos == Map.end()? Vector.end() : - (Vector.begin() + Pos->second); + return Pos == Map.end() ? Vector.end() : (Vector.begin() + Pos->second); } [[nodiscard]] const_iterator find(const KeyT &Key) const { typename MapType::const_iterator Pos = Map.find(Key); - return Pos == Map.end()? Vector.end() : - (Vector.begin() + Pos->second); + return Pos == Map.end() ? Vector.end() : (Vector.begin() + Pos->second); + } + + /// at - Return the entry for the specified key, or abort if no such + /// entry exists. + [[nodiscard]] ValueT &at(const KeyT &Key) { + auto Pos = Map.find(Key); + assert(Pos != Map.end() && "MapVector::at failed due to a missing key"); + return Vector[Pos->second].second; + } + + /// at - Return the entry for the specified key, or abort if no such + /// entry exists. + [[nodiscard]] const ValueT &at(const KeyT &Key) const { + auto Pos = Map.find(Key); + assert(Pos != Map.end() && "MapVector::at failed due to a missing key"); + return Vector[Pos->second].second; } /// Remove the last element from the vector. diff --git a/llvm/unittests/ADT/DenseMapTest.cpp b/llvm/unittests/ADT/DenseMapTest.cpp index 273ee09fc1e28..553d159d33b1a 100644 --- a/llvm/unittests/ADT/DenseMapTest.cpp +++ b/llvm/unittests/ADT/DenseMapTest.cpp @@ -200,6 +200,14 @@ TYPED_TEST(DenseMapTest, AtTest) { EXPECT_EQ(this->getValue(0), this->Map.at(this->getKey(0))); EXPECT_EQ(this->getValue(1), this->Map.at(this->getKey(1))); EXPECT_EQ(this->getValue(2), this->Map.at(this->getKey(2))); + + this->Map.at(this->getKey(0)) = this->getValue(1); + EXPECT_EQ(this->getValue(1), this->Map.at(this->getKey(0))); + + const auto &ConstMap = this->Map; + EXPECT_EQ(this->getValue(1), ConstMap.at(this->getKey(0))); + EXPECT_EQ(this->getValue(1), ConstMap.at(this->getKey(1))); + EXPECT_EQ(this->getValue(2), ConstMap.at(this->getKey(2))); } // Test clear() method diff --git a/llvm/unittests/ADT/MapVectorTest.cpp b/llvm/unittests/ADT/MapVectorTest.cpp index 639e1a14e8178..e0589445e3271 100644 --- a/llvm/unittests/ADT/MapVectorTest.cpp +++ b/llvm/unittests/ADT/MapVectorTest.cpp @@ -292,6 +292,21 @@ TEST(MapVectorTest, GetArrayRef) { EXPECT_TRUE(MV.getArrayRef().equals({std::pair(100, 99), std::pair(99, 98)})); } +TEST(MapVectorTest, AtTest) { + MapVector<int, int> MV; + MV[0] = 10; + MV[1] = 11; + EXPECT_EQ(MV.at(0), 10); + EXPECT_EQ(MV.at(1), 11); + + MV.at(1) = 12; + EXPECT_EQ(MV.at(1), 12); + + const auto &ConstMV = MV; + EXPECT_EQ(ConstMV.at(0), 10); + EXPECT_EQ(ConstMV.at(1), 12); +} + template <class IntType> struct MapVectorMappedTypeTest : ::testing::Test { using int_type = IntType; };