[clang] [LifetimeSafety] Add UseFacts for function arguments and assignment RHS (PR #180446)

[Utkarsh Saxena via cfe-commits]

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

>From 49026a72270c8ca815098279bd38d609168cfa62 Mon Sep 17 00:00:00 2001
From: Utkarsh Saxena <usx at google.com>
Date: Sun, 8 Feb 2026 22:29:52 +0000
Subject: [PATCH] Improve liveness to detect more invaldiations

---
 .../Analyses/LifetimeSafety/FactsGenerator.h  |  8 ++---
 .../LifetimeSafety/FactsGenerator.cpp         | 29 ++++++++++---------
 clang/test/Sema/Inputs/lifetime-analysis.h    | 16 ++++++----
 .../warn-lifetime-safety-invalidations.cpp    | 19 ++++++++++--
 4 files changed, 47 insertions(+), 25 deletions(-)

diff --git a/clang/include/clang/Analysis/Analyses/LifetimeSafety/FactsGenerator.h b/clang/include/clang/Analysis/Analyses/LifetimeSafety/FactsGenerator.h
index fb7d5ad91db79..43551e48c8ac0 100644
--- a/clang/include/clang/Analysis/Analyses/LifetimeSafety/FactsGenerator.h
+++ b/clang/include/clang/Analysis/Analyses/LifetimeSafety/FactsGenerator.h
@@ -102,10 +102,10 @@ class FactsGenerator : public ConstStmtVisitor<FactsGenerator> {
   /// If so, creates a `TestPointFact` and returns true.
   bool handleTestPoint(const CXXFunctionalCastExpr *FCE);
 
-  // A DeclRefExpr will be treated as a use of the referenced decl. It will be
+  // Treats an expression as a use of the referenced object. It will be
   // checked for use-after-free unless it is later marked as being written to
-  // (e.g. on the left-hand side of an assignment).
-  void handleUse(const DeclRefExpr *DRE);
+  // (e.g. on the left-hand side of an assignment in the case of a DeclRefExpr).
+  void handleUse(const Expr *E);
 
   void markUseAsWrite(const DeclRefExpr *DRE);
 
@@ -122,7 +122,7 @@ class FactsGenerator : public ConstStmtVisitor<FactsGenerator> {
   // `DeclRefExpr`s as "read" uses. When an assignment is processed, the use
   // corresponding to the left-hand side is updated to be a "write", thereby
   // exempting it from the check.
-  llvm::DenseMap<const DeclRefExpr *, UseFact *> UseFacts;
+  llvm::DenseMap<const Expr *, UseFact *> UseFacts;
 };
 
 } // namespace clang::lifetimes::internal
diff --git a/clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp b/clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp
index b69f69ddbae34..be62acf475fd1 100644
--- a/clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp
+++ b/clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp
@@ -366,6 +366,7 @@ void FactsGenerator::VisitBinaryOperator(const BinaryOperator *BO) {
   // result should have the same loans as the pointer operand.
   if (BO->isCompoundAssignmentOp())
     return;
+  handleUse(BO->getRHS());
   if (BO->isAssignmentOp())
     handleAssignment(BO->getLHS(), BO->getRHS());
   // TODO: Handle assignments involving dereference like `*p = q`.
@@ -582,7 +583,9 @@ void FactsGenerator::handleFunctionCall(const Expr *Call,
   FD = getDeclWithMergedLifetimeBoundAttrs(FD);
   if (!FD)
     return;
-
+  // All arguments to a function are a use of the corresponding expressions.
+  for (const Expr *Arg : Args)
+    handleUse(Arg);
   handleInvalidatingCall(Call, FD, Args);
   handleMovedArgsInCall(FD, Args);
   if (!CallList)
@@ -677,24 +680,24 @@ bool FactsGenerator::handleTestPoint(const CXXFunctionalCastExpr *FCE) {
   return false;
 }
 
-// A DeclRefExpr will be treated as a use of the referenced decl. It will be
-// checked for use-after-free unless it is later marked as being written to
-// (e.g. on the left-hand side of an assignment).
-void FactsGenerator::handleUse(const DeclRefExpr *DRE) {
-  OriginList *List = getOriginsList(*DRE);
+void FactsGenerator::handleUse(const Expr *E) {
+  OriginList *List = getOriginsList(*E);
   if (!List)
     return;
-  // Remove the outer layer of origin which borrows from the decl directly
-  // (e.g., when this is not a reference). This is a use of the underlying decl.
-  if (!DRE->getDecl()->getType()->isReferenceType())
+  // For DeclRefExpr: Remove the outer layer of origin which borrows from the
+  // decl directly (e.g., when this is not a reference). This is a use of the
+  // underlying decl.
+  if (auto *DRE = dyn_cast<DeclRefExpr>(E);
+      DRE && !DRE->getDecl()->getType()->isReferenceType())
     List = getRValueOrigins(DRE, List);
   // Skip if there is no inner origin (e.g., when it is not a pointer type).
   if (!List)
     return;
-  UseFact *UF = FactMgr.createFact<UseFact>(DRE, List);
-  CurrentBlockFacts.push_back(UF);
-  assert(!UseFacts.contains(DRE));
-  UseFacts[DRE] = UF;
+  if (!UseFacts.contains(E)) {
+    UseFact *UF = FactMgr.createFact<UseFact>(E, List);
+    CurrentBlockFacts.push_back(UF);
+    UseFacts[E] = UF;
+  }
 }
 
 void FactsGenerator::markUseAsWrite(const DeclRefExpr *DRE) {
diff --git a/clang/test/Sema/Inputs/lifetime-analysis.h b/clang/test/Sema/Inputs/lifetime-analysis.h
index f30db1a29b149..880e4650f21d0 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,17 @@ 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 find(const Key& key);
+  void 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..414541eb031c5 100644
--- a/clang/test/Sema/warn-lifetime-safety-invalidations.cpp
+++ b/clang/test/Sema/warn-lifetime-safety-invalidations.cpp
@@ -260,9 +260,22 @@ void PointerToVectorElement() {
 }
 
 void SelfInvalidatingMap() {
-  std::unordered_map<int, int> mp;
-  mp[1] = 1;
-  mp[2] = mp[1];  // FIXME: Detect this. We are mising a UseFact for the assignment params.
+  std::unordered_map<int, std::string> mp;
+  // TODO: We do not have a way to differentiate between pointer stability and iterator stability!
+  // std::unordered_map and other containers provide pointer/reference stability. Therefore the
+  // following is safe in practice.
+  // On the other hand, std::flat_hash_map (since C++23) does not provide pointer stability on
+  // insertion and following is unsafe for this container.
+  mp[1] = "42";
+  mp[2] = mp[1];  // expected-warning {{object whose reference is captured is later invalidated}} \
+                  // expected-note {{invalidated here}} \
+                  // expected-note {{later used here}}
+
+  // None of these containers provide iterator stability. So following is unsafe:
+  auto it = mp.find(3); // expected-warning {{object whose reference is captured is later invalidated}}
+  mp.erase(mp.find(4)); // expected-note {{invalidated here}} 
+  if (it != mp.end())   // expected-note {{later used here}}
+    *it;
 }
 } // namespace ElementReferences