diff --git a/clang-tools-extra/clangd/SemanticSelection.cpp b/clang-tools-extra/clangd/SemanticSelection.cpp
--- a/clang-tools-extra/clangd/SemanticSelection.cpp
+++ b/clang-tools-extra/clangd/SemanticSelection.cpp
@@ -160,8 +160,7 @@
 // Related issue: https://github.com/clangd/clangd/issues/310
 llvm::Expected<std::vector<FoldingRange>> getFoldingRanges(ParsedAST &AST) {
   syntax::Arena A(AST.getSourceManager(), AST.getLangOpts(), AST.getTokens());
-  const auto *SyntaxTree =
-      syntax::buildSyntaxTree(A, *AST.getASTContext().getTranslationUnitDecl());
+  const auto *SyntaxTree = syntax::buildSyntaxTree(A, AST.getASTContext());
   return collectFoldingRanges(SyntaxTree, AST.getSourceManager());
 }
 
diff --git a/clang/include/clang/Tooling/Syntax/BuildTree.h b/clang/include/clang/Tooling/Syntax/BuildTree.h
--- a/clang/include/clang/Tooling/Syntax/BuildTree.h
+++ b/clang/include/clang/Tooling/Syntax/BuildTree.h
@@ -19,8 +19,9 @@
 namespace syntax {
 
 /// Build a syntax tree for the main file.
-syntax::TranslationUnit *buildSyntaxTree(Arena &A,
-                                         const clang::TranslationUnitDecl &TU);
+/// This usually covers the whole TranslationUnitDecl, but can be restricted by
+/// the ASTContext's traversal scope.
+syntax::TranslationUnit *buildSyntaxTree(Arena &A, ASTContext &Context);
 
 // Create syntax trees from subtrees not backed by the source code.
 
diff --git a/clang/lib/Tooling/Syntax/BuildTree.cpp b/clang/lib/Tooling/Syntax/BuildTree.cpp
--- a/clang/lib/Tooling/Syntax/BuildTree.cpp
+++ b/clang/lib/Tooling/Syntax/BuildTree.cpp
@@ -1711,9 +1711,9 @@
   return It->second;
 }
 
-syntax::TranslationUnit *
-syntax::buildSyntaxTree(Arena &A, const TranslationUnitDecl &TU) {
+syntax::TranslationUnit *syntax::buildSyntaxTree(Arena &A,
+                                                 ASTContext &Context) {
   TreeBuilder Builder(A);
-  BuildTreeVisitor(TU.getASTContext(), Builder).TraverseAST(TU.getASTContext());
+  BuildTreeVisitor(Context, Builder).TraverseAST(Context);
   return std::move(Builder).finalize();
 }
diff --git a/clang/unittests/Tooling/Syntax/TreeTestBase.cpp b/clang/unittests/Tooling/Syntax/TreeTestBase.cpp
--- a/clang/unittests/Tooling/Syntax/TreeTestBase.cpp
+++ b/clang/unittests/Tooling/Syntax/TreeTestBase.cpp
@@ -81,7 +81,7 @@
       Tokens = nullptr; // make sure we fail if this gets called twice.
       Arena = std::make_unique<syntax::Arena>(Ctx.getSourceManager(),
                                               Ctx.getLangOpts(), *TB);
-      Root = syntax::buildSyntaxTree(*Arena, *Ctx.getTranslationUnitDecl());
+      Root = syntax::buildSyntaxTree(*Arena, Ctx);
     }
 
   private: