[AST] RecursiveASTVisitor visits lambda classes when implicit visitation is on.

Summary:
This fixes ASTContext's parent map for nodes in such classes (e.g. operator()).
https://bugs.llvm.org/show_bug.cgi?id=39949

This also changes the observed shape of the AST for implicit RAVs.
- this includes AST MatchFinder: cxxRecordDecl() now matches lambda classes,
functionDecl() matches the call operator, and the parent chain is body -> call
operator -> lambda class -> lambdaexpr rather than body -> lambdaexpr.
- this appears not to matter for the ASTImporterLookupTable builder
- this doesn't matter for the other RAVs in-tree.

In order to do this, we remove the TraverseLambdaBody hook. The problem is it's
hard/weird to ensure this hook is called when traversing via the implicit class.
There were just two users of this hook in-tree, who use it to skip bodies.
I replaced these with explicitly traversing the captures only. Another approach
would be recording the bodies when the lambda is visited, and then recognizing
them later.
I'd be open to suggestion on how to preserve this hook, instead.

Reviewers: aaron.ballman, JonasToth

Subscribers: cfe-commits, rsmith, jdennett

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

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

5 files changed, 82 insertions, 46 deletions

diff --git a/include/clang/AST/RecursiveASTVisitor.h b/include/clang/AST/RecursiveASTVisitor.h
index 0d88517010..96f3d24c6d 100644
--- a/include/clang/AST/RecursiveASTVisitor.h
+++ b/include/clang/AST/RecursiveASTVisitor.h
@@ -298,14 +298,6 @@ public:
   bool TraverseLambdaCapture(LambdaExpr *LE, const LambdaCapture *C,
                              Expr *Init);
 
-  /// 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, DataRecursionQueue *Queue = nullptr);
-
   /// Recursively visit the syntactic or semantic form of an
   /// initialization list.
   ///
@@ -936,13 +928,6 @@ RecursiveASTVisitor<Derived>::TraverseLambdaCapture(LambdaExpr *LE,
   return true;
 }
 
-template <typename Derived>
-bool RecursiveASTVisitor<Derived>::TraverseLambdaBody(
-    LambdaExpr *LE, DataRecursionQueue *Queue) {
-  TRY_TO_TRAVERSE_OR_ENQUEUE_STMT(LE->getBody());
-  return true;
-}
-
 // ----------------- Type traversal -----------------
 
 // This macro makes available a variable T, the passed-in type.
@@ -2404,6 +2389,7 @@ DEF_TRAVERSE_STMT(CXXTemporaryObjectExpr, {
 
 // Walk only the visible parts of lambda expressions.
 DEF_TRAVERSE_STMT(LambdaExpr, {
+  // Visit the capture list.
   for (unsigned I = 0, N = S->capture_size(); I != N; ++I) {
     const LambdaCapture *C = S->capture_begin() + I;
     if (C->isExplicit() || getDerived().shouldVisitImplicitCode()) {
@@ -2411,25 +2397,31 @@ DEF_TRAVERSE_STMT(LambdaExpr, {
     }
   }
 
-  TypeLoc TL = S->getCallOperator()->getTypeSourceInfo()->getTypeLoc();
-  FunctionProtoTypeLoc Proto = TL.getAsAdjusted<FunctionProtoTypeLoc>();
+  if (getDerived().shouldVisitImplicitCode()) {
+    // The implicit model is simple: everything else is in the lambda class.
+    TRY_TO(TraverseDecl(S->getLambdaClass()));
+  } else {
+    // We need to poke around to find the bits that might be explicitly written.
+    TypeLoc TL = S->getCallOperator()->getTypeSourceInfo()->getTypeLoc();
+    FunctionProtoTypeLoc Proto = TL.getAsAdjusted<FunctionProtoTypeLoc>();
 
-  if (S->hasExplicitParameters()) {
-    // Visit parameters.
-    for (unsigned I = 0, N = Proto.getNumParams(); I != N; ++I)
-      TRY_TO(TraverseDecl(Proto.getParam(I)));
-  }
-  if (S->hasExplicitResultType())
-    TRY_TO(TraverseTypeLoc(Proto.getReturnLoc()));
+    if (S->hasExplicitParameters()) {
+      // Visit parameters.
+      for (unsigned I = 0, N = Proto.getNumParams(); I != N; ++I)
+        TRY_TO(TraverseDecl(Proto.getParam(I)));
+    }
+    if (S->hasExplicitResultType())
+      TRY_TO(TraverseTypeLoc(Proto.getReturnLoc()));
 
-  auto *T = Proto.getTypePtr();
-  for (const auto &E : T->exceptions())
-    TRY_TO(TraverseType(E));
+    auto *T = Proto.getTypePtr();
+    for (const auto &E : T->exceptions())
+      TRY_TO(TraverseType(E));
 
-  if (Expr *NE = T->getNoexceptExpr())
-    TRY_TO_TRAVERSE_OR_ENQUEUE_STMT(NE);
+    if (Expr *NE = T->getNoexceptExpr())
+      TRY_TO_TRAVERSE_OR_ENQUEUE_STMT(NE);
 
-  ReturnValue = TRAVERSE_STMT_BASE(LambdaBody, LambdaExpr, S, Queue);
+    TRY_TO_TRAVERSE_OR_ENQUEUE_STMT(S->getBody());
+  }
   ShouldVisitChildren = false;
 })
 
diff --git a/lib/CodeGen/CodeGenPGO.cpp b/lib/CodeGen/CodeGenPGO.cpp
index 48900ac3e8..776060743a 100644
--- a/lib/CodeGen/CodeGenPGO.cpp
+++ b/lib/CodeGen/CodeGenPGO.cpp
@@ -165,7 +165,12 @@ struct MapRegionCounters : public RecursiveASTVisitor<MapRegionCounters> {
   // Blocks and lambdas are handled as separate functions, so we need not
   // traverse them in the parent context.
   bool TraverseBlockExpr(BlockExpr *BE) { return true; }
-  bool TraverseLambdaBody(LambdaExpr *LE) { return true; }
+  bool TraverseLambdaExpr(LambdaExpr *LE) {
+    // Traverse the captures, but not the body.
+    for (const auto &C : zip(LE->captures(), LE->capture_inits()))
+      TraverseLambdaCapture(LE, &std::get<0>(C), std::get<1>(C));
+    return true;
+  }
   bool TraverseCapturedStmt(CapturedStmt *CS) { return true; }
 
   bool VisitDecl(const Decl *D) {
diff --git a/lib/Sema/AnalysisBasedWarnings.cpp b/lib/Sema/AnalysisBasedWarnings.cpp
index 3b6cbe9469..c818d40c77 100644
--- a/lib/Sema/AnalysisBasedWarnings.cpp
+++ b/lib/Sema/AnalysisBasedWarnings.cpp
@@ -1153,7 +1153,12 @@ namespace {
     bool TraverseDecl(Decl *D) { return true; }
 
     // We analyze lambda bodies separately. Skip them here.
-    bool TraverseLambdaBody(LambdaExpr *LE) { return true; }
+    bool TraverseLambdaExpr(LambdaExpr *LE) {
+      // Traverse the captures, but not the body.
+      for (const auto &C : zip(LE->captures(), LE->capture_inits()))
+        TraverseLambdaCapture(LE, &std::get<0>(C), std::get<1>(C));
+      return true;
+    }
 
   private:
 
diff --git a/unittests/AST/ASTContextParentMapTest.cpp b/unittests/AST/ASTContextParentMapTest.cpp
index f06f32bf76..fb9d517069 100644
--- a/unittests/AST/ASTContextParentMapTest.cpp
+++ b/unittests/AST/ASTContextParentMapTest.cpp
@@ -106,5 +106,16 @@ TEST(GetParents, RespectsTraversalScope) {
   EXPECT_THAT(Ctx.getParents(Foo), ElementsAre(DynTypedNode::create(TU)));
 }
 
+TEST(GetParents, ImplicitLambdaNodes) {
+  MatchVerifier<Decl> LambdaVerifier;
+  EXPECT_TRUE(LambdaVerifier.match(
+      "auto x = []{int y;};",
+      varDecl(hasName("y"), hasAncestor(functionDecl(
+                                hasOverloadedOperatorName("()"),
+                                hasParent(cxxRecordDecl(
+                                    isImplicit(), hasParent(lambdaExpr())))))),
+      Lang_CXX11));
+}
+
 } // end namespace ast_matchers
 } // end namespace clang
diff --git a/unittests/Tooling/RecursiveASTVisitorTests/LambdaExpr.cpp b/unittests/Tooling/RecursiveASTVisitorTests/LambdaExpr.cpp
index 80aeb43528..a965632153 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,28 @@ 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, 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, TraverseLambdaBodyCanBeOverridden) {
+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) {