[LifetimeSafety] Track origins through std::function (llvm#191123)

aeft · web-flow · commit 1b804e20b9d2 · 2026-04-14T15:01:35.000+05:30
1. Recognizes `std::function` and `std::move_only_function` as types that can carry origins from a wrapped lambda's captures, propagating origins through both construction and assignment. 2. Adds a kill-only mechanism (i.e., a new `KillOriginFact`) to clear old loans when the RHS has no origins. Fixes llvm#186009
diff --git a/clang/include/clang/Analysis/Analyses/LifetimeSafety/Facts.h b/clang/include/clang/Analysis/Analyses/LifetimeSafety/Facts.h
@@ -55,6 +55,8 @@ class Fact {
     OriginEscapes,
     /// An origin is invalidated (e.g. vector resized).
     InvalidateOrigin,
+    /// All loans of an origin are cleared.
+    KillOrigin,
   };
 
 private:
@@ -316,6 +318,24 @@ class TestPointFact : public Fact {
             const OriginManager &) const override;
 };
 
+/// All loans are cleared from an origin (e.g., assigning a callable without
+/// tracked origins to std::function).
+class KillOriginFact : public Fact {
+  OriginID OID;
+
+public:
+  static bool classof(const Fact *F) {
+    return F->getKind() == Kind::KillOrigin;
+  }
+
+  KillOriginFact(OriginID OID) : Fact(Kind::KillOrigin), OID(OID) {}
+
+  OriginID getKilledOrigin() const { return OID; }
+
+  void dump(llvm::raw_ostream &OS, const LoanManager &,
+            const OriginManager &OM) const override;
+};
+
 class FactManager {
 public:
   FactManager(const AnalysisDeclContext &AC, const CFG &Cfg) : OriginMgr(AC) {
diff --git a/clang/include/clang/Analysis/Analyses/LifetimeSafety/LifetimeAnnotations.h b/clang/include/clang/Analysis/Analyses/LifetimeSafety/LifetimeAnnotations.h
@@ -80,6 +80,10 @@ bool isUniquePtrRelease(const CXXMethodDecl &MD);
 // https://en.cppreference.com/w/cpp/container#Iterator_invalidation
 bool isContainerInvalidationMethod(const CXXMethodDecl &MD);
 
+/// Returns true for standard library callable wrappers (e.g., std::function)
+/// that can propagate the stored lambda's origins.
+bool isStdCallableWrapperType(const CXXRecordDecl *RD);
+
 } // namespace clang::lifetimes
 
 #endif // LLVM_CLANG_ANALYSIS_ANALYSES_LIFETIMEANNOTATIONS_H
diff --git a/clang/lib/Analysis/LifetimeSafety/Dataflow.h b/clang/lib/Analysis/LifetimeSafety/Dataflow.h
@@ -180,6 +180,8 @@ class DataflowAnalysis {
       return D->transfer(In, *F->getAs<TestPointFact>());
     case Fact::Kind::InvalidateOrigin:
       return D->transfer(In, *F->getAs<InvalidateOriginFact>());
+    case Fact::Kind::KillOrigin:
+      return D->transfer(In, *F->getAs<KillOriginFact>());
     }
     llvm_unreachable("Unknown fact kind");
   }
@@ -193,6 +195,7 @@ class DataflowAnalysis {
   Lattice transfer(Lattice In, const UseFact &) { return In; }
   Lattice transfer(Lattice In, const TestPointFact &) { return In; }
   Lattice transfer(Lattice In, const InvalidateOriginFact &) { return In; }
+  Lattice transfer(Lattice In, const KillOriginFact &) { return In; }
 };
 } // namespace clang::lifetimes::internal
 #endif // LLVM_CLANG_ANALYSIS_ANALYSES_LIFETIMESAFETY_DATAFLOW_H
diff --git a/clang/lib/Analysis/LifetimeSafety/Facts.cpp b/clang/lib/Analysis/LifetimeSafety/Facts.cpp
@@ -103,6 +103,13 @@ void TestPointFact::dump(llvm::raw_ostream &OS, const LoanManager &,
   OS << "TestPoint (Annotation: \"" << getAnnotation() << "\")\n";
 }
 
