[llvm,clang] Don't assume non-erased DenseMap entries remain valid after erase. NFC (llvm#198982)

In preparation for switching DenseMap from tombstone deletion to
backward-shift deletion, update call sites that reuse an iterator or a
bucket reference after erasing another entry from the same map.

These work under tombstone deletion because unrelated buckets stay put,
but backward-shift deletion relocates entries to close the gap.

Add DenseMap::remove_if, similar to SmallPtrSet::remove_if, as
replacement for erase-while-iterating, and use it where applicable.

Aided by Claude Opus 4.7

Author: MaskRay

clang/lib/AST/ASTImporter.cpp

@@ -9864,10 +9864,15 @@ Expected<Decl *> ASTImporter::Import(Decl *FromD) {
     // Failed to import.
 
     auto Pos = ImportedDecls.find(FromD);
-    if (Pos != ImportedDecls.end()) {
+    bool ToDWasCreated = Pos != ImportedDecls.end();
+    // Capture the mapped decl before erasing: the iterator is invalidated by
+    // the erase below under backward-shift deletion, but it is still needed
+    // further down to record the import error.
+    Decl *CreatedToD = ToDWasCreated ? Pos->second : nullptr;
+    if (ToDWasCreated) {
       // Import failed after the object was created.
       // Remove all references to it.
-      auto *ToD = Pos->second;
+      auto *ToD = CreatedToD;
       ImportedDecls.erase(Pos);
 
       // ImportedDecls and ImportedFromDecls are not symmetric.  It may happen
@@ -9903,8 +9908,8 @@ Expected<Decl *> ASTImporter::Import(Decl *FromD) {
                     [&ErrOut](const ASTImportError &E) { ErrOut = E; });
     setImportDeclError(FromD, ErrOut);
     // Set the error for the mapped to Decl, which is in the "to" context.
-    if (Pos != ImportedDecls.end())
-      SharedState->setImportDeclError(Pos->second, ErrOut);
+    if (ToDWasCreated)
+      SharedState->setImportDeclError(CreatedToD, ErrOut);
 
     // Set the error for all nodes which have been created before we
     // recognized the error.

clang/lib/CodeGen/CoverageMappingGen.cpp

@@ -2260,13 +2260,9 @@ struct CounterCoverageMappingBuilder
     (void)FoundCount;
 
     // Tell CodeGenPGO not to instrument.
-    for (auto I = MCDCState.BranchByStmt.begin(),
-              E = MCDCState.BranchByStmt.end();
-         I != E;) {
-      auto II = I++;
-      if (II->second.DecisionStmt == Decision)
-        MCDCState.BranchByStmt.erase(II);
-    }
+    MCDCState.BranchByStmt.remove_if([&](const auto &Entry) {
+      return Entry.second.DecisionStmt == Decision;
+    });
     MCDCState.DecisionByStmt.erase(Decision);
   }
 

clang/lib/Interpreter/IncrementalParser.cpp

@@ -174,6 +174,9 @@ IncrementalParser::Parse(llvm::StringRef input) {
 
 void IncrementalParser::CleanUpPTU(TranslationUnitDecl *MostRecentTU) {
   if (StoredDeclsMap *Map = MostRecentTU->getPrimaryContext()->getLookupPtr()) {
+    // Collect the keys to erase: erasing during iteration invalidates the map
+    // iterator under backward-shift deletion.
+    llvm::SmallVector<DeclarationName, 16> KeysToErase;
     for (auto &&[Key, List] : *Map) {
       DeclContextLookupResult R = List.getLookupResult();
       std::vector<NamedDecl *> NamedDeclsToRemove;
@@ -185,12 +188,14 @@ void IncrementalParser::CleanUpPTU(TranslationUnitDecl *MostRecentTU) {
           RemoveAll = false;
       }
       if (LLVM_LIKELY(RemoveAll)) {
-        Map->erase(Key);
+        KeysToErase.push_back(Key);
       } else {
         for (NamedDecl *D : NamedDeclsToRemove)
           List.remove(D);
       }
     }
+    for (DeclarationName Key : KeysToErase)
+      Map->erase(Key);
   }
 
   ExternCContextDecl *ECCD = S.getASTContext().getExternCContextDecl();

clang/lib/Sema/HLSLExternalSemaSource.cpp

@@ -726,6 +726,10 @@ void HLSLExternalSemaSource::CompleteType(TagDecl *Tag) {
   auto It = Completions.find(Record);
   if (It == Completions.end())
     return;
-  It->second(Record);
+  // Move out the callback and erase before invoking it: the callback can
+  // re-enter CompleteType and mutate Completions, which invalidates It under
+  // backward-shift deletion.
+  CompletionFunction Fn = std::move(It->second);
   Completions.erase(It);
+  Fn(Record);
 }

llvm/include/llvm/ADT/DenseMap.h

@@ -344,6 +344,33 @@ class DenseMapBase : public DebugEpochBase {
     incrementNumTombstones();
   }
 
+  /// Remove entries that match the given predicate. \p Pred is invoked
+  /// with a reference to each live bucket and must not access the map being
+  /// modified. This is the safe replacement for erase-while-iterating.
+  ///
+  /// Returns whether anything was removed. If so, all iterators and references
+  /// into the map are invalidated.
+  template <typename Predicate> bool remove_if(Predicate Pred) {
+    const KeyT EmptyKey = KeyInfoT::getEmptyKey();
+    const KeyT TombstoneKey = KeyInfoT::getTombstoneKey();
+    bool Removed = false;
+    for (BucketT &B : buckets()) {
+      if (KeyInfoT::isEqual(B.getFirst(), EmptyKey) ||
+          KeyInfoT::isEqual(B.getFirst(), TombstoneKey))
+        continue;
+      if (Pred(B)) {
+        B.getSecond().~ValueT();
+        B.getFirst() = TombstoneKey;
+        decrementNumEntries();
+        incrementNumTombstones();
+        Removed = true;
+      }
+    }
+    if (Removed)
+      incrementEpoch();
+    return Removed;
+  }
+
   ValueT &operator[](const KeyT &Key) {
     return lookupOrInsertIntoBucket(Key).first->second;
   }

llvm/include/llvm/ADT/DenseSet.h

@@ -99,6 +99,16 @@ class DenseSetImpl {
 
   bool erase(const ValueT &V) { return TheMap.erase(V); }
 
+  /// Remove all elements for which \p Pred returns true.  This is the safe
+  /// replacement for erase-while-iterating; see DenseMap::remove_if.  The
+  /// predicate must not access the set being modified.  Returns whether
+  /// anything was removed; if so, all iterators are invalidated.
+  template <typename Predicate> bool remove_if(Predicate Pred) {
+    return TheMap.remove_if([&](const typename MapTy::value_type &KV) {
+      return Pred(KV.getFirst());
+    });
+  }
+
   void swap(DenseSetImpl &RHS) { TheMap.swap(RHS.TheMap); }
 
 private:

llvm/lib/Analysis/IRSimilarityIdentifier.cpp

@@ -720,7 +720,12 @@ bool IRSimilarityCandidate::compareAssignmentMapping(
   if (!WasInserted && !ValueMappingIt->second.contains(InstValB))
     return false;
   else if (ValueMappingIt->second.size() != 1) {
-    for (unsigned OtherVal : ValueMappingIt->second) {
+    // Snapshot the set before iterating: when InstValA maps to itself the
+    // erase below removes InstValA from the very set being iterated, which
+    // invalidates the range iterator under backward-shift deletion.
+    SmallVector<unsigned> OtherVals(ValueMappingIt->second.begin(),
+                                    ValueMappingIt->second.end());
+    for (unsigned OtherVal : OtherVals) {
       if (OtherVal == InstValB)
         continue;
       auto OtherValIt = ValueNumberMappingA.find(OtherVal);

llvm/lib/Analysis/LoopAccessAnalysis.cpp

@@ -3229,12 +3229,11 @@ void LoopAccessInfoManager::clear() {
   // analyzed loop or SCEVs that may have been modified or invalidated. At the
   // moment, that is loops requiring memory or SCEV runtime checks, as those cache
   // SCEVs, e.g. for pointer expressions.
-  for (const auto &[L, LAI] : LoopAccessInfoMap) {
-    if (LAI->getRuntimePointerChecking()->getChecks().empty() &&
-        LAI->getPSE().getPredicate().isAlwaysTrue())
-      continue;
-    LoopAccessInfoMap.erase(L);
-  }
+  LoopAccessInfoMap.remove_if([](const auto &Entry) {
+    const auto &LAI = Entry.second;
+    return !(LAI->getRuntimePointerChecking()->getChecks().empty() &&
+             LAI->getPSE().getPredicate().isAlwaysTrue());
+  });
 }
 
 bool LoopAccessInfoManager::invalidate(

llvm/lib/Analysis/ScalarEvolution.cpp

@@ -8763,8 +8763,8 @@ void ScalarEvolution::visitAndClearUsers(
     ValueExprMapType::iterator It =
         ValueExprMap.find_as(static_cast<Value *>(I));
     if (It != ValueExprMap.end()) {
-      eraseValueFromMap(It->first);
       ToForget.push_back(It->second);
+      eraseValueFromMap(It->first);
       if (PHINode *PN = dyn_cast<PHINode>(I))
         ConstantEvolutionLoopExitValue.erase(PN);
     }
@@ -8788,14 +8788,8 @@ void ScalarEvolution::forgetLoop(const Loop *L) {
     forgetBackedgeTakenCounts(CurrL, /* Predicated */ true);
 
     // Drop information about predicated SCEV rewrites for this loop.
-    for (auto I = PredicatedSCEVRewrites.begin();
-         I != PredicatedSCEVRewrites.end();) {
-      std::pair<const SCEV *, const Loop *> Entry = I->first;
-      if (Entry.second == CurrL)
-        PredicatedSCEVRewrites.erase(I++);
-      else
-        ++I;
-    }
+    PredicatedSCEVRewrites.remove_if(
+        [&](const auto &Entry) { return Entry.first.second == CurrL; });
 
     auto LoopUsersItr = LoopUsers.find(CurrL);
     if (LoopUsersItr != LoopUsers.end())
@@ -14581,14 +14575,8 @@ void ScalarEvolution::forgetMemoizedResults(ArrayRef<SCEVUse> SCEVs) {
   for (const auto *S : ToForget)
     forgetMemoizedResultsImpl(S);
 
-  for (auto I = PredicatedSCEVRewrites.begin();
-       I != PredicatedSCEVRewrites.end();) {
-    std::pair<const SCEV *, const Loop *> Entry = I->first;
-    if (ToForget.count(Entry.first))
-      PredicatedSCEVRewrites.erase(I++);
-    else
-      ++I;
-  }
+  PredicatedSCEVRewrites.remove_if(
+      [&](const auto &Entry) { return ToForget.count(Entry.first.first); });
 }
 
 void ScalarEvolution::forgetMemoizedResultsImpl(const SCEV *S) {

llvm/lib/CodeGen/AssignmentTrackingAnalysis.cpp

@@ -543,18 +543,15 @@ class MemLocFragmentFill {
     // Meet A and B.
     //
     // Result = meet(a, b) for a in A, b in B where Var(a) == Var(b)
-    for (auto It = A.begin(), End = A.end(); It != End; ++It) {
-      unsigned AVar = It->first;
-      FragsInMemMap &AFrags = It->second;
-      auto BIt = B.find(AVar);
-      if (BIt == B.end()) {
-        A.erase(It);
-        continue; // Var has no bits defined in B.
-      }
+    A.remove_if([&](auto &Entry) {
+      auto BIt = B.find(Entry.first);
+      if (BIt == B.end())
+        return true; // Var has no bits defined in B.
       LLVM_DEBUG(dbgs() << "meet fragment maps for "
-                        << Aggregates[AVar].first->getName() << "\n");
-      AFrags = meetFragments(AFrags, BIt->second);
-    }
+                        << Aggregates[Entry.first].first->getName() << "\n");
+      Entry.second = meetFragments(Entry.second, BIt->second);
+      return false;
+    });
   }
 
   bool meet(const BasicBlock &BB,