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

1 files changed, 25 insertions, 0 deletions

diff --git a/unittests/Tooling/RecursiveASTVisitorTest.cpp b/unittests/Tooling/RecursiveASTVisitorTest.cpp
index 2d226fddfc..c97ee0c8c4 100644
--- 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