From 86e67a2fc40860f081f87537bcf55c2faaf1b942 Mon Sep 17 00:00:00 2001
From: James Dennett <jdennett@google.com>
Date: Wed, 10 Jul 2013 18:29:15 +0000
Subject: [PATCH] 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.

llvm-svn: 186024
---
 clang/include/clang/AST/RecursiveASTVisitor.h | 19 ++++++++++++--
 .../Tooling/RecursiveASTVisitorTest.cpp       | 25 +++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/clang/include/clang/AST/RecursiveASTVisitor.h b/clang/include/clang/AST/RecursiveASTVisitor.h
index 09ac4766968c..6af9e55b8366 100644
--- a/clang/include/clang/AST/RecursiveASTVisitor.h
+++ b/clang/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/clang/unittests/Tooling/RecursiveASTVisitorTest.cpp b/clang/unittests/Tooling/RecursiveASTVisitorTest.cpp
index 2d226fddfc47..c97ee0c8c476 100644
--- a/clang/unittests/Tooling/RecursiveASTVisitorTest.cpp
+++ b/clang/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