+void KillOriginFact::dump(llvm::raw_ostream &OS, const LoanManager &,
+                          const OriginManager &OM) const {
+  OS << "KillOrigin (";
+  OM.dump(getKilledOrigin(), OS);
+  OS << ")\n";
+}
+
 llvm::StringMap<ProgramPoint> FactManager::getTestPoints() const {
   llvm::StringMap<ProgramPoint> AnnotationToPointMap;
   for (const auto &BlockFacts : BlockToFacts) {
diff --git a/clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp b/clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp
@@ -190,9 +190,8 @@ void FactsGenerator::VisitCXXConstructExpr(const CXXConstructExpr *CCE) {
     return;
   }
   // For defaulted (implicit or `= default`) copy/move constructors, propagate
-  // origins directly. User-defined copy/move constructors have opaque semantics
-  // and fall through to `handleFunctionCall`, where [[clang::lifetimebound]] is
-  // needed to propagate origins.
+  // origins directly. User-defined copy/move constructors are not handled here
+  // as they have opaque semantics.
   if (CCE->getConstructor()->isCopyOrMoveConstructor() &&
       CCE->getConstructor()->isDefaulted() && CCE->getNumArgs() == 1 &&
       hasOrigins(CCE->getType())) {
@@ -202,6 +201,16 @@ void FactsGenerator::VisitCXXConstructExpr(const CXXConstructExpr *CCE) {
       return;
     }
   }
+  // Standard library callable wrappers (e.g., std::function) propagate the
+  // stored lambda's origins.
+  if (const auto *RD = CCE->getType()->getAsCXXRecordDecl();
+      RD && isStdCallableWrapperType(RD) && CCE->getNumArgs() == 1) {
+    const Expr *Arg = CCE->getArg(0);
+    if (OriginList *ArgList = getRValueOrigins(Arg, getOriginsList(*Arg))) {
+      flow(getOriginsList(*CCE), ArgList, /*Kill=*/true);
+      return;
+    }
+  }
   handleFunctionCall(CCE, CCE->getConstructor(),
                      {CCE->getArgs(), CCE->getNumArgs()},
                      /*IsGslConstruction=*/false);
@@ -400,6 +409,15 @@ void FactsGenerator::handleAssignment(const Expr *LHSExpr,
     } else
       markUseAsWrite(DRE_LHS);
   }
+  if (!RHSList) {
+    // RHS has no tracked origins (e.g., assigning a callable without origins
+    // to std::function). Clear loans of the destination.
+    for (OriginList *LHSInner = LHSList->peelOuterOrigin(); LHSInner;
+         LHSInner = LHSInner->peelOuterOrigin())
+      CurrentBlockFacts.push_back(
+          FactMgr.createFact<KillOriginFact>(LHSInner->getOuterOriginID()));
+    return;
+  }
   // Kill the old loans of the destination origin and flow the new loans
   // from the source origin.
   flow(LHSList->peelOuterOrigin(), RHSList, /*Kill=*/true);
@@ -493,6 +511,13 @@ void FactsGenerator::VisitCXXOperatorCallExpr(const CXXOperatorCallExpr *OCE) {
       handleAssignment(OCE->getArg(0), OCE->getArg(1));
       return;
     }
+    // Standard library callable wrappers (e.g., std::function) can propagate
+    // the stored lambda's origins.
+    if (const auto *RD = LHSTy->getAsCXXRecordDecl();
+        RD && isStdCallableWrapperType(RD)) {
+      handleAssignment(OCE->getArg(0), OCE->getArg(1));
+      return;
+    }
     // Other tracked types: only defaulted operator= propagates origins.
     // User-defined operator= has opaque semantics, so don't handle them now.
     if (const auto *MD =
diff --git a/clang/lib/Analysis/LifetimeSafety/LifetimeAnnotations.cpp b/clang/lib/Analysis/LifetimeSafety/LifetimeAnnotations.cpp
@@ -382,4 +382,12 @@ bool isContainerInvalidationMethod(const CXXMethodDecl &MD) {
 
   return InvalidatingMethods->contains(MD.getName());
 }
+
+bool isStdCallableWrapperType(const CXXRecordDecl *RD) {
+  if (!RD || !isInStlNamespace(RD))
+    return false;
+  StringRef Name = getName(*RD);
+  return Name == "function" || Name == "move_only_function";
+}
+
 } // namespace clang::lifetimes
diff --git a/clang/lib/Analysis/LifetimeSafety/LiveOrigins.cpp b/clang/lib/Analysis/LifetimeSafety/LiveOrigins.cpp
@@ -166,6 +166,10 @@ class AnalysisImpl
     return Lattice(Factory.remove(In.LiveOrigins, OF.getDestOriginID()));
   }
 
+  Lattice transfer(Lattice In, const KillOriginFact &F) {
+    return Lattice(Factory.remove(In.LiveOrigins, F.getKilledOrigin()));
+  }
+
   Lattice transfer(Lattice In, const ExpireFact &F) {
     if (auto OID = F.getOriginID())
       return Lattice(Factory.remove(In.LiveOrigins, *OID));
diff --git a/clang/lib/Analysis/LifetimeSafety/LoanPropagation.cpp b/clang/lib/Analysis/LifetimeSafety/LoanPropagation.cpp
@@ -63,6 +63,9 @@ static llvm::BitVector computePersistentOrigins(const FactManager &FactMgr,
              Cur = Cur->peelOuterOrigin())
           CheckOrigin(Cur->getOuterOriginID());
         break;
+      case Fact::Kind::KillOrigin:
+        CheckOrigin(F->getAs<KillOriginFact>()->getKilledOrigin());
+        break;
       case Fact::Kind::MovedOrigin:
       case Fact::Kind::OriginEscapes:
       case Fact::Kind::Expire:
@@ -181,6 +184,10 @@ class AnalysisImpl
     return setLoans(In, DestOID, MergedLoans);
   }
 
+  Lattice transfer(Lattice In, const KillOriginFact &F) {
+    return setLoans(In, F.getKilledOrigin(), LoanSetFactory.getEmptySet());
+  }
+
   Lattice transfer(Lattice In, const ExpireFact &F) {
     if (auto OID = F.getOriginID())
       return setLoans(In, *OID, LoanSetFactory.getEmptySet());
diff --git a/clang/lib/Analysis/LifetimeSafety/Origins.cpp b/clang/lib/Analysis/LifetimeSafety/Origins.cpp
@@ -107,6 +107,10 @@ bool OriginManager::hasOrigins(QualType QT) const {
   const auto *RD = QT->getAsCXXRecordDecl();
   if (!RD)
     return false;
+  // Standard library callable wrappers (e.g., std::function) can propagate the
+  // stored lambda's origins.
+  if (isStdCallableWrapperType(RD))
+    return true;
   // TODO: Limit to lambdas for now. This will be extended to user-defined
   // structs with pointer-like fields.
   if (!RD->isLambda())
diff --git a/clang/test/Sema/Inputs/lifetime-analysis.h b/clang/test/Sema/Inputs/lifetime-analysis.h
@@ -268,4 +268,18 @@ struct true_type {
 template<class T> struct is_pointer : false_type {};
 template<class T> struct is_pointer<T*> : true_type {};
 template<class T> struct is_pointer<T* const> : true_type {};
+
+template<class> class function;
+template<class R, class... Args>
+class function<R(Args...)> {
+public:
+  template<class F> function(F) {}
+  function(const function&) {}
+  function(function&&) {}
+  template<class F> function& operator=(F) { return *this; }
+  function& operator=(const function&) { return *this; }
+  function& operator=(function&&) { return *this; }
+  ~function();
+};
+
 }
diff --git a/clang/test/Sema/warn-lifetime-safety-dangling-field.cpp b/clang/test/Sema/warn-lifetime-safety-dangling-field.cpp
@@ -196,3 +196,16 @@ struct DerivedWithCtor : BaseWithPointer {
   }
 };
 } // namespace DanglingPointerFieldInBaseClass
+
+namespace callable_wrappers {
+
+struct HasCallback {
+  std::function<void()> callback; // expected-note {{this field dangles}}
+
+  void set_callback() {
+    int local;
+    callback = [&local]() { (void)local; }; // expected-warning {{address of stack memory escapes to a field}}
+  }
+};
+
+} // namespace callable_wrappers
diff --git a/clang/test/Sema/warn-lifetime-safety-invalidations.cpp b/clang/test/Sema/warn-lifetime-safety-invalidations.cpp
@@ -527,3 +527,15 @@ struct S {
   }
 };
 } // namespace method_call_uses_field_invalidation
+
+namespace callable_wrappers {
+
+void function_captured_ref_invalidated() {
+  std::vector<int> v;
+  v.push_back(1);
+  std::function<void()> f = [&r = v[0]]() { (void)r; }; // expected-warning {{object whose reference is captured is later invalidated}}
+  v.push_back(2); // expected-note {{invalidated here}}
+  (void)f; // expected-note {{later used here}}
+}
+
+} // namespace callable_wrappers
diff --git a/clang/test/Sema/warn-lifetime-safety-noescape.cpp b/clang/test/Sema/warn-lifetime-safety-noescape.cpp
@@ -191,3 +191,11 @@ MyObj& return_ref_from_noescape_ptr(
 int* return_spaced_brackets(int* p [ [clang::noescape] /*some comment*/ ]) { // expected-warning {{parameter is marked [[clang::noescape]] but escapes}}
   return p; // expected-note {{returned here}}
 }
+
+namespace callable_wrappers {
+
+std::function<void()> escape_noescape_via_function(int &x [[clang::noescape]]) { // expected-warning {{parameter is marked [[clang::noescape]] but escapes}}
+  return [&]() { (void)x; }; // expected-note {{returned here}}
+}
+
+} // namespace callable_wrappers
diff --git a/clang/test/Sema/warn-lifetime-safety-suggestions.cpp b/clang/test/Sema/warn-lifetime-safety-suggestions.cpp
@@ -1,8 +1,8 @@
 // RUN: rm -rf %t
 // RUN: split-file %s %t
-// RUN: %clang_cc1 -fsyntax-only -flifetime-safety-inference -fexperimental-lifetime-safety-tu-analysis -Wlifetime-safety-suggestions -Wlifetime-safety -Wno-dangling -I%t -verify %t/test_source.cpp
-// RUN: %clang_cc1 -flifetime-safety-inference -fexperimental-lifetime-safety-tu-analysis -Wlifetime-safety-suggestions -Wlifetime-safety -Wno-dangling -I%t -fixit %t/test_source.cpp
-// RUN: %clang_cc1 -fsyntax-only -flifetime-safety-inference -fexperimental-lifetime-safety-tu-analysis -Wlifetime-safety-suggestions -Wno-dangling -I%t -Werror=lifetime-safety-suggestions %t/test_source.cpp
+// RUN: %clang_cc1 -fsyntax-only -flifetime-safety-inference -fexperimental-lifetime-safety-tu-analysis -Wlifetime-safety-suggestions -Wlifetime-safety -Wno-dangling -I%t -I%S -verify %t/test_source.cpp
+// RUN: %clang_cc1 -flifetime-safety-inference -fexperimental-lifetime-safety-tu-analysis -Wlifetime-safety-suggestions -Wlifetime-safety -Wno-dangling -I%t -I%S -fixit %t/test_source.cpp
+// RUN: %clang_cc1 -fsyntax-only -flifetime-safety-inference -fexperimental-lifetime-safety-tu-analysis -Wlifetime-safety-suggestions -Wno-dangling -I%t -I%S -Werror=lifetime-safety-suggestions %t/test_source.cpp
 
 View definition_before_header(View a);
 
@@ -67,6 +67,7 @@ struct ReturnThisPointer {
 //--- test_source.cpp
 
 #include "test_header.h"
+#include "Inputs/lifetime-analysis.h"
 
 View definition_before_header(View a) {
   return a;                               // expected-note {{param returned here}}
@@ -529,3 +530,20 @@ CaptureRefToBaseView test_ref_to_base_view() {
   return x; // expected-note {{returned here}}
 }
 } // namespace capturing_constructor
+
+namespace callable_wrappers {
+
+std::function<void()> return_lambda_capturing_param(int &x) { // expected-warning {{parameter in intra-TU function should be marked [[clang::lifetimebound]]}}
+  return [&]() { (void)x; }; // expected-note {{param returned here}}
+}
+
+void uaf_via_inferred_lifetimebound() {
+  std::function<void()> f = []() {};
+  {
+    int local;
+    f = return_lambda_capturing_param(local); // expected-warning {{object whose reference is captured does not live long enough}}
+  } // expected-note {{destroyed here}}
+  (void)f; // expected-note {{later used here}}
+}
+
+} // namespace callable_wrappers
diff --git a/clang/test/Sema/warn-lifetime-safety.cpp b/clang/test/Sema/warn-lifetime-safety.cpp
