[clang][dataflow] Add captured parameters to ReferencedDecls for lamb…

clang/include/clang/Analysis/FlowSensitive/ASTOps.h

@@ -146,6 +146,10 @@ struct ReferencedDecls {
   /// Free functions and member functions which are referenced (but not
   /// necessarily called).
   llvm::DenseSet<const FunctionDecl *> Functions;
+  /// When analyzing a lambda's call operator, the set of all parameters (from
+  /// the surrounding function) that the lambda captures. Captured local
+  /// variables are already included in `Locals` above.
+  llvm::DenseSet<const ParmVarDecl *> LambdaCapturedParams;
 };
 
 /// Returns declarations that are declared in or referenced from `FD`.

clang/lib/Analysis/FlowSensitive/ASTOps.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/Analysis/FlowSensitive/ASTOps.h"
+#include "clang/AST/ASTLambda.h"
 #include "clang/AST/ComputeDependence.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
@@ -282,6 +283,28 @@ ReferencedDecls getReferencedDecls(const FunctionDecl &FD) {
   if (const auto *CtorDecl = dyn_cast<CXXConstructorDecl>(&FD))
     Visitor.traverseConstructorInits(CtorDecl);
 
+  // If analyzing a lambda call operator, collect all captures of parameters (of
+  // the surrounding function). This collects them even if they are not
+  // referenced in the body of the lambda call operator. Non-parameter local
+  // variables that are captured are already collected into
+  // `ReferencedDecls.Locals` when traversing the call operator body, but we
+  // collect parameters here to avoid needing to check at each referencing node
+  // whether the parameter is a lambda capture from a surrounding function or is
+  // a parameter of the current function. If it becomes necessary to limit this
+  // set to the parameters actually referenced in the body, alternative
+  // optimizations can be implemented to minimize duplicative work.
+  if (const auto *Method = dyn_cast<CXXMethodDecl>(&FD);
+      Method && isLambdaCallOperator(Method)) {
+    for (const auto &Capture : Method->getParent()->captures()) {
+      if (Capture.capturesVariable()) {
+        if (const auto *Param =
+                dyn_cast<ParmVarDecl>(Capture.getCapturedVar())) {
+          Result.LambdaCapturedParams.insert(Param);
+        }
+      }
+    }
+  }
+
   return Result;
 }
 

clang/unittests/Analysis/FlowSensitive/ASTOpsTest.cpp

@@ -8,15 +8,23 @@
 
 #include "clang/Analysis/FlowSensitive/ASTOps.h"
 #include "TestingSupport.h"
+#include "clang/AST/Decl.h"
+#include "clang/AST/DeclCXX.h"
+#include "clang/ASTMatchers/ASTMatchers.h"
+#include "clang/Basic/LLVM.h"
+#include "clang/Frontend/ASTUnit.h"
+#include "clang/Tooling/Tooling.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <memory>
+#include <string>
 
 namespace {
 
 using namespace clang;
 using namespace dataflow;
 
+using ast_matchers::cxxMethodDecl;
 using ast_matchers::cxxRecordDecl;
 using ast_matchers::hasName;
 using ast_matchers::hasType;
@@ -107,4 +115,37 @@ TEST(ASTOpsTest, ReferencedDeclsLocalsNotParamsOrStatics) {
               UnorderedElementsAre(LocalDecl));
 }
 
+TEST(ASTOpsTest, LambdaCaptures) {
+  std::string Code = R"cc(
+    void func(int CapturedByRef, int CapturedByValue, int NotCaptured) {
+    int Local;
+      auto Lambda = [&CapturedByRef, CapturedByValue, &Local](int LambdaParam) {
+      };
+    }
+  )cc";
+  std::unique_ptr<ASTUnit> Unit =
+      tooling::buildASTFromCodeWithArgs(Code, {"-fsyntax-only", "-std=c++17"});
+  auto &ASTCtx = Unit->getASTContext();
+
+  ASSERT_EQ(ASTCtx.getDiagnostics().getClient()->getNumErrors(), 0U);
+
+  auto *LambdaCallOp = selectFirst<CXXMethodDecl>(
+      "l", match(cxxMethodDecl(hasName("operator()")).bind("l"), ASTCtx));
+  ASSERT_NE(LambdaCallOp, nullptr);
+  auto *Func = cast<FunctionDecl>(findValueDecl(ASTCtx, "func"));
+  ASSERT_NE(Func, nullptr);
+  auto *CapturedByRefDecl = Func->getParamDecl(0);
+  ASSERT_NE(CapturedByRefDecl, nullptr);
+  auto *CapturedByValueDecl = Func->getParamDecl(1);
+  ASSERT_NE(CapturedByValueDecl, nullptr);
+
+  EXPECT_THAT(getReferencedDecls(*Func).LambdaCapturedParams, IsEmpty());
+  ReferencedDecls ForLambda = getReferencedDecls(*LambdaCallOp);
+  EXPECT_THAT(ForLambda.LambdaCapturedParams,
+              UnorderedElementsAre(CapturedByRefDecl, CapturedByValueDecl));
+  // Captured locals must be seen in the body for them to appear in
+  // ReferencedDecls.
+  EXPECT_THAT(ForLambda.Locals, IsEmpty());
+}
+
 } // namespace