invalidations and eraseusers/usx95/02-12-invalidations_and_erase

3 files changed, 40 insertions, 2 deletions

diff --git a/clang/lib/Analysis/LifetimeSafety/LifetimeAnnotations.cpp b/clang/lib/Analysis/LifetimeSafety/LifetimeAnnotations.cpp
index 01666f4ac271..67d06f17ffd7 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 f1c84ff0a3f6..a5b6d1c34d82 100644
--- a/clang/test/Sema/Inputs/lifetime-analysis.h
+++ b/clang/test/Sema/Inputs/lifetime-analysis.h
@@ -90,6 +90,7 @@ struct unordered_map {
   T& operator[](const Key& key);
   iterator begin();
   iterator end();
+  iterator erase(iterator);
 };
 
 template<typename T>
diff --git a/clang/test/Sema/warn-lifetime-safety-invalidations.cpp b/clang/test/Sema/warn-lifetime-safety-invalidations.cpp
index d90c952e26e5..d23894e68fcb 100644
--- a/clang/test/Sema/warn-lifetime-safety-invalidations.cpp
+++ b/clang/test/Sema/warn-lifetime-safety-invalidations.cpp
@@ -148,6 +148,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}}