[LifetimeSafety] Add parameter lifetime tracking in CFG (#169320)

This PR enhances the CFG builder to properly handle function parameters
in lifetime analysis:

1. Added code to include parameters in the initial scope during CFG
construction for both `FunctionDecl` and `BlockDecl` types
2. Added a special case to skip reference parameters, as they don't need
automatic destruction
3. Fixed several test cases that were previously marked as "FIXME" due
to missing parameter lifetime tracking

Previously, Clang's lifetime analysis was not properly tracking the
lifetime of function parameters, causing it to miss important
use-after-return bugs when parameter values were returned by reference
or address. This change ensures that parameters are properly tracked in
the CFG, allowing the analyzer to correctly identify when stack memory
associated with parameters is returned.

Fixes https://github.com/llvm/llvm-project/issues/169014

clang/lib/Analysis/CFG.cpp

@@ -1666,6 +1666,12 @@ std::unique_ptr<CFG> CFGBuilder::buildCFG(const Decl *D, Stmt *Statement) {
   assert(Succ == &cfg->getExit());
   Block = nullptr;  // the EXIT block is empty.  Create all other blocks lazily.
 
+  // Add parameters to the initial scope to handle their dtos and lifetime ends.
+  LocalScope *paramScope = nullptr;
+  if (const auto *FD = dyn_cast_or_null<FunctionDecl>(D))
+    for (ParmVarDecl *PD : FD->parameters())
+      paramScope = addLocalScopeForVarDecl(PD, paramScope);
+
   if (BuildOpts.AddImplicitDtors)
     if (const CXXDestructorDecl *DD = dyn_cast_or_null<CXXDestructorDecl>(D))
       addImplicitDtorsForDestructor(DD);
@@ -2246,6 +2252,11 @@ LocalScope* CFGBuilder::addLocalScopeForVarDecl(VarDecl *VD,
   if (!VD->hasLocalStorage())
     return Scope;
 
+  // Reference parameters are aliases to objects that live elsewhere, so they
+  // don't require automatic destruction or lifetime tracking.
+  if (isa<ParmVarDecl>(VD) && VD->getType()->isReferenceType())
+    return Scope;
+
   if (!BuildOpts.AddLifetime && !BuildOpts.AddScopes &&
       !needsAutomaticDestruction(VD)) {
     assert(BuildOpts.AddImplicitDtors);
@@ -5616,8 +5627,15 @@ public:
   bool handleDecl(const Decl *D, raw_ostream &OS) {
     DeclMapTy::iterator I = DeclMap.find(D);
 
-    if (I == DeclMap.end())
+    if (I == DeclMap.end()) {
+      // ParmVarDecls are not declared in the CFG itself, so they do not appear
+      // in DeclMap.
+      if (auto *PVD = dyn_cast_or_null<ParmVarDecl>(D)) {
+        OS << "[Parm: " << PVD->getNameAsString() << "]";
+        return true;
+      }
       return false;
+    }
 
     if (currentBlock >= 0 && I->second.first == (unsigned) currentBlock
                           && I->second.second == currStmt) {

clang/test/Analysis/lifetime-cfg-output.cpp

@@ -935,3 +935,31 @@ label:
   goto label;
   i++;
 }
+
+// CHECK:       [B2 (ENTRY)]
+// CHECK-NEXT:    Succs (1): B1
+// CHECK:       [B1]
+// CHECK-NEXT:    1: a
+// CHECK-NEXT:    2: [B1.1] (ImplicitCastExpr, LValueToRValue, int)
+// CHECK-NEXT:    3: b
+// CHECK-NEXT:    4: [B1.3] (ImplicitCastExpr, LValueToRValue, int)
+// CHECK-NEXT:    5: [B1.2] + [B1.4]
+// CHECK-NEXT:    6: c
+// CHECK-NEXT:    7: [B1.6] (ImplicitCastExpr, LValueToRValue, int)
+// CHECK-NEXT:    8: [B1.5] + [B1.7]
+// CHECK-NEXT:    9: int res = a + b + c;
+// CHECK-NEXT:    10: res
+// CHECK-NEXT:    11: [B1.10] (ImplicitCastExpr, LValueToRValue, int)
+// CHECK-NEXT:    12: return [B1.11];
+// CHECK-NEXT:    13: [B1.9] (Lifetime ends)
+// CHECK-NEXT:    14: [Parm: c] (Lifetime ends)
+// CHECK-NEXT:    15: [Parm: b] (Lifetime ends)
+// CHECK-NEXT:    16: [Parm: a] (Lifetime ends)
+// CHECK-NEXT:    Preds (1): B2
+// CHECK-NEXT:    Succs (1): B0
+// CHECK:       [B0 (EXIT)]
+// CHECK-NEXT:    Preds (1): B1
+int test_param_scope_end_order(int a, int b, int c) {
+  int res = a + b + c;
+  return res; 
+}

clang/test/Analysis/scopes-cfg-output.cpp

@@ -1437,12 +1437,14 @@ void test_cleanup_functions() {
 // CHECK-NEXT:    4: return;
 // CHECK-NEXT:    5: CleanupFunction (cleanup_int)
 // CHECK-NEXT:    6: CFGScopeEnd(i)
+// CHECK-NEXT:    7: CFGScopeEnd(m)
 // CHECK-NEXT:    Preds (1): B3
 // CHECK-NEXT:    Succs (1): B0
 // CHECK:      [B2]
 // CHECK-NEXT:    1: return;
 // CHECK-NEXT:    2: CleanupFunction (cleanup_int)
 // CHECK-NEXT:    3: CFGScopeEnd(i)
+// CHECK-NEXT:    4: CFGScopeEnd(m)
 // CHECK-NEXT:    Preds (1): B3
 // CHECK-NEXT:    Succs (1): B0
 // CHECK:      [B3]

clang/test/Sema/warn-lifetime-safety.cpp

@@ -529,14 +529,14 @@ TriviallyDestructedClass* trivial_class_uar () {
   return ptr;     // expected-note {{returned here}}
 }
 
-// FIXME: No lifetime warning for this as no expire facts are generated for parameters
 const int& return_parameter(int a) { 
-  return a; 
+  return a; // expected-warning {{address of stack memory is returned later}}
+            // expected-note@-1 {{returned here}}
 }
 
-// FIXME: No lifetime warning for this as no expire facts are generated for parameters
 int* return_pointer_to_parameter(int a) {
-    return &a;
+    return &a;  // expected-warning {{address of stack memory is returned later}}
+                // expected-note@-1 {{returned here}}
 }
 
 const int& return_reference_to_parameter(int a)
@@ -788,9 +788,52 @@ const MyObj& lifetimebound_return_ref_to_local() {
                              // expected-note@-1 {{returned here}}
 }
 
-// FIXME: Fails to diagnose UAR when a reference to a by-value param escapes via the return value.
-View lifetimebound_return_of_by_value_param(MyObj stack_param) {
-  return Identity(stack_param); 
+View lifetimebound_return_by_value_param(MyObj stack_param) {
+  return Identity(stack_param); // expected-warning {{address of stack memory is returned later}}
+                                // expected-note@-1 {{returned here}}
+}
+
+View lifetimebound_return_by_value_multiple_param(int cond, MyObj a, MyObj b, MyObj c) {
+  if (cond == 1) 
+    return Identity(a); // expected-warning {{address of stack memory is returned later}}
+                        // expected-note@-1 {{returned here}}
+  if (cond == 2) 
+    return Identity(b); // expected-warning {{address of stack memory is returned later}}
+                        // expected-note@-1 {{returned here}}
+  return Identity(c); // expected-warning {{address of stack memory is returned later}}
+                      // expected-note@-1 {{returned here}}
+}
+
+template<class T>
+View lifetimebound_return_by_value_param_template(T t) {
+  return Identity(t); // expected-warning {{address of stack memory is returned later}}
+                      // expected-note@-1 {{returned here}}
+}
+void use_lifetimebound_return_by_value_param_template() { 
+  lifetimebound_return_by_value_param_template(MyObj{}); // expected-note {{in instantiation of}}
+}
+
+void lambda_uar_param() {
+  auto lambda = [](MyObj stack_param) {
+    return Identity(stack_param); // expected-warning {{address of stack memory is returned later}}
+                                  // expected-note@-1 {{returned here}}
+  };
+  lambda(MyObj{});
+}
+
+// FIXME: This should be detected. We see correct destructors but origin flow breaks somewhere.
+namespace VariadicTemplatedParamsUAR{
+
+template<typename... Args>
+View Max(Args... args [[clang::lifetimebound]]);
+
+template<typename... Args>
+View lifetimebound_return_of_variadic_param(Args... args) {
+  return Max(args...);
+}
+void test_variadic() {
+  lifetimebound_return_of_variadic_param(MyObj{1}, MyObj{2}, MyObj{3});
+}
 }
 
 // FIXME: Fails to diagnose UAF when a reference to a by-value param escapes via an out-param.

clang/unittests/Analysis/FlowSensitive/LoggerTest.cpp

@@ -149,9 +149,18 @@ recordState(Elements=8, Branches=2, Joins=1)
 enterElement(return b ? p : q;)
 transfer()
 recordState(Elements=9, Branches=2, Joins=1)
+enterElement([Parm: q] (Lifetime ends))
+transfer()
+recordState(Elements=10, Branches=2, Joins=1)
+enterElement([Parm: p] (Lifetime ends))
+transfer()
+recordState(Elements=11, Branches=2, Joins=1)
+enterElement([Parm: b] (Lifetime ends))
+transfer()
+recordState(Elements=12, Branches=2, Joins=1)
 
 enterBlock(0, false)
-recordState(Elements=9, Branches=2, Joins=1)
+recordState(Elements=12, Branches=2, Joins=1)
 
 endAnalysis()
 )");