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

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 {

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.

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

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