[AST] Include the TranslationUnitDecl when traversing with TraversalScope

[sammccall]

Summary

Given int foo, bar;, TraverseAST reveals this tree:

TranslationUnitDecl
 - foo
 - bar

Before this patch, with the TraversalScope set to {foo}, TraverseAST yields:

foo

After this patch it yields:

TranslationUnitDecl
- foo

Also, TraverseDecl(TranslationUnitDecl) now respects the traversal scope.

---

The main effect of this today is that clang-tidy checks that match the
translationUnitDecl(), either in order to traverse it or check
parentage, should work.

[sammccall]

This includes both the traversal change and the clangd changes to adapt to it, since there ended up being few.

I chickened out of the idea of changing the behavior of TranslationUnitDecl::decls(), for a couple of reasons:

• having it apply to RecursiveASTVisitor only is similar to the previous version, so this is a slightly less invasive/risky change
• it's actually slightly awkward to change the node's behavior as decls() is not virtual in DeclContext

[hokein]

In D104071#2811827, @sammccall wrote:

This includes both the traversal change and the clangd changes to adapt to it, since there ended up being few.

I chickened out of the idea of changing the behavior of TranslationUnitDecl::decls(), for a couple of reasons:

• having it apply to RecursiveASTVisitor only is similar to the previous version, so this is a slightly less invasive/risky change
• it's actually slightly awkward to change the node's behavior as decls() is not virtual in DeclContext

Thanks, this seems a nice approach considering all trade-offs.

IIUC, the key point of the approach is that we change the RAV behavior on traversing TUDecl, to make it respect decls from traversal scope, which seems a reasonable mental model.
I think we should mention this in the patch description, and document somewhere around the RAV?

While it solves the critical problem in clang-tidy skip-headers, we should bear in mind that:

1. The default TU-decl behavior in RAV can be overrode in RAV's subclass (I bet it should be very rare, and we probably don't worry about this case )
2. users can still use "TU->decls()" to traverse all decls, which might cause clangd to deserialize preamble and get performance hit (this is an existing problem)

All in all, I'm up for this patch.

clang/include/clang/AST/RecursiveASTVisitor.h

1502

I would use a more meaningful name -- it is unclear to me what "trivial" means here without reading the code very hard.

IIUC it distinguishes the conventional case ( a single translation-unit decl) vs the case (where e.g. the TraversalScope is set to main-file decls).

[sammccall]

Expand/clarify comments and variable names

[hokein]

Per https://reviews.llvm.org/harbormaster/unit/view/779496/, window test seems to be failing...

Thanks! I think we should probably document this RAV behavior change in the release note.

[sammccall]

In D104071#2812870, @hokein wrote:

Per https://reviews.llvm.org/harbormaster/unit/view/779496/, window test seems to be failing...

Oops, there was a test I forgot to update (failed on all platforms). Fixed.

Thanks! I think we should probably document this RAV behavior change in the release note.

Happy to do this if you like, but I wouldn't be inclined - there are no in-tree uses of setTraversalScope apart from clangd, and in google's internal codebase (quite a lot of clang API usage) I only found one caller and no changes were needed.

clang-tools-extra/clangd/DumpAST.cpp

public:
   ASTNode Root;
   DumpVisitor(const syntax::TokenBuffer &Tokens, const ASTContext &Ctx)
       : Tokens(Tokens), Ctx(Ctx) {}
 
   // Override traversal to record the nodes we care about.
   // Generally, these are nodes with position information (TypeLoc, not Type).
 
   bool TraverseDecl(Decl *D) {
-    return !D || isInjectedClassName(D) || traverseNode("declaration", D, [&] {
-      if (isa<TranslationUnitDecl>(D))
-        Base::TraverseAST(const_cast<ASTContext &>(Ctx));
-      else
-        Base::TraverseDecl(D);
-    });
+    return !D || isInjectedClassName(D) ||
+           traverseNode("declaration", D, [&] { Base::TraverseDecl(D); });
   }
   bool TraverseTypeLoc(TypeLoc TL) {
     return !TL || traverseNode("type", TL, [&] { Base::TraverseTypeLoc(TL); });
   }
   bool TraverseTemplateName(const TemplateName &TN) {
     return traverseNode("template name", TN,
                         [&] { Base::TraverseTemplateName(TN); });
   }

clang-tools-extra/clangd/refactor/tweaks/AddUsing.cpp

bool VisitUsingDecl(UsingDecl *D) {
     }
     return true;
   }
 
   bool TraverseDecl(Decl *Node) {
     // There is no need to go deeper into nodes that do not enclose selection,
     // since "using" there will not affect selection, nor would it make a good
     // insertion point.
-    if (Node->getDeclContext()->Encloses(SelectionDeclContext)) {
+    if (!Node->getDeclContext() ||
+        Node->getDeclContext()->Encloses(SelectionDeclContext)) {
       return RecursiveASTVisitor<UsingFinder>::TraverseDecl(Node);
     }
     return true;
   }
 
 private:
   std::vector<const UsingDecl *> &Results;
   const DeclContext *SelectionDeclContext;

clang-tools-extra/clangd/unittests/DiagnosticsTests.cpp

Annotations Test(R"cpp(
     #include $deprecated[["assert.h"]]
 
     #define $macrodef[[SQUARE]](X) (X)*(X)
     int $main[[main]]() {
       int y = 4;
       return SQUARE($macroarg[[++]]y);
       return $doubled[[sizeof]](sizeof(int));
     }
+
+    // misc-no-recursion uses a custom traversal from the TUDecl
+    void foo();
+    void $bar[[bar]]() {
+      foo();
+    }
+    void $foo[[foo]]() {
+      bar();
+    }
   )cpp");
   auto TU = TestTU::withCode(Test.code());
   TU.HeaderFilename = "assert.h"; // Suppress "not found" error.
   TU.ClangTidyProvider = addTidyChecks("bugprone-sizeof-expression,"
                                        "bugprone-macro-repeated-side-effects,"
                                        "modernize-deprecated-headers,"
-                                       "modernize-use-trailing-return-type");
+                                       "modernize-use-trailing-return-type,"
+                                       "misc-no-recursion");
   EXPECT_THAT(
       *TU.build().getDiagnostics(),
       UnorderedElementsAre(
           AllOf(Diag(Test.range("deprecated"),
                      "inclusion of deprecated C++ header 'assert.h'; consider "
                      "using 'cassert' instead"),
                 DiagSource(Diag::ClangTidy),
                 DiagName("modernize-deprecated-headers"),
EXPECT_THAT(
           Diag(Test.range("macroarg"),
                "multiple unsequenced modifications to 'y'"),
           AllOf(
               Diag(Test.range("main"),
                    "use a trailing return type for this function"),
               DiagSource(Diag::ClangTidy),
               DiagName("modernize-use-trailing-return-type"),
               // Verify that we don't have "[check-name]" suffix in the message.
-              WithFix(FixMessage(
-                  "use a trailing return type for this function")))));
+              WithFix(
+                  FixMessage("use a trailing return type for this function"))),
+          Diag(Test.range("foo"),
+               "function 'foo' is within a recursive call chain"),
+          Diag(Test.range("bar"),
+               "function 'bar' is within a recursive call chain")));
 }
 
 TEST(DiagnosticsTest, ClangTidyEOF) {
   // clang-format off
   Annotations Test(R"cpp(
   [[#]]include <b.h>
   #include "a.h")cpp");
   // clang-format on

clang-tools-extra/clangd/unittests/TestTU.cpp

const Symbol &findSymbol(const SymbolSlab &Slab, llvm::StringRef QName) {
   if (!Result) {
     llvm::errs() << "No symbol named " << QName << " in "
                  << llvm::to_string(Slab);
     assert(false && "No symbol with QName");
   }
   return *Result;
 }
 
+// RAII scoped class to disable TraversalScope for a ParsedAST.
+class TraverseHeadersToo {
+  ASTContext &Ctx;
+  std::vector<Decl *> ScopeToRestore;
+
+public:
+  TraverseHeadersToo(ParsedAST &AST)
+      : Ctx(AST.getASTContext()), ScopeToRestore(Ctx.getTraversalScope()) {
+    Ctx.setTraversalScope({Ctx.getTranslationUnitDecl()});
+  }
+  ~TraverseHeadersToo() { Ctx.setTraversalScope(std::move(ScopeToRestore)); }
+};
+
 const NamedDecl &findDecl(ParsedAST &AST, llvm::StringRef QName) {
   auto &Ctx = AST.getASTContext();
   auto LookupDecl = [&Ctx](const DeclContext &Scope,
                            llvm::StringRef Name) -> const NamedDecl & {
     auto LookupRes = Scope.lookup(DeclarationName(&Ctx.Idents.get(Name)));
     assert(!LookupRes.empty() && "Lookup failed");
     assert(LookupRes.isSingleResult() && "Lookup returned multiple results");
     return *LookupRes.front();
   };
 
   const DeclContext *Scope = Ctx.getTranslationUnitDecl();
 
   StringRef Cur, Rest;
   for (std::tie(Cur, Rest) = QName.split("::"); !Rest.empty();
        std::tie(Cur, Rest) = Rest.split("::")) {
     Scope = &cast<DeclContext>(LookupDecl(*Scope, Cur));
   }
   return LookupDecl(*Scope, Cur);
 }
 
 const NamedDecl &findDecl(ParsedAST &AST,
                           std::function<bool(const NamedDecl &)> Filter) {
+  TraverseHeadersToo Too(AST);
   struct Visitor : RecursiveASTVisitor<Visitor> {
     decltype(Filter) F;
     llvm::SmallVector<const NamedDecl *, 1> Decls;
     bool VisitNamedDecl(const NamedDecl *ND) {
       if (F(*ND))
         Decls.push_back(ND);
       return true;
     }

clang/include/clang/AST/ASTContext.h

public:
 
   /// Returns the clang bytecode interpreter context.
   interp::Context &getInterpContext();
 
   /// Returns the dynamic AST node parent map context.
   ParentMapContext &getParentMapContext();
 
   // A traversal scope limits the parts of the AST visible to certain analyses.
-  // RecursiveASTVisitor::TraverseAST will only visit reachable nodes, and
+  // RecursiveASTVisitor only visits specified children of TranslationUnitDecl.
   // getParents() will only observe reachable parent edges.
   //
-  // The scope is defined by a set of "top-level" declarations.
-  // Initially, it is the entire TU: {getTranslationUnitDecl()}.
+  // The scope is defined by a set of "top-level" declarations which will be
+  // visible under the TranslationUnitDecl.
+  // Initially, it is the entire TU, represented by {getTranslationUnitDecl()}.
+  //
+  // After setTraversalScope({foo, bar}), the exposed AST looks like:
+  // TranslationUnitDecl
+  //  - foo
+  //    - ...
+  //  - bar
+  //    - ...
+  // All other siblings of foo and bar are pruned from the tree.
+  // (However they are still accessible via TranslationUnitDecl->decls())
+  //
   // Changing the scope clears the parent cache, which is expensive to rebuild.
   std::vector<Decl *> getTraversalScope() const { return TraversalScope; }
   void setTraversalScope(const std::vector<Decl *> &);
 
   /// Forwards to get node parents from the ParentMapContext. New callers should
   /// use ParentMapContext::getParents() directly.
   template <typename NodeT> DynTypedNodeList getParents(const NodeT &Node);
 

clang/include/clang/AST/RecursiveASTVisitor.h

public:
   bool shouldVisitImplicitCode() const { return false; }
 
   /// Return whether this visitor should recurse into lambda body
   bool shouldVisitLambdaBody() const { return true; }
 
   /// Return whether this visitor should traverse post-order.
   bool shouldTraversePostOrder() const { return false; }
 
-  /// Recursively visits an entire AST, starting from the top-level Decls
-  /// in the AST traversal scope (by default, the TranslationUnitDecl).
+  /// Recursively visits an entire AST, starting from the TranslationUnitDecl.
   /// \returns false if visitation was terminated early.
   bool TraverseAST(ASTContext &AST) {
-    for (Decl *D : AST.getTraversalScope())
-      if (!getDerived().TraverseDecl(D))
-        return false;
-    return true;
+    // Currently just an alias for TraverseDecl(TUDecl), but kept in case
+    // we change the implementation again.
+    return getDerived().TraverseDecl(AST.getTranslationUnitDecl());
   }
 
   /// Recursively visit a statement or expression, by
   /// dispatching to Traverse*() based on the argument's dynamic type.
   ///
   /// \returns false if the visitation was terminated early, true
   /// otherwise (including when the argument is nullptr).
   bool TraverseStmt(Stmt *S, DataRecursionQueue *Queue = nullptr);

 DEF_TRAVERSE_DECL(ObjCPropertyImplDecl, {// FIXME: implement this
                                         })
 
 DEF_TRAVERSE_DECL(StaticAssertDecl, {
   TRY_TO(TraverseStmt(D->getAssertExpr()));
   TRY_TO(TraverseStmt(D->getMessage()));
 })
 
-DEF_TRAVERSE_DECL(
-    TranslationUnitDecl,
-    {// Code in an unnamed namespace shows up automatically in
-     // decls_begin()/decls_end().  Thus we don't need to recurse on
-     // D->getAnonymousNamespace().
-    })
+DEF_TRAVERSE_DECL(TranslationUnitDecl, {
+  // Code in an unnamed namespace shows up automatically in
+
+  // If the traversal scope is set, then consider them to be the children of
+  // the TUDecl, rather than traversing (and loading?) all top-level decls.
+  auto Scope = D->getASTContext().getTraversalScope();
+  bool HasLimitedScope =
+      Scope.size() != 1 || !isa<TranslationUnitDecl>(Scope.front());
+  if (HasLimitedScope) {
+    ShouldVisitChildren = false; // we'll do that here instead
+    for (auto *Child : Scope) {
+      if (!canIgnoreChildDeclWhileTraversingDeclContext(Child))
+        TRY_TO(TraverseDecl(Child));
+    }
+  }
 
 DEF_TRAVERSE_DECL(PragmaCommentDecl, {})
 
 DEF_TRAVERSE_DECL(PragmaDetectMismatchDecl, {})
 
 DEF_TRAVERSE_DECL(ExternCContextDecl, {})
 
 DEF_TRAVERSE_DECL(NamespaceAliasDecl, {

clang/unittests/AST/ASTContextParentMapTest.cpp

EXPECT_TRUE(TemplateVerifier.match(
       "template<typename T> struct C { void f() {} };"
       "void g() { C<int> c; c.f(); }",
       compoundStmt(allOf(
           hasAncestor(cxxRecordDecl(isTemplateInstantiation())),
           hasAncestor(cxxRecordDecl(unless(isTemplateInstantiation())))))));
 }
 
 TEST(GetParents, RespectsTraversalScope) {
-  auto AST =
-      tooling::buildASTFromCode("struct foo { int bar; };", "foo.cpp",
-                                std::make_shared<PCHContainerOperations>());
+  auto AST = tooling::buildASTFromCode(
+      "struct foo { int bar; }; struct baz{};", "foo.cpp",
   auto &Ctx = AST->getASTContext();
   auto &TU = *Ctx.getTranslationUnitDecl();
   auto &Foo = *TU.lookup(&Ctx.Idents.get("foo")).front();
   auto &Bar = *cast<DeclContext>(Foo).lookup(&Ctx.Idents.get("bar")).front();
+  auto &Baz = *TU.lookup(&Ctx.Idents.get("baz")).front();
 
   // Initially, scope is the whole TU.
   EXPECT_THAT(Ctx.getParents(Bar), ElementsAre(DynTypedNode::create(Foo)));
   EXPECT_THAT(Ctx.getParents(Foo), ElementsAre(DynTypedNode::create(TU)));
+  EXPECT_THAT(Ctx.getParents(Baz), ElementsAre(DynTypedNode::create(TU)));
 
   // Restrict the scope, now some parents are gone.
   Ctx.setTraversalScope({&Foo});
   EXPECT_THAT(Ctx.getParents(Bar), ElementsAre(DynTypedNode::create(Foo)));
-  EXPECT_THAT(Ctx.getParents(Foo), ElementsAre());
+  EXPECT_THAT(Ctx.getParents(Foo), ElementsAre(DynTypedNode::create(TU)));
+  EXPECT_THAT(Ctx.getParents(Baz), ElementsAre());
 
   // Reset the scope, we get back the original results.
   Ctx.setTraversalScope({&TU});
   EXPECT_THAT(Ctx.getParents(Bar), ElementsAre(DynTypedNode::create(Foo)));
   EXPECT_THAT(Ctx.getParents(Foo), ElementsAre(DynTypedNode::create(TU)));
+  EXPECT_THAT(Ctx.getParents(Baz), 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

clang/unittests/Tooling/RecursiveASTVisitorTests/TraversalScope.cpp

 using namespace clang;
 
 namespace {
 
 class Visitor : public ExpectedLocationVisitor<Visitor, clang::TestVisitor> {
 public:
   Visitor(ASTContext *Context) { this->Context = Context; }
 
+  bool VisitTranslationUnitDecl(TranslationUnitDecl *D) {
+    auto &SM = D->getParentASTContext().getSourceManager();
+    Match("TU", SM.getLocForStartOfFile(SM.getMainFileID()));
+    return true;
+  }
+
   bool VisitNamedDecl(NamedDecl *D) {
     if (!D->isImplicit())
       Match(D->getName(), D->getLocation());
     return true;
   }
 };
 
 TEST(RecursiveASTVisitor, RespectsTraversalScope) {
struct foo {
   auto &Ctx = AST->getASTContext();
   auto &TU = *Ctx.getTranslationUnitDecl();
   auto &Foo = *TU.lookup(&Ctx.Idents.get("foo")).front();
   auto &Bar = *cast<DeclContext>(Foo).lookup(&Ctx.Idents.get("bar")).front();
 
   Ctx.setTraversalScope({&Bar});
 
   Visitor V(&Ctx);
+  V.ExpectMatch("TU", 1, 1);
   V.DisallowMatch("foo", 2, 8);
   V.ExpectMatch("bar", 3, 10);
   V.ExpectMatch("baz", 4, 12);
   V.TraverseAST(Ctx);
 }
 
 } // end anonymous namespace