Add a hook RecursiveASTVisitor::TraverseLambdaBody, to enable visitors to

use/maintain additional state from the LambdaExpr while visiting the body
of a LambdaExpr.

One use for this arises because Clang's AST currently holds lambda bodies
in a form prior to their adjustment to refer to captured copies of local
variables, and so some clients will need access to the lambda's closure
type in order to query how to map VarDecl*s to the FieldDecls of their
by-copy captures. This hook is sufficient for at least one such client;
to do this without such a hook would require the client to re-implement
the whole of TraverseLambdaExpr, which is non-trivial and would likely be
more brittle.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@186024 91177308-0d34-0410-b5e6-96231b3b80d8

diff --git a/include/clang/AST/RecursiveASTVisitor.h b/include/clang/AST/RecursiveASTVisitor.h
index 09ac4766968c73f252ee7378ef222cb02d1a7840..6af9e55b8366bce4c75cd7b060db3770a90997d7 100644 (file)
--- a/include/clang/AST/RecursiveASTVisitor.h
+++ b/include/clang/AST/RecursiveASTVisitor.h
@@ -244,7 +244,15 @@ public:
   ///
   /// \returns false if the visitation was terminated early, true otherwise.
   bool TraverseLambdaCapture(LambdaExpr *LE, const LambdaExpr::Capture *C);
-  
+
+  /// \brief Recursively visit the body of a lambda expression.
+  ///
+  /// This provides a hook for visitors that need more context when visiting
+  /// \c LE->getBody().
+  ///
+  /// \returns false if the visitation was terminated early, true otherwise.
+  bool TraverseLambdaBody(LambdaExpr *LE);
+
   // ---- Methods on Stmts ----
 
   // Declare Traverse*() for all concrete Stmt classes.
@@ -809,6 +817,13 @@ bool RecursiveASTVisitor<Derived>::TraverseLambdaCapture(
   return true;
 }
 
+template<typename Derived>
+bool RecursiveASTVisitor<Derived>::TraverseLambdaBody(LambdaExpr *LE) {
+  TRY_TO(TraverseStmt(LE->getBody()));
+  return true;
+}
+
+
 // ----------------- Type traversal -----------------
 
 // This macro makes available a variable T, the passed-in type.
@@ -2153,7 +2168,7 @@ bool RecursiveASTVisitor<Derived>::TraverseLambdaExpr(LambdaExpr *S) {
     }
   }
 
-  TRY_TO(TraverseStmt(S->getBody()));
+  TRY_TO(TraverseLambdaBody(S));
   return true;
 }
 

diff --git a/unittests/Tooling/RecursiveASTVisitorTest.cpp b/unittests/Tooling/RecursiveASTVisitorTest.cpp
index 2d226fddfc4743203baee93a05ebeb3c97877aa5..c97ee0c8c4764825ba0e01dbd98142404badf30b 100644 (file)
--- a/unittests/Tooling/RecursiveASTVisitorTest.cpp
+++ b/unittests/Tooling/RecursiveASTVisitorTest.cpp
@@ -9,6 +9,8 @@
 
 #include "TestVisitor.h"
 
+#include <stack>
+
 namespace clang {
 
 class TypeLocVisitor : public ExpectedLocationVisitor<TypeLocVisitor> {
@@ -82,9 +84,25 @@ public:
 class LambdaExprVisitor : public ExpectedLocationVisitor<LambdaExprVisitor> {
 public:
   bool VisitLambdaExpr(LambdaExpr *Lambda) {
+    PendingBodies.push(Lambda);
     Match("", Lambda->getIntroducerRange().getBegin());
     return true;
   }
+  /// For each call to VisitLambdaExpr, we expect a subsequent call (with
+  /// proper nesting) to TraverseLambdaBody.
+  bool TraverseLambdaBody(LambdaExpr *Lambda) {
+    EXPECT_FALSE(PendingBodies.empty());
+    EXPECT_EQ(PendingBodies.top(), Lambda);
+    PendingBodies.pop();
+    return TraverseStmt(Lambda->getBody());
+  }
+  /// Determine whether TraverseLambdaBody has been called for every call to
+  /// VisitLambdaExpr.
+  bool allBodiesHaveBeenTraversed() const {
+    return PendingBodies.empty();
+  }
+private:
+  std::stack<LambdaExpr *> PendingBodies; 
 };
 
 class TemplateArgumentLocTraverser
@@ -478,4 +496,11 @@ TEST(RecursiveASTVisitor, VisitsLambdaExpr) {
                              LambdaExprVisitor::Lang_CXX11));
 }
 
+TEST(RecursiveASTVisitor, TraverseLambdaBodyCanBeOverridden) {
+  LambdaExprVisitor Visitor;
+  EXPECT_TRUE(Visitor.runOver("void f() { []{ return; }(); }",
+                             LambdaExprVisitor::Lang_CXX11));
+  EXPECT_TRUE(Visitor.allBodiesHaveBeenTraversed());
+}
+
 } // end namespace clang