[clang-tools-extra] [clang-tidy] loop convert can handle lambda init capture (PR #109159)

[Congcong Cai via cfe-commits]

https://github.com/HerrCai0907 updated https://github.com/llvm/llvm-project/pull/109159

>From 4fc1d24c4ff22a8da22454aebe7053ea76419767 Mon Sep 17 00:00:00 2001
From: Congcong Cai <congcongcai0907 at 163.com>
Date: Wed, 18 Sep 2024 23:26:02 +0800
Subject: [PATCH] [clang-tidy] loop convert can handle lambda init capture

Fixes: #109083
---
 .../clang-tidy/modernize/LoopConvertUtils.cpp | 25 ++++++++++-------
 .../clang-tidy/modernize/LoopConvertUtils.h   |  2 ++
 clang-tools-extra/docs/ReleaseNotes.rst       |  4 +++
 .../checkers/modernize/loop-convert-basic.cpp | 27 +++++++++++++++++++
 4 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/clang-tools-extra/clang-tidy/modernize/LoopConvertUtils.cpp b/clang-tools-extra/clang-tidy/modernize/LoopConvertUtils.cpp
index c0bf4903ec3911..90c16bfddd84f8 100644
--- a/clang-tools-extra/clang-tidy/modernize/LoopConvertUtils.cpp
+++ b/clang-tools-extra/clang-tidy/modernize/LoopConvertUtils.cpp
@@ -777,7 +777,7 @@ bool ForLoopIndexUseVisitor::TraverseLambdaCapture(LambdaExpr *LE,
                                                    const LambdaCapture *C,
                                                    Expr *Init) {
   if (C->capturesVariable()) {
-    const ValueDecl *VDecl = C->getCapturedVar();
+    ValueDecl *VDecl = C->getCapturedVar();
     if (areSameVariable(IndexVar, VDecl)) {
       // FIXME: if the index is captured, it will count as an usage and the
       // alias (if any) won't work, because it is only used in case of having
@@ -787,6 +787,8 @@ bool ForLoopIndexUseVisitor::TraverseLambdaCapture(LambdaExpr *LE,
                                                        : Usage::UK_CaptureByRef,
                      C->getLocation()));
     }
+    if (VDecl->isInitCapture())
+      TraverseStmtImpl(cast<VarDecl>(VDecl)->getInit());
   }
   return VisitorBase::TraverseLambdaCapture(LE, C, Init);
 }
@@ -816,6 +818,17 @@ bool ForLoopIndexUseVisitor::VisitDeclStmt(DeclStmt *S) {
   return true;
 }
 
+bool ForLoopIndexUseVisitor::TraverseStmtImpl(Stmt *S) {
+  // All this pointer swapping is a mechanism for tracking immediate parentage
+  // of Stmts.
+  const Stmt *OldNextParent = NextStmtParent;
+  CurrStmtParent = NextStmtParent;
+  NextStmtParent = S;
+  bool Result = VisitorBase::TraverseStmt(S);
+  NextStmtParent = OldNextParent;
+  return Result;
+}
+
 bool ForLoopIndexUseVisitor::TraverseStmt(Stmt *S) {
   // If this is an initialization expression for a lambda capture, prune the
   // traversal so that we don't end up diagnosing the contained DeclRefExpr as
@@ -828,15 +841,7 @@ bool ForLoopIndexUseVisitor::TraverseStmt(Stmt *S) {
       return true;
     }
   }
-
-  // All this pointer swapping is a mechanism for tracking immediate parentage
-  // of Stmts.
-  const Stmt *OldNextParent = NextStmtParent;
-  CurrStmtParent = NextStmtParent;
-  NextStmtParent = S;
-  bool Result = VisitorBase::TraverseStmt(S);
-  NextStmtParent = OldNextParent;
-  return Result;
+  return TraverseStmtImpl(S);
 }
 
 std::string VariableNamer::createIndexName() {
diff --git a/clang-tools-extra/clang-tidy/modernize/LoopConvertUtils.h b/clang-tools-extra/clang-tidy/modernize/LoopConvertUtils.h
index cfceb3b586842d..ca9c1855038b53 100644
--- a/clang-tools-extra/clang-tidy/modernize/LoopConvertUtils.h
+++ b/clang-tools-extra/clang-tidy/modernize/LoopConvertUtils.h
@@ -354,6 +354,8 @@ class ForLoopIndexUseVisitor
   bool VisitDeclStmt(DeclStmt *S);
   bool TraverseStmt(Stmt *S);
 
+  bool TraverseStmtImpl(Stmt *S);
+
   /// Add an expression to the list of expressions on which the container
   /// expression depends.
   void addComponent(const Expr *E);
diff --git a/clang-tools-extra/docs/ReleaseNotes.rst b/clang-tools-extra/docs/ReleaseNotes.rst
index d284bb62f7c7f4..2aa2fd961e721f 100644
--- a/clang-tools-extra/docs/ReleaseNotes.rst
+++ b/clang-tools-extra/docs/ReleaseNotes.rst
@@ -139,6 +139,10 @@ Changes in existing checks
   as a replacement for parameters of incomplete C array type in C++20 and 
   ``std::array`` or ``std::vector`` before C++20.
 
+- Improved :doc:`modernize-loop-convert
+  <clang-tidy/checks/modernize/loop-convert>` check to fix false positive when
+  using loop variable in initializer of lambda capture.
+
 - Improved :doc:`modernize-use-std-format
   <clang-tidy/checks/modernize/use-std-format>` check to support replacing
   member function calls too.
diff --git a/clang-tools-extra/test/clang-tidy/checkers/modernize/loop-convert-basic.cpp b/clang-tools-extra/test/clang-tidy/checkers/modernize/loop-convert-basic.cpp
index 695925a5d877c9..608e580591205c 100644
--- a/clang-tools-extra/test/clang-tidy/checkers/modernize/loop-convert-basic.cpp
+++ b/clang-tools-extra/test/clang-tidy/checkers/modernize/loop-convert-basic.cpp
@@ -980,3 +980,30 @@ namespace PR78381 {
     }
   }
 }
+
+namespace GH109083 {
+void test() {
+  const int N = 6;
+  int Arr[N] = {1, 2, 3, 4, 5, 6};
+
+  for (int I = 0; I < N; ++I) {
+    auto V = [T = Arr[I]]() {};
+  }
+  // CHECK-MESSAGES: :[[@LINE-3]]:3: warning: use range-based for loop instead [modernize-loop-convert]
+  // CHECK-FIXES: for (int I : Arr)
+  // CHECK-FIXES-NEXT: auto V = [T = I]() {};
+  for (int I = 0; I < N; ++I) {
+    auto V = [T = 10 + Arr[I]]() {};
+  }
+  // CHECK-MESSAGES: :[[@LINE-3]]:3: warning: use range-based for loop instead [modernize-loop-convert]
+  // CHECK-FIXES: for (int I : Arr)
+  // CHECK-FIXES-NEXT: auto V = [T = 10 + I]() {};
+
+  for (int I = 0; I < N; ++I) {
+    auto V = [T = I]() {};
+  }
+  for (int I = 0; I < N; ++I) {
+    auto V = [T = I + 10]() {};
+  }
+}
+} // namespace GH109083
\ No newline at end of file