Customize simplified dumping and matching of LambdaExpr

Reviewers: aaron.ballman

Subscribers: cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D71680

clang/include/clang/AST/ASTNodeTraverser.h

@@ -128,9 +128,12 @@ public:
       ConstStmtVisitor<Derived>::Visit(S);
 
       // Some statements have custom mechanisms for dumping their children.
-      if (isa<DeclStmt>(S) || isa<GenericSelectionExpr>(S)) {
+      if (isa<DeclStmt>(S) || isa<GenericSelectionExpr>(S))
+        return;
+
+      if (isa<LambdaExpr>(S) &&
+          Traversal == ast_type_traits::TK_IgnoreUnlessSpelledInSource)
         return;
-      }
 
       for (const Stmt *SubStmt : S->children())
         Visit(SubStmt);
@@ -646,7 +649,23 @@ public:
   }
 
   void VisitLambdaExpr(const LambdaExpr *Node) {
-    Visit(Node->getLambdaClass());
+    if (Traversal == ast_type_traits::TK_IgnoreUnlessSpelledInSource) {
+      for (unsigned I = 0, N = Node->capture_size(); I != N; ++I) {
+        const auto *C = Node->capture_begin() + I;
+        if (!C->isExplicit())
+          continue;
+        if (Node->isInitCapture(C))
+          Visit(C->getCapturedVar());
+        else
+          Visit(Node->capture_init_begin()[I]);
+      }
+      dumpTemplateParameters(Node->getTemplateParameterList());
+      for (const auto *P : Node->getCallOperator()->parameters())
+        Visit(P);
+      Visit(Node->getBody());
+    } else {
+      return Visit(Node->getLambdaClass());
+    }
   }
 
   void VisitSizeOfPackExpr(const SizeOfPackExpr *Node) {

clang/lib/ASTMatchers/ASTMatchFinder.cpp

@@ -138,20 +138,32 @@ public:
     ScopedIncrement ScopedDepth(&CurrentDepth);
     return (DeclNode == nullptr) || traverse(*DeclNode);
   }
+
+  Stmt *getStmtToTraverse(Stmt *StmtNode) {
+    Stmt *StmtToTraverse = StmtNode;
+    if (auto *ExprNode = dyn_cast_or_null<Expr>(StmtNode)) {
+      auto *LambdaNode = dyn_cast_or_null<LambdaExpr>(StmtNode);
+      if (LambdaNode && Finder->getASTContext().getTraversalKind() ==
+                          ast_type_traits::TK_IgnoreUnlessSpelledInSource)
+        StmtToTraverse = LambdaNode;
+      else
+        StmtToTraverse = Finder->getASTContext().traverseIgnored(ExprNode);
+    }
+    if (Traversal ==
+        ast_type_traits::TraversalKind::TK_IgnoreImplicitCastsAndParentheses) {
+      if (Expr *ExprNode = dyn_cast_or_null<Expr>(StmtNode))
+        StmtToTraverse = ExprNode->IgnoreParenImpCasts();
+    }
+    return StmtToTraverse;
+  }
+
   bool TraverseStmt(Stmt *StmtNode, DataRecursionQueue *Queue = nullptr) {
     // If we need to keep track of the depth, we can't perform data recursion.
     if (CurrentDepth == 0 || (CurrentDepth <= MaxDepth && MaxDepth < INT_MAX))
       Queue = nullptr;
 
     ScopedIncrement ScopedDepth(&CurrentDepth);
-    Stmt *StmtToTraverse = StmtNode;
+    Stmt *StmtToTraverse = getStmtToTraverse(StmtNode);
-    if (auto *ExprNode = dyn_cast_or_null<Expr>(StmtNode))
-      StmtToTraverse = Finder->getASTContext().traverseIgnored(ExprNode);
-    if (Traversal ==
-        ast_type_traits::TraversalKind::TK_IgnoreImplicitCastsAndParentheses) {
-      if (Expr *ExprNode = dyn_cast_or_null<Expr>(StmtNode))
-        StmtToTraverse = ExprNode->IgnoreParenImpCasts();
-    }
     if (!StmtToTraverse)
       return true;
     if (!match(*StmtToTraverse))
@@ -203,6 +215,41 @@ public:
     ScopedIncrement ScopedDepth(&CurrentDepth);
     return traverse(*CtorInit);
   }
+  bool TraverseLambdaExpr(LambdaExpr *Node) {
+    if (Finder->getASTContext().getTraversalKind() !=
+        ast_type_traits::TK_IgnoreUnlessSpelledInSource)
+      return VisitorBase::TraverseLambdaExpr(Node);
+    if (!Node)
+      return true;
+    ScopedIncrement ScopedDepth(&CurrentDepth);
+
+    for (unsigned I = 0, N = Node->capture_size(); I != N; ++I) {
+      const auto *C = Node->capture_begin() + I;
+      if (!C->isExplicit())
+        continue;
+      if (Node->isInitCapture(C) && !match(*C->getCapturedVar()))
+        return false;
+      if (!match(*Node->capture_init_begin()[I]))
+        return false;
+    }
+
+    if (const auto *TPL = Node->getTemplateParameterList()) {
+      for (const auto *TP : *TPL) {
+        if (!match(*TP))
+          return false;
+      }
+    }
+
+    for (const auto *P : Node->getCallOperator()->parameters()) {
+      if (!match(*P))
+        return false;
+    }
+
+    if (!match(*Node->getBody()))
+      return false;
+
+    return false;
+  }
 
   bool shouldVisitTemplateInstantiations() const { return true; }
   bool shouldVisitImplicitCode() const { return true; }

clang/unittests/AST/ASTTraverserTest.cpp

@@ -479,4 +479,123 @@ FunctionDecl 'func12'
 )cpp");
 }
 
+TEST(Traverse, LambdaUnlessSpelledInSource) {
+
+  auto AST =
+      buildASTFromCodeWithArgs(R"cpp(
+
+void captures() {
+  int a = 0;
+  int b = 0;
+  int d = 0;
+  int f = 0;
+
+  [a, &b, c = d, &e = f](int g, int h = 42) {};
+}
+
+void templated() {
+  int a = 0;
+  [a]<typename T>(T t) {};
+}
+
+struct SomeStruct {
+    int a = 0;
+    void capture_this() {
+        [this]() {};
+    }
+    void capture_this_copy() {
+        [self = *this]() {};
+    }
+};
+)cpp",
+                               {"-Wno-unused-value", "-Wno-c++2a-extensions"});
+
+  auto getLambdaNode = [&AST](const std::string &name) {
+    auto BN = ast_matchers::match(
+        lambdaExpr(hasAncestor(functionDecl(hasName(name)))).bind("lambda"),
+        AST->getASTContext());
+    EXPECT_EQ(BN.size(), 1u);
+    return BN[0].getNodeAs<LambdaExpr>("lambda");
+  };
+
+  {
+    auto L = getLambdaNode("captures");
+
+    EXPECT_EQ(dumpASTString(ast_type_traits::TK_IgnoreUnlessSpelledInSource, L),
+              R"cpp(
+LambdaExpr
+|-DeclRefExpr 'a'
+|-DeclRefExpr 'b'
+|-VarDecl 'c'
+| `-DeclRefExpr 'd'
+|-VarDecl 'e'
+| `-DeclRefExpr 'f'
+|-ParmVarDecl 'g'
+|-ParmVarDecl 'h'
+| `-IntegerLiteral
+`-CompoundStmt
+)cpp");
+
+    EXPECT_EQ(dumpASTString(ast_type_traits::TK_AsIs, L),
+              R"cpp(
+LambdaExpr
+|-CXXRecordDecl ''
+| |-CXXMethodDecl 'operator()'
+| | |-ParmVarDecl 'g'
+| | |-ParmVarDecl 'h'
+| | | `-IntegerLiteral
+| | `-CompoundStmt
+| |-FieldDecl ''
+| |-FieldDecl ''
+| |-FieldDecl ''
+| |-FieldDecl ''
+| `-CXXDestructorDecl '~'
+|-ImplicitCastExpr
+| `-DeclRefExpr 'a'
+|-DeclRefExpr 'b'
+|-ImplicitCastExpr
+| `-DeclRefExpr 'd'
+|-DeclRefExpr 'f'
+`-CompoundStmt
+)cpp");
+  }
+
+  {
+    auto L = getLambdaNode("templated");
+
+    EXPECT_EQ(dumpASTString(ast_type_traits::TK_IgnoreUnlessSpelledInSource, L),
+              R"cpp(
+LambdaExpr
+|-DeclRefExpr 'a'
+|-TemplateTypeParmDecl 'T'
+|-ParmVarDecl 't'
+`-CompoundStmt
+)cpp");
+  }
+
+  {
+    auto L = getLambdaNode("capture_this");
+
+    EXPECT_EQ(dumpASTString(ast_type_traits::TK_IgnoreUnlessSpelledInSource, L),
+              R"cpp(
+LambdaExpr
+|-CXXThisExpr
+`-CompoundStmt
+)cpp");
+  }
+
+  {
+    auto L = getLambdaNode("capture_this_copy");
+
+    EXPECT_EQ(dumpASTString(ast_type_traits::TK_IgnoreUnlessSpelledInSource, L),
+              R"cpp(
+LambdaExpr
+|-VarDecl 'self'
+| `-UnaryOperator
+|   `-CXXThisExpr
+`-CompoundStmt
+)cpp");
+  }
+}
+
 } // namespace clang

clang/unittests/ASTMatchers/ASTMatchersTraversalTest.cpp

@@ -1751,6 +1751,17 @@ B func12() {
   return c;
 }
 
+void func13() {
+  int a = 0;
+  int c = 0;
+
+  [a, b = c](int d) { int e = d; };
+}
+
+void func14() {
+  [] <typename TemplateType> (TemplateType t, TemplateType u) { int e = t + u; };
+}
+
 )cpp";
 
   EXPECT_TRUE(matches(
@@ -1821,6 +1832,23 @@ B func12() {
                      returnStmt(forFunction(functionDecl(hasName("func12"))),
                                 hasReturnValue(
                                     declRefExpr(to(varDecl(hasName("c")))))))));
+
+  EXPECT_TRUE(matches(
+      Code,
+      traverse(
+          ast_type_traits::TK_IgnoreUnlessSpelledInSource,
+          lambdaExpr(forFunction(functionDecl(hasName("func13"))),
+                     has(compoundStmt(hasDescendant(varDecl(hasName("e"))))),
+                     has(declRefExpr(to(varDecl(hasName("a"))))),
+                     has(varDecl(hasName("b"), hasInitializer(declRefExpr(to(
+                                                   varDecl(hasName("c"))))))),
+                     has(parmVarDecl(hasName("d")))))));
+
+  EXPECT_TRUE(matches(
+      Code, traverse(ast_type_traits::TK_IgnoreUnlessSpelledInSource,
+                     lambdaExpr(
+                         forFunction(functionDecl(hasName("func14"))),
+                         has(templateTypeParmDecl(hasName("TemplateType")))))));
 }
 
 TEST(IgnoringImpCasts, MatchesImpCasts) {