[LifetimeSafety] Detect use-after-return (#165370)

Adding "use-after-return" in Lifetime Analysis.

Detecting when a function returns a reference to its own stack memory:
[UAR Design
Doc](https://docs.google.com/document/d/1Wxjn_rJD_tuRdejP81dlb9VOckTkCq5-aE1nGcerb_o/edit?usp=sharing)

Consider the following example:

```cpp
std::string_view foo() {
    std::string_view a;
    std::string str = "small scoped string";
    a = str;
    return a;
}
```

The code adds a new Fact "OriginEscape" in the end of the CFG to
determine any loan that is escaping the function as shown below:

```
Function: foo
  Block B2:
  End of Block
  Block B1:
    OriginFlow (Dest: 0 (Decl: a), Src: 1 (Expr: CXXConstructExpr))
    OriginFlow (Dest: 2 (Expr: ImplicitCastExpr), Src: 3 (Expr: StringLiteral))
    Issue (0 (Path: operator=), ToOrigin: 4 (Expr: DeclRefExpr))
    OriginFlow (Dest: 5 (Expr: ImplicitCastExpr), Src: 4 (Expr: DeclRefExpr))
    Use (0 (Decl: a), Write)
    Issue (1 (Path: str), ToOrigin: 6 (Expr: DeclRefExpr))
    OriginFlow (Dest: 7 (Expr: ImplicitCastExpr), Src: 6 (Expr: DeclRefExpr))
    OriginFlow (Dest: 8 (Expr: CXXMemberCallExpr), Src: 7 (Expr: ImplicitCastExpr))
    OriginFlow (Dest: 9 (Expr: ImplicitCastExpr), Src: 8 (Expr: CXXMemberCallExpr))
    OriginFlow (Dest: 10 (Expr: ImplicitCastExpr), Src: 9 (Expr: ImplicitCastExpr))
    OriginFlow (Dest: 11 (Expr: MaterializeTemporaryExpr), Src: 10 (Expr: ImplicitCastExpr))
    OriginFlow (Dest: 0 (Decl: a), Src: 11 (Expr: MaterializeTemporaryExpr))
    Use (0 (Decl: a), Read)
    OriginFlow (Dest: 12 (Expr: ImplicitCastExpr), Src: 0 (Decl: a))
    OriginFlow (Dest: 13 (Expr: CXXConstructExpr), Src: 12 (Expr: ImplicitCastExpr))
    Expire (1 (Path: str))
    OriginEscapes (13 (Expr: CXXConstructExpr))
  End of Block
  Block B0:
  End of Block
```

The confidence of the report is determined by checking if at least one
of the loans returned is not expired (strict). If all loans are expired
it is considered permissive.

More information [UAR Design
Doc](https://docs.google.com/document/d/1Wxjn_rJD_tuRdejP81dlb9VOckTkCq5-aE1nGcerb_o/edit?usp=sharing)

Author: kashika0112

clang/include/clang/Analysis/Analyses/LifetimeSafety/Facts.h

@@ -42,12 +42,12 @@ class Fact {
     /// it. Otherwise, the source's loan set is merged into the destination's
     /// loan set.
     OriginFlow,
-    /// An origin escapes the function by flowing into the return value.
-    ReturnOfOrigin,
     /// An origin is used (eg. appears as l-value expression like DeclRefExpr).
     Use,
     /// A marker for a specific point in the code, for testing.
     TestPoint,
+    /// An origin that escapes the function scope (e.g., via return).
+    OriginEscapes,
   };
 
 private:
@@ -136,16 +136,19 @@ class OriginFlowFact : public Fact {
             const OriginManager &OM) const override;
 };
 
-class ReturnOfOriginFact : public Fact {
+class OriginEscapesFact : public Fact {
   OriginID OID;
+  const Expr *EscapeExpr;
 
 public:
   static bool classof(const Fact *F) {
-    return F->getKind() == Kind::ReturnOfOrigin;
+    return F->getKind() == Kind::OriginEscapes;
   }
 
-  ReturnOfOriginFact(OriginID OID) : Fact(Kind::ReturnOfOrigin), OID(OID) {}
-  OriginID getReturnedOriginID() const { return OID; }
+  OriginEscapesFact(OriginID OID, const Expr *EscapeExpr)
+      : Fact(Kind::OriginEscapes), OID(OID), EscapeExpr(EscapeExpr) {}
+  OriginID getEscapedOriginID() const { return OID; }
+  const Expr *getEscapeExpr() const { return EscapeExpr; };
   void dump(llvm::raw_ostream &OS, const LoanManager &,
             const OriginManager &OM) const override;
 };
@@ -225,6 +228,9 @@ class FactManager {
   /// user-defined locations in the code.
   /// \note This is intended for testing only.
   llvm::StringMap<ProgramPoint> getTestPoints() const;
+  /// Retrieves all the facts in the block containing Program Point P.
+  /// \note This is intended for testing only.
+  llvm::ArrayRef<const Fact *> getBlockContaining(ProgramPoint P) const;
 
   unsigned getNumFacts() const { return NextFactID.Value; }
 

clang/include/clang/Analysis/Analyses/LifetimeSafety/FactsGenerator.h

@@ -94,6 +94,10 @@ class FactsGenerator : public ConstStmtVisitor<FactsGenerator> {
   FactManager &FactMgr;
   AnalysisDeclContext &AC;
   llvm::SmallVector<Fact *> CurrentBlockFacts;
+  // Collect origins that escape the function in this block (OriginEscapesFact),
+  // appended at the end of CurrentBlockFacts to ensure they appear after
+  // ExpireFact entries.
+  llvm::SmallVector<Fact *> EscapesInCurrentBlock;
   // To distinguish between reads and writes for use-after-free checks, this map
   // stores the `UseFact` for each `DeclRefExpr`. We initially identify all
   // `DeclRefExpr`s as "read" uses. When an assignment is processed, the use

clang/include/clang/Analysis/Analyses/LifetimeSafety/LifetimeSafety.h

@@ -42,6 +42,11 @@ class LifetimeSafetyReporter {
   virtual void reportUseAfterFree(const Expr *IssueExpr, const Expr *UseExpr,
                                   SourceLocation FreeLoc,
                                   Confidence Confidence) {}
+
+  virtual void reportUseAfterReturn(const Expr *IssueExpr,
+                                    const Expr *EscapeExpr,
+                                    SourceLocation ExpiryLoc,
+                                    Confidence Confidence) {}
 };
 
 /// The main entry point for the analysis.

clang/include/clang/Analysis/Analyses/LifetimeSafety/LiveOrigins.h

@@ -31,6 +31,9 @@
 
 namespace clang::lifetimes::internal {
 
+using CausingFactType =
+    ::llvm::PointerUnion<const UseFact *, const OriginEscapesFact *>;
+
 enum class LivenessKind : uint8_t {
   Dead,  // Not alive
   Maybe, // Live on some path but not all paths (may-be-live)
@@ -43,7 +46,7 @@ struct LivenessInfo {
   /// multiple uses along different paths, this will point to the use appearing
   /// earlier in the translation unit.
   /// This is 'null' when the origin is not live.
-  const UseFact *CausingUseFact;
+  CausingFactType CausingFact;
 
   /// The kind of liveness of the origin.
   /// `Must`: The origin is live on all control-flow paths from the current
@@ -56,17 +59,16 @@ struct LivenessInfo {
   /// while `Maybe`-be-alive suggests a potential one on some paths.
   LivenessKind Kind;
 
-  LivenessInfo() : CausingUseFact(nullptr), Kind(LivenessKind::Dead) {}
-  LivenessInfo(const UseFact *UF, LivenessKind K)
-      : CausingUseFact(UF), Kind(K) {}
+  LivenessInfo() : CausingFact(nullptr), Kind(LivenessKind::Dead) {}
+  LivenessInfo(CausingFactType CF, LivenessKind K) : CausingFact(CF), Kind(K) {}
 
   bool operator==(const LivenessInfo &Other) const {
-    return CausingUseFact == Other.CausingUseFact && Kind == Other.Kind;
+    return CausingFact == Other.CausingFact && Kind == Other.Kind;
   }
   bool operator!=(const LivenessInfo &Other) const { return !(*this == Other); }
 
   void Profile(llvm::FoldingSetNodeID &IDBuilder) const {
-    IDBuilder.AddPointer(CausingUseFact);
+    IDBuilder.AddPointer(CausingFact.getOpaqueValue());
     IDBuilder.Add(Kind);
   }
 };

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -10744,8 +10744,19 @@ def warn_lifetime_safety_loan_expires_permissive : Warning<
 def warn_lifetime_safety_loan_expires_strict : Warning<
    "object whose reference is captured may not live long enough">,
    InGroup<LifetimeSafetyStrict>, DefaultIgnore;
+
+def warn_lifetime_safety_return_stack_addr_permissive
+    : Warning<"address of stack memory is returned later">,
+      InGroup<LifetimeSafetyPermissive>,
+      DefaultIgnore;
+def warn_lifetime_safety_return_stack_addr_strict
+    : Warning<"address of stack memory may be returned later">,
+      InGroup<LifetimeSafetyStrict>,
+      DefaultIgnore;
+
 def note_lifetime_safety_used_here : Note<"later used here">;
 def note_lifetime_safety_destroyed_here : Note<"destroyed here">;
+def note_lifetime_safety_returned_here : Note<"returned here">;
 
 // For non-floating point, expressions of the form x == x or x != x
 // should result in a warning, since these always evaluate to a constant.

clang/lib/Analysis/LifetimeSafety/Checker.cpp

@@ -43,7 +43,7 @@ namespace {
 /// Struct to store the complete context for a potential lifetime violation.
 struct PendingWarning {
   SourceLocation ExpiryLoc; // Where the loan expired.
-  const Expr *UseExpr;      // Where the origin holding this loan was used.
+  llvm::PointerUnion<const UseFact *, const OriginEscapesFact *> CausingFact;
   Confidence ConfidenceLevel;
 };
 
@@ -68,7 +68,7 @@ class LifetimeChecker {
     issuePendingWarnings();
   }
 
-  /// Checks for use-after-free errors when a loan expires.
+  /// Checks for use-after-free & use-after-return errors when a loan expires.
   ///
   /// This method examines all live origins at the expiry point and determines
   /// if any of them hold the expiring loan. If so, it creates a pending
@@ -83,7 +83,11 @@ class LifetimeChecker {
     LoanID ExpiredLoan = EF->getLoanID();
     LivenessMap Origins = LiveOrigins.getLiveOriginsAt(EF);
     Confidence CurConfidence = Confidence::None;
-    const UseFact *BadUse = nullptr;
+    // The UseFact or OriginEscapesFact most indicative of a lifetime error,
+    // prioritized by earlier source location.
+    llvm::PointerUnion<const UseFact *, const OriginEscapesFact *>
+        BestCausingFact = nullptr;
+
     for (auto &[OID, LiveInfo] : Origins) {
       LoanSet HeldLoans = LoanPropagation.getLoans(OID, EF);
       if (!HeldLoans.contains(ExpiredLoan))
@@ -92,17 +96,17 @@ class LifetimeChecker {
       Confidence NewConfidence = livenessKindToConfidence(LiveInfo.Kind);
       if (CurConfidence < NewConfidence) {
         CurConfidence = NewConfidence;
-        BadUse = LiveInfo.CausingUseFact;
+        BestCausingFact = LiveInfo.CausingFact;
       }
     }
-    if (!BadUse)
+    if (!BestCausingFact)
       return;
     // We have a use-after-free.
     Confidence LastConf = FinalWarningsMap.lookup(ExpiredLoan).ConfidenceLevel;
     if (LastConf >= CurConfidence)
       return;
     FinalWarningsMap[ExpiredLoan] = {/*ExpiryLoc=*/EF->getExpiryLoc(),
-                                     /*UseExpr=*/BadUse->getUseExpr(),
+                                     /*BestCausingFact=*/BestCausingFact,
                                      /*ConfidenceLevel=*/CurConfidence};
   }
 
@@ -112,8 +116,20 @@ class LifetimeChecker {
     for (const auto &[LID, Warning] : FinalWarningsMap) {
       const Loan &L = FactMgr.getLoanMgr().getLoan(LID);
       const Expr *IssueExpr = L.IssueExpr;
-      Reporter->reportUseAfterFree(IssueExpr, Warning.UseExpr,
-                                   Warning.ExpiryLoc, Warning.ConfidenceLevel);
+      llvm::PointerUnion<const UseFact *, const OriginEscapesFact *>
+          CausingFact = Warning.CausingFact;
+      Confidence Confidence = Warning.ConfidenceLevel;
+      SourceLocation ExpiryLoc = Warning.ExpiryLoc;
+
+      if (const auto *UF = CausingFact.dyn_cast<const UseFact *>())
+        Reporter->reportUseAfterFree(IssueExpr, UF->getUseExpr(), ExpiryLoc,
+                                     Confidence);
+      else if (const auto *OEF =
+                   CausingFact.dyn_cast<const OriginEscapesFact *>())
+        Reporter->reportUseAfterReturn(IssueExpr, OEF->getEscapeExpr(),
+                                       ExpiryLoc, Confidence);
+      else
+        llvm_unreachable("Unhandled CausingFact type");
     }
   }
 };

clang/lib/Analysis/LifetimeSafety/Dataflow.h

@@ -170,8 +170,8 @@ class DataflowAnalysis {
       return D->transfer(In, *F->getAs<ExpireFact>());
     case Fact::Kind::OriginFlow:
       return D->transfer(In, *F->getAs<OriginFlowFact>());
-    case Fact::Kind::ReturnOfOrigin:
-      return D->transfer(In, *F->getAs<ReturnOfOriginFact>());
+    case Fact::Kind::OriginEscapes:
+      return D->transfer(In, *F->getAs<OriginEscapesFact>());
     case Fact::Kind::Use:
       return D->transfer(In, *F->getAs<UseFact>());
     case Fact::Kind::TestPoint:
@@ -184,7 +184,7 @@ class DataflowAnalysis {
   Lattice transfer(Lattice In, const IssueFact &) { return In; }
   Lattice transfer(Lattice In, const ExpireFact &) { return In; }
   Lattice transfer(Lattice In, const OriginFlowFact &) { return In; }
-  Lattice transfer(Lattice In, const ReturnOfOriginFact &) { return In; }
+  Lattice transfer(Lattice In, const OriginEscapesFact &) { return In; }
   Lattice transfer(Lattice In, const UseFact &) { return In; }
   Lattice transfer(Lattice In, const TestPointFact &) { return In; }
 };

clang/lib/Analysis/LifetimeSafety/Facts.cpp

@@ -43,10 +43,10 @@ void OriginFlowFact::dump(llvm::raw_ostream &OS, const LoanManager &,
   OS << ")\n";
 }
 
-void ReturnOfOriginFact::dump(llvm::raw_ostream &OS, const LoanManager &,
-                              const OriginManager &OM) const {
-  OS << "ReturnOfOrigin (";
-  OM.dump(getReturnedOriginID(), OS);
+void OriginEscapesFact::dump(llvm::raw_ostream &OS, const LoanManager &,
+                             const OriginManager &OM) const {
+  OS << "OriginEscapes (";
+  OM.dump(getEscapedOriginID(), OS);
   OS << ")\n";
 }
 
@@ -95,4 +95,14 @@ void FactManager::dump(const CFG &Cfg, AnalysisDeclContext &AC) const {
   }
 }
 
+llvm::ArrayRef<const Fact *>
+FactManager::getBlockContaining(ProgramPoint P) const {
+  for (const auto &BlockToFactsVec : BlockToFacts) {
+    for (const Fact *F : BlockToFactsVec)
+      if (F == P)
+        return BlockToFactsVec;
+  }
+  return {};
+}
+
 } // namespace clang::lifetimes::internal

clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp

@@ -58,6 +58,7 @@ void FactsGenerator::run() {
   // initializations and destructions are processed in the correct sequence.
   for (const CFGBlock *Block : *AC.getAnalysis<PostOrderCFGView>()) {
     CurrentBlockFacts.clear();
+    EscapesInCurrentBlock.clear();
     for (unsigned I = 0; I < Block->size(); ++I) {
       const CFGElement &Element = Block->Elements[I];
       if (std::optional<CFGStmt> CS = Element.getAs<CFGStmt>())
@@ -66,6 +67,8 @@ void FactsGenerator::run() {
                    Element.getAs<CFGAutomaticObjDtor>())
         handleDestructor(*DtorOpt);
     }
+    CurrentBlockFacts.append(EscapesInCurrentBlock.begin(),
+                             EscapesInCurrentBlock.end());
     FactMgr.addBlockFacts(Block, CurrentBlockFacts);
   }
 }
@@ -166,7 +169,8 @@ void FactsGenerator::VisitReturnStmt(const ReturnStmt *RS) {
   if (const Expr *RetExpr = RS->getRetValue()) {
     if (hasOrigin(RetExpr)) {
       OriginID OID = FactMgr.getOriginMgr().getOrCreate(*RetExpr);
-      CurrentBlockFacts.push_back(FactMgr.createFact<ReturnOfOriginFact>(OID));
+      EscapesInCurrentBlock.push_back(
+          FactMgr.createFact<OriginEscapesFact>(OID, RetExpr));
     }
   }
 }

clang/lib/Analysis/LifetimeSafety/LiveOrigins.cpp

@@ -53,6 +53,14 @@ struct Lattice {
   }
 };
 
+static SourceLocation GetFactLoc(CausingFactType F) {
+  if (const auto *UF = F.dyn_cast<const UseFact *>())
+    return UF->getUseExpr()->getExprLoc();
+  if (const auto *OEF = F.dyn_cast<const OriginEscapesFact *>())
+    return OEF->getEscapeExpr()->getExprLoc();
+  llvm_unreachable("unhandled causing fact in PointerUnion");
+}
+
 /// The analysis that tracks which origins are live, with granular information
 /// about the causing use fact and confidence level. This is a backward
 /// analysis.
@@ -74,11 +82,14 @@ class AnalysisImpl
   /// one.
   Lattice join(Lattice L1, Lattice L2) const {
     LivenessMap Merged = L1.LiveOrigins;
-    // Take the earliest UseFact to make the join hermetic and commutative.
-    auto CombineUseFact = [](const UseFact &A,
-                             const UseFact &B) -> const UseFact * {
-      return A.getUseExpr()->getExprLoc() < B.getUseExpr()->getExprLoc() ? &A
-                                                                         : &B;
+    // Take the earliest Fact to make the join hermetic and commutative.
+    auto CombineCausingFact = [](CausingFactType A,
+                                 CausingFactType B) -> CausingFactType {
+      if (!A)
+        return B;
+      if (!B)
+        return A;
+      return GetFactLoc(A) < GetFactLoc(B) ? A : B;
     };
     auto CombineLivenessKind = [](LivenessKind K1,
                                   LivenessKind K2) -> LivenessKind {
@@ -93,12 +104,11 @@ class AnalysisImpl
                                    const LivenessInfo *L2) -> LivenessInfo {
       assert((L1 || L2) && "unexpectedly merging 2 empty sets");
       if (!L1)
-        return LivenessInfo(L2->CausingUseFact, LivenessKind::Maybe);
+        return LivenessInfo(L2->CausingFact, LivenessKind::Maybe);
       if (!L2)
-        return LivenessInfo(L1->CausingUseFact, LivenessKind::Maybe);
-      return LivenessInfo(
-          CombineUseFact(*L1->CausingUseFact, *L2->CausingUseFact),
-          CombineLivenessKind(L1->Kind, L2->Kind));
+        return LivenessInfo(L1->CausingFact, LivenessKind::Maybe);
+      return LivenessInfo(CombineCausingFact(L1->CausingFact, L2->CausingFact),
+                          CombineLivenessKind(L1->Kind, L2->Kind));
     };
     return Lattice(utils::join(
         L1.LiveOrigins, L2.LiveOrigins, Factory, CombineLivenessInfo,
@@ -120,6 +130,14 @@ class AnalysisImpl
                                LivenessInfo(&UF, LivenessKind::Must)));
   }
 
+  /// An escaping origin (e.g., via return) makes the origin live with definite
+  /// confidence, as it dominates this program point.
+  Lattice transfer(Lattice In, const OriginEscapesFact &OEF) {
+    OriginID OID = OEF.getEscapedOriginID();
+    return Lattice(Factory.add(In.LiveOrigins, OID,
+                               LivenessInfo(&OEF, LivenessKind::Must)));
+  }
+
   /// Issuing a new loan to an origin kills its liveness.
   Lattice transfer(Lattice In, const IssueFact &IF) {
     return Lattice(Factory.remove(In.LiveOrigins, IF.getOriginID()));