[clangd] Collect tokens of main files when building the AST

[ilya-biryukov]

Summary

The first use of this is a code tweak to expand macro calls.
Will later be used to build syntax trees.

The memory overhead is small as we only store tokens of the main file.

[sammccall]

Can we add a test using TestTU that does a very basic verification of expanded/spelled tokens (first after preamble, last token in file)?

clang-tools-extra/clangd/ClangdUnit.h

140 ↗

(On Diff #203357)

Not sure "expanded" is accurate here - it has both.
(And it'd be worth spelling out if expanded tokens or spelled tokens or both are missing for the preamble)

[ilya-biryukov]

• Rename tokens() to getTokens() for consistency with the rest of ParsedAST.
• Update comments.
• Add a test.

[ilya-biryukov]

In D62956#1547826, @sammccall wrote:

Can we add a test using TestTU that does a very basic verification of expanded/spelled tokens (first after preamble, last token in file)?

Done. Will land this tomorrow, want to run it over all of LLVM first to give the token-building code some coverage.

clang-tools-extra/trunk/clangd/CMakeLists.txt

add_clang_library(clangDaemon
   clangTidyPerformanceModule
   clangTidyPortabilityModule
   clangTidyReadabilityModule
   clangTidyZirconModule
   clangTooling
   clangToolingCore
   clangToolingInclusions
   clangToolingRefactoring
+  clangToolingSyntax
   ${LLVM_PTHREAD_LIB}
   ${CLANGD_ATOMIC_LIB}
   )
 
 add_subdirectory(refactor/tweaks)
 if (${CMAKE_SYSTEM_NAME} STREQUAL "Linux")
   # FIXME: Make fuzzer not use linux-specific APIs, build it everywhere.
   add_subdirectory(fuzzer)

clang-tools-extra/trunk/clangd/ClangdUnit.h

 #include "index/CanonicalIncludes.h"
 #include "index/Index.h"
 #include "clang/Frontend/FrontendAction.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Serialization/ASTBitCodes.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
+#include "clang/Tooling/Syntax/Tokens.h"
 #include <memory>
 #include <string>
 #include <vector>
 
 namespace llvm {
 class raw_ostream;
 
 namespace vfs {
public:
   const std::vector<Diag> &getDiagnostics() const;
 
   /// Returns the esitmated size of the AST and the accessory structures, in
   /// bytes. Does not include the size of the preamble.
   std::size_t getUsedBytes() const;
   const IncludeStructure &getIncludeStructure() const;
   const CanonicalIncludes &getCanonicalIncludes() const;
 
+  /// Tokens recorded while parsing the main file.
+  /// (!) does not have tokens from the preamble.
+  const syntax::TokenBuffer &getTokens() const { return Tokens; }
+
 private:
   ParsedAST(std::shared_ptr<const PreambleData> Preamble,
             std::unique_ptr<CompilerInstance> Clang,
-            std::unique_ptr<FrontendAction> Action,
+            std::unique_ptr<FrontendAction> Action, syntax::TokenBuffer Tokens,
             std::vector<Decl *> LocalTopLevelDecls, std::vector<Diag> Diags,
             IncludeStructure Includes, CanonicalIncludes CanonIncludes);
 
   // In-memory preambles must outlive the AST, it is important that this member
   // goes before Clang and Action.
   std::shared_ptr<const PreambleData> Preamble;
   // We store an "incomplete" FrontendAction (i.e. no EndSourceFile was called
   // on it) and CompilerInstance used to run it. That way we don't have to do
   // complex memory management of all Clang structures on our own. (They are
   // stored in CompilerInstance and cleaned up by
   // FrontendAction.EndSourceFile).
   std::unique_ptr<CompilerInstance> Clang;
   std::unique_ptr<FrontendAction> Action;
+  /// Tokens recorded after the preamble finished.
+  ///   - Includes all spelled tokens for the main file.
+  ///   - Includes expanded tokens produced **after** preabmle.
+  ///   - Does not have spelled or expanded tokens for files from preamble.
+  syntax::TokenBuffer Tokens;
 
   // Data, stored after parsing.
   std::vector<Diag> Diags;
   // Top-level decls inside the current file. Not that this does not include
   // top-level decls from the preamble.
   std::vector<Decl *> LocalTopLevelDecls;
   IncludeStructure Includes;
   CanonicalIncludes CanonIncludes;

clang-tools-extra/trunk/clangd/ClangdUnit.cpp

 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Sema/Sema.h"
 #include "clang/Serialization/ASTWriter.h"
 #include "clang/Serialization/PCHContainerOperations.h"
 #include "clang/Tooling/CompilationDatabase.h"
+#include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <memory>
 
ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
   if (Preamble)
     CanonIncludes = Preamble->CanonIncludes;
   else
     addSystemHeadersMapping(&CanonIncludes);
   std::unique_ptr<CommentHandler> IWYUHandler =
       collectIWYUHeaderMaps(&CanonIncludes);
   Clang->getPreprocessor().addCommentHandler(IWYUHandler.get());
 
+  // Collect tokens of the main file.
+  syntax::TokenCollector Tokens(Clang->getPreprocessor());
+
   if (!Action->Execute())
     log("Execute() failed when building AST for {0}", MainInput.getFile());
 
   std::vector<Decl *> ParsedDecls = Action->takeTopLevelDecls();
   // AST traversals should exclude the preamble, to avoid performance cliffs.
   Clang->getASTContext().setTraversalScope(ParsedDecls);
   {
     // Run the AST-dependent part of the clang-tidy checks.
ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
   // So just inform the preprocessor of EOF, while keeping everything alive.
   Clang->getPreprocessor().EndSourceFile();
 
   std::vector<Diag> Diags = ASTDiags.take(CTContext.getPointer());
   // Add diagnostics from the preamble, if any.
   if (Preamble)
     Diags.insert(Diags.begin(), Preamble->Diags.begin(), Preamble->Diags.end());
   return ParsedAST(std::move(Preamble), std::move(Clang), std::move(Action),
-                   std::move(ParsedDecls), std::move(Diags),
-                   std::move(Includes), std::move(CanonIncludes));
+                   std::move(Tokens).consume(), std::move(ParsedDecls),
+                   std::move(Diags), std::move(Includes),
+                   std::move(CanonIncludes));
 }
 
 ParsedAST::ParsedAST(ParsedAST &&Other) = default;
 
 ParsedAST &ParsedAST::operator=(ParsedAST &&Other) = default;
 
 ParsedAST::~ParsedAST() {
   if (Action) {
PreambleData::PreambleData(PrecompiledPreamble Preamble,
     : Preamble(std::move(Preamble)), Diags(std::move(Diags)),
       Includes(std::move(Includes)), MainFileMacros(std::move(MainFileMacros)),
       StatCache(std::move(StatCache)), CanonIncludes(std::move(CanonIncludes)) {
 }
 
 ParsedAST::ParsedAST(std::shared_ptr<const PreambleData> Preamble,
                      std::unique_ptr<CompilerInstance> Clang,
                      std::unique_ptr<FrontendAction> Action,
+                     syntax::TokenBuffer Tokens,
                      std::vector<Decl *> LocalTopLevelDecls,
                      std::vector<Diag> Diags, IncludeStructure Includes,
                      CanonicalIncludes CanonIncludes)
     : Preamble(std::move(Preamble)), Clang(std::move(Clang)),
-      Action(std::move(Action)), Diags(std::move(Diags)),
+      Action(std::move(Action)), Tokens(std::move(Tokens)),
+      Diags(std::move(Diags)),
       LocalTopLevelDecls(std::move(LocalTopLevelDecls)),
       Includes(std::move(Includes)), CanonIncludes(std::move(CanonIncludes)) {
   assert(this->Clang);
   assert(this->Action);
 }
 
 std::shared_ptr<const PreambleData>
 buildPreamble(PathRef FileName, CompilerInvocation &CI,

clang-tools-extra/trunk/clangd/tool/CMakeLists.txt

target_link_libraries(clangd
   clangBasic
   clangTidy
   clangDaemon
   clangFormat
   clangFrontend
   clangSema
   clangTooling
   clangToolingCore
+  clangToolingSyntax
   ${CLANGD_XPC_LIBS}
   )

clang-tools-extra/trunk/clangd/unittests/CMakeLists.txt

target_link_libraries(ClangdTests
   clangIndex
   clangLex
   clangSema
   clangSerialization
   clangTidy
   clangTooling
   clangToolingCore
   clangToolingInclusions
+  clangToolingSyntax
   LLVMSupport
   LLVMTestingSupport
   )
 
 if (CLANGD_BUILD_XPC)
   add_subdirectory(xpc)
 endif ()
 
 configure_lit_site_cfg(
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in
   ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py)

clang-tools-extra/trunk/clangd/unittests/ClangdUnitTests.cpp

 //===-- ClangdUnitTests.cpp - ClangdUnit tests ------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Annotations.h"
 #include "ClangdUnit.h"
 #include "SourceCode.h"
 #include "TestTU.h"
+#include "clang/Basic/TokenKinds.h"
+#include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
TU.HeaderCode = R"(
     int header1();
     int header2;
   )";
   TU.Code = "int main();";
   auto AST = TU.build();
   EXPECT_THAT(AST.getLocalTopLevelDecls(), ElementsAre(DeclNamed("main")));
 }
 
+TEST(ClangdUnitTest, TokensAfterPreamble) {
+  TestTU TU;
+  TU.AdditionalFiles["foo.h"] = R"(
+    int foo();
+  )";
+  TU.Code = R"cpp(
+      #include "foo.h"
+      first_token;
+      void test() {
+      }
+      last_token
+)cpp";
+  auto AST = TU.build();
+  const syntax::TokenBuffer &T = AST.getTokens();
+  const auto &SM = AST.getSourceManager();
+
+  ASSERT_GT(T.expandedTokens().size(), 2u);
+  // Check first token after the preamble.
+  EXPECT_EQ(T.expandedTokens().front().text(SM), "first_token");
+  // Last token is always 'eof'.
+  EXPECT_EQ(T.expandedTokens().back().kind(), tok::eof);
+  // Check the token before 'eof'.
+  EXPECT_EQ(T.expandedTokens().drop_back().back().text(SM), "last_token");
+
+  // The spelled tokens for the main file should have everything.
+  auto Spelled = T.spelledTokens(SM.getMainFileID());
+  ASSERT_FALSE(Spelled.empty());
+  EXPECT_EQ(Spelled.front().kind(), tok::hash);
+  EXPECT_EQ(Spelled.back().text(SM), "last_token");
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang