[clang] clang: fix lifetime_capture_by indexing for member functions (PR #212264)

[Aditya prakash jha via cfe-commits]

https://github.com/AdityaOP007 updated https://github.com/llvm/llvm-project/pull/212264

>From 1eee2fd8d6653fe62adc50176136c759e6928814 Mon Sep 17 00:00:00 2001
From: Aditya Prakash Jha <adityaprakashjha321 at gmail.com>
Date: Mon, 27 Jul 2026 20:14:59 +0530
Subject: [PATCH] clang: fix lifetime_capture_by indexing for member functions

---
 .../LifetimeSafety/FactsGenerator.cpp         | 100 ++++++++++--------
 clang/test/Sema/LifetimeSafety/safety.cpp     |  37 +++++++
 2 files changed, 92 insertions(+), 45 deletions(-)

diff --git a/clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp b/clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp
index ac6267dabf48e..65aaa1ed02cad 100644
--- a/clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp
+++ b/clang/lib/Analysis/LifetimeSafety/FactsGenerator.cpp
@@ -1019,54 +1019,64 @@ void FactsGenerator::handleLifetimeCaptureBy(const FunctionDecl *FD,
   // FIXME: Add support for capture_by on constructors.
   if (isa<CXXConstructorDecl>(FD))
     return;
-  const auto *Method = dyn_cast<CXXMethodDecl>(FD);
-  bool IsInstance =
-      Method && Method->isInstance() && !isa<CXXConstructorDecl>(FD);
-  auto getArgCaptureBy = [FD,
-                          IsInstance](unsigned I) -> LifetimeCaptureByAttr * {
-    const ParmVarDecl *PVD = nullptr;
-    if (IsInstance) {
-      // FIXME: Add support for I == 0 i.e. capture_by on function declarations
-      if (I > 0 && I - 1 < FD->getNumParams())
-        PVD = FD->getParamDecl(I - 1);
-    } else {
-      if (I < FD->getNumParams())
-        PVD = FD->getParamDecl(I);
+  bool HasImplicitThisParam = hasImplicitObjectParameter(FD);
+  auto processArgAttrs = [&](unsigned I, auto Callback) {
+    if (HasImplicitThisParam) {
+      if (I == 0) {
+        if (TypeSourceInfo *TSI = FD->getTypeSourceInfo()) {
+          AttributedTypeLoc ATL;
+          for (TypeLoc TL = TSI->getTypeLoc();
+               (ATL = TL.getAsAdjusted<AttributedTypeLoc>());
+               TL = ATL.getModifiedLoc())
+            if (auto *A = ATL.getAttrAs<LifetimeCaptureByAttr>())
+              Callback(A);
+        }
+        return;
+      }
+      --I;
     }
-    return PVD ? PVD->getAttr<LifetimeCaptureByAttr>() : nullptr;
+    if (I < FD->getNumParams())
+      for (const auto *A :
+           FD->getParamDecl(I)->specific_attrs<LifetimeCaptureByAttr>())
+        Callback(A);
   };
   for (unsigned I = 0; I < Args.size(); ++I) {
-    const LifetimeCaptureByAttr *Attr = getArgCaptureBy(I);
-    if (!Attr)
-      continue;
-    OriginList *CapturedOriginList = getOriginsList(*Args[I]);
-    if (!CapturedOriginList)
-      continue;
-    for (int CapturingArgIdx : Attr->params()) {
-      // FIXME: Add support for capturing to Global/unknown.
-      if (CapturingArgIdx == LifetimeCaptureByAttr::Global ||
-          CapturingArgIdx == LifetimeCaptureByAttr::Unknown ||
-          CapturingArgIdx == LifetimeCaptureByAttr::Invalid)
-        continue;
-      ArrayRef<const Expr *> CallArgs = IsInstance ? Args.drop_front() : Args;
-      const Expr *CapturedByArg =
-          (CapturingArgIdx == LifetimeCaptureByAttr::This)
-              ? Args[0]
-              : CallArgs[CapturingArgIdx];
-      assert(CapturedByArg && "Capturer expression must be valid");
-
-      OriginList *CapturingOriginList = getOriginsList(*CapturedByArg);
-      OriginList *Dest = getRValueOrigins(CapturedByArg, CapturingOriginList);
-      if (!Dest)
-        continue;
-      // KillDest=false because we cannot know if previous captures are being
-      // replaced or accumulated. Multiple successive captures into the same
-      // destination must all be tracked, so captured lifetimes are always
-      // merged.
-      CurrentBlockFacts.push_back(FactMgr.createFact<OriginFlowFact>(
-          Dest->getOuterOriginID(), CapturedOriginList->getOuterOriginID(),
-          /*KillDest=*/false));
-    }
+    OriginList *CapturedOriginList = nullptr;
+    processArgAttrs(I, [&](const LifetimeCaptureByAttr *Attr) {
+      if (!Attr)
+        return;
+      if (!CapturedOriginList) {
+        CapturedOriginList = getOriginsList(*Args[I]);
+        if (!CapturedOriginList)
+          return;
+      }
+      for (int CapturingArgIdx : Attr->params()) {
+        // FIXME: Add support for capturing to Global/unknown.
+        if (CapturingArgIdx == LifetimeCaptureByAttr::Global ||
+            CapturingArgIdx == LifetimeCaptureByAttr::Unknown ||
+            CapturingArgIdx == LifetimeCaptureByAttr::Invalid)
+          continue;
+        if (CapturingArgIdx < 0 ||
+            CapturingArgIdx >= static_cast<int>(Args.size()))
+          continue;
+
+        const Expr *CapturedByArg = Args[CapturingArgIdx];
+        assert(CapturedByArg && "Capturer expression must be valid");
+
+        OriginList *CapturingOriginList = getOriginsList(*CapturedByArg);
+        OriginList *Dest =
+            getRValueOrigins(CapturedByArg, CapturingOriginList);
+        if (!Dest)
+          continue;
+        // KillDest=false because we cannot know if previous captures are being
+        // replaced or accumulated. Multiple successive captures into the same
+        // destination must all be tracked, so captured lifetimes are always
+        // merged.
+        CurrentBlockFacts.push_back(FactMgr.createFact<OriginFlowFact>(
+            Dest->getOuterOriginID(), CapturedOriginList->getOuterOriginID(),
+            /*KillDest=*/false));
+      }
+    });
   }
 }
 
diff --git a/clang/test/Sema/LifetimeSafety/safety.cpp b/clang/test/Sema/LifetimeSafety/safety.cpp
index f3b042cce74b2..1f2673348780c 100644
--- a/clang/test/Sema/LifetimeSafety/safety.cpp
+++ b/clang/test/Sema/LifetimeSafety/safety.cpp
@@ -4272,3 +4272,40 @@ void test_cyclic_cfg(int n) {
   }         // expected-note {{local variable 'a' is destroyed here}}
   v.use();  // expected-note {{later used here}}
 }
+
+namespace member_capture_by_param {
+struct S {
+  void cap(int *p [[clang::lifetime_capture_by(c)]], int *&c);
+  void cap_const(int *p [[clang::lifetime_capture_by(c)]], int *&c) const;
+  static void cap_static(int *p [[clang::lifetime_capture_by(c)]], int *&c);
+
+  template <typename T>
+  void cap_template(T *p [[clang::lifetime_capture_by(c)]], T *&c);
+
+  void cap_overload(int *p [[clang::lifetime_capture_by(c)]], int *&c);
+  void cap_overload(double *p [[clang::lifetime_capture_by(c)]], double *&c);
+
+  void cap_multi(int *p [[clang::lifetime_capture_by(c1, c2)]], int *&c1, int *&c2);
+  void cap_reordered(int *&c, int *p [[clang::lifetime_capture_by(c)]]);
+};
+
+namespace ns {
+  void cap_ns(int *p [[clang::lifetime_capture_by(c)]], int *&c);
+}
+
+void test_member_capture_by() {
+  S s;
+  int i;
+  int *c;
+  s.cap(&i, c);
+  s.cap_const(&i, c);
+  S::cap_static(&i, c);
+  s.cap_template(&i, c);
+  s.cap_overload(&i, c);
+  int *c2;
+  s.cap_multi(&i, c, c2);
+  s.cap_reordered(c, &i);
+  ns::cap_ns(&i, c);
+}
+} // namespace member_capture_by_param
+