[clang] [LifetimeSafety] 'erase' does not invaldiate node-based containers (PR #181216)

[Utkarsh Saxena via cfe-commits]

https://github.com/usx95 updated https://github.com/llvm/llvm-project/pull/181216

>From 342377ae219913e063ce1aa5a4fac3604c63fdc0 Mon Sep 17 00:00:00 2001
From: Utkarsh Saxena <usx at google.com>
Date: Thu, 12 Feb 2026 19:28:17 +0000
Subject: [PATCH] invalidations and erase

---
 .../LifetimeSafety/LifetimeAnnotations.cpp    | 26 +++++++++++++++++--
 clang/test/Sema/Inputs/lifetime-analysis.h    | 14 ++++++----
 .../warn-lifetime-safety-invalidations.cpp    | 15 +++++++++++
 3 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/clang/lib/Analysis/LifetimeSafety/LifetimeAnnotations.cpp b/clang/lib/Analysis/LifetimeSafety/LifetimeAnnotations.cpp
index 01666f4ac271c..67d06f17ffd74 100644
--- a/clang/lib/Analysis/LifetimeSafety/LifetimeAnnotations.cpp
+++ b/clang/lib/Analysis/LifetimeSafety/LifetimeAnnotations.cpp
@@ -317,12 +317,33 @@ bool isContainerInvalidationMethod(const CXXMethodDecl &MD) {
 
   if (!MD.getIdentifier())
     return false;
+
+  StringRef MethodName = MD.getName();
+
+  // Special handling for 'erase':
+  // It invalidates the whole container (effectively) for contiguous/flat
+  // storage, but is safe for other iterators in node-based containers.
+  if (MethodName == "erase") {
+    static const llvm::StringSet<> NodeBasedContainers = {"map",
+                                                          "set",
+                                                          "multimap",
+                                                          "multiset",
+                                                          "unordered_map",
+                                                          "unordered_set",
+                                                          "unordered_multimap",
+                                                          "unordered_multiset"};
+
+    // 'erase' invalidates for non node-based containers (vector, deque, string,
+    // flat_map).
+    return !NodeBasedContainers.contains(ContainerName);
+  }
+
   static const llvm::StringSet<> InvalidatingMembers = {
       // Basic Insertion/Emplacement
       "push_front", "push_back", "emplace_front", "emplace_back", "insert",
       "emplace", "push",
       // Basic Removal/Clearing
-      "pop_front", "pop_back", "pop", "erase", "clear",
+      "pop_front", "pop_back", "pop", "clear",
       // Memory Management
       "reserve", "resize", "shrink_to_fit",
       // Assignment (Named)
@@ -331,6 +352,7 @@ bool isContainerInvalidationMethod(const CXXMethodDecl &MD) {
       "append", "replace",
       // Modern C++ (C++17/23)
       "extract", "try_emplace", "insert_range", "append_range", "assign_range"};
-  return InvalidatingMembers.contains(MD.getName());
+
+  return InvalidatingMembers.contains(MethodName);
 }
 } // namespace clang::lifetimes
diff --git a/clang/test/Sema/Inputs/lifetime-analysis.h b/clang/test/Sema/Inputs/lifetime-analysis.h
index f30db1a29b149..5946cee49b5dc 100644
--- a/clang/test/Sema/Inputs/lifetime-analysis.h
+++ b/clang/test/Sema/Inputs/lifetime-analysis.h
@@ -75,11 +75,6 @@ struct vector {
   void clear();
 };
 
-template<class Key,class T>
-struct unordered_map {
-  T& operator[](const Key& key);
-};
-
 template<class T>
 void swap( T& a, T& b );
 
@@ -89,6 +84,15 @@ struct pair {
   B second;
 };
 
+template<class Key,class T>
+struct unordered_map {
+  using iterator = __gnu_cxx::basic_iterator<std::pair<const Key, T>>;
+  T& operator[](const Key& key);
+  iterator begin();
+  iterator end();
+  iterator erase(iterator);
+};
+
 template<typename T>
 struct basic_string_view {
   basic_string_view();
diff --git a/clang/test/Sema/warn-lifetime-safety-invalidations.cpp b/clang/test/Sema/warn-lifetime-safety-invalidations.cpp
index c9ce0c35c53d2..cc48b3d034468 100644
--- a/clang/test/Sema/warn-lifetime-safety-invalidations.cpp
+++ b/clang/test/Sema/warn-lifetime-safety-invalidations.cpp
@@ -138,6 +138,21 @@ void IteratorInvalidationInAWhileLoop(std::vector<int> v) {
   }
 }
 
+void NoIteratorInvalidationInAWhileLoopErase(std::unordered_map<int, int> mp) {
+  auto it = mp.begin();
+  while (it != std::end(mp)) {
+    if (Bool()) {
+      auto next = it;
+      ++next;
+      mp.erase(it); // Ok. 'next' remains valid.
+      it = next;
+    }
+    else {
+      ++it;
+    }
+  }
+}
+
 void IteratorInvalidationInAForeachLoop(std::vector<int> v) {
   for (int& x : v) { // expected-warning {{object whose reference is captured is later invalidated}} \
                      // expected-note {{later used here}}