[llvm] ADT: Complete the at() methods for DenseMap and MapVector (PR #169147)

[Nicolai Hähnle via llvm-commits]

https://github.com/nhaehnle created https://github.com/llvm/llvm-project/pull/169147

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

>From f6fcb2a75ddac2bf8cec2169dd10e1e8e5a1eb40 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicolai=20H=C3=A4hnle?= <nicolai.haehnle at amd.com>
Date: Thu, 20 Nov 2025 09:21:38 -0800
Subject: [PATCH] ADT: Complete the at() methods for DenseMap and MapVector

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

commit-id:8b1d1826
---
 llvm/include/llvm/ADT/DenseMap.h     |  8 ++++++++
 llvm/include/llvm/ADT/MapVector.h    | 22 ++++++++++++++++++----
 llvm/unittests/ADT/DenseMapTest.cpp  |  8 ++++++++
 llvm/unittests/ADT/MapVectorTest.cpp | 15 +++++++++++++++
 4 files changed, 49 insertions(+), 4 deletions(-)

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