1 files changed, 48 insertions, 14 deletions
diff --git a/unittests/Tooling/RecursiveASTVisitorTests/LambdaExpr.cpp b/unittests/Tooling/RecursiveASTVisitorTests/LambdaExpr.cpp
index 80aeb43528..d3a3eba15d 100644
--- a/unittests/Tooling/RecursiveASTVisitorTests/LambdaExpr.cpp
+++ b/unittests/Tooling/RecursiveASTVisitorTests/LambdaExpr.cpp
@@ -17,25 +17,33 @@ namespace {
 class LambdaExprVisitor : public ExpectedLocationVisitor<LambdaExprVisitor> {
 public:
   bool VisitLambdaExpr(LambdaExpr *Lambda) {
-    PendingBodies.push(Lambda);
+    PendingBodies.push(Lambda->getBody());
+    PendingClasses.push(Lambda->getLambdaClass());
     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());
+  /// For each call to VisitLambdaExpr, we expect a subsequent call to visit
+  /// the body (and maybe the lambda class, which is implicit).
+  bool VisitStmt(Stmt *S) {
+    if (!PendingBodies.empty() && S == PendingBodies.top())
+      PendingBodies.pop();
+    return true;
   }
-  /// Determine whether TraverseLambdaBody has been called for every call to
-  /// VisitLambdaExpr.
-  bool allBodiesHaveBeenTraversed() const {
-    return PendingBodies.empty();
+  bool VisitDecl(Decl *D) {
+    if (!PendingClasses.empty() && D == PendingClasses.top())
+      PendingClasses.pop();
+    return true;
   }
+  /// Determine whether parts of lambdas (VisitLambdaExpr) were later traversed.
+  bool allBodiesHaveBeenTraversed() const { return PendingBodies.empty(); }
+  bool allClassesHaveBeenTraversed() const { return PendingClasses.empty(); }
+
+  bool VisitImplicitCode = false;
+  bool shouldVisitImplicitCode() const { return VisitImplicitCode; }
+
 private:
-  std::stack<LambdaExpr *> PendingBodies;
+  std::stack<Stmt *> PendingBodies;
+  std::stack<Decl *> PendingClasses;
 };
 
 TEST(RecursiveASTVisitor, VisitsLambdaExpr) {
@@ -43,13 +51,39 @@ TEST(RecursiveASTVisitor, VisitsLambdaExpr) {
   Visitor.ExpectMatch("", 1, 12);
   EXPECT_TRUE(Visitor.runOver("void f() { []{ return; }(); }",
                               LambdaExprVisitor::Lang_CXX11));
+  EXPECT_TRUE(Visitor.allBodiesHaveBeenTraversed());
+  EXPECT_FALSE(Visitor.allClassesHaveBeenTraversed());
 }
 
-TEST(RecursiveASTVisitor, TraverseLambdaBodyCanBeOverridden) {
+TEST(RecursiveASTVisitor, LambdaInLambda) {
   LambdaExprVisitor Visitor;
+  Visitor.ExpectMatch("", 1, 12);
+  Visitor.ExpectMatch("", 1, 16);
+  EXPECT_TRUE(Visitor.runOver("void f() { []{ []{ return; }; }(); }",
+                              LambdaExprVisitor::Lang_CXX11));
+  EXPECT_TRUE(Visitor.allBodiesHaveBeenTraversed());
+  EXPECT_FALSE(Visitor.allClassesHaveBeenTraversed());
+}
+
+TEST(RecursiveASTVisitor, TopLevelLambda) {
+  LambdaExprVisitor Visitor;
+  Visitor.VisitImplicitCode = true;
+  Visitor.ExpectMatch("", 1, 10);
+  Visitor.ExpectMatch("", 1, 14);
+  EXPECT_TRUE(Visitor.runOver("auto x = []{ [] {}; };",
+                              LambdaExprVisitor::Lang_CXX11));
+  EXPECT_TRUE(Visitor.allBodiesHaveBeenTraversed());
+  EXPECT_TRUE(Visitor.allClassesHaveBeenTraversed());
+}
+
+TEST(RecursiveASTVisitor, VisitsLambdaExprAndImplicitClass) {
+  LambdaExprVisitor Visitor;
+  Visitor.VisitImplicitCode = true;
+  Visitor.ExpectMatch("", 1, 12);
   EXPECT_TRUE(Visitor.runOver("void f() { []{ return; }(); }",
                               LambdaExprVisitor::Lang_CXX11));
   EXPECT_TRUE(Visitor.allBodiesHaveBeenTraversed());
+  EXPECT_TRUE(Visitor.allClassesHaveBeenTraversed());
 }
 
 TEST(RecursiveASTVisitor, VisitsAttributedLambdaExpr) {