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

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.

Reviewers: sammccall

Reviewed By: sammccall

Subscribers: mgorny, MaskRay, jkorous, arphaman, kadircet, cfe-commits

Tags: #clang

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

llvm-svn: 363803

6 files changed, 57 insertions, 4 deletions

diff --git a/clang-tools-extra/clangd/CMakeLists.txt b/clang-tools-extra/clangd/CMakeLists.txt
index f110d0ca5ca..12f20b54724 100644
--- a/clang-tools-extra/clangd/CMakeLists.txt
+++ b/clang-tools-extra/clangd/CMakeLists.txt
@@ -128,6 +128,7 @@ add_clang_library(clangDaemon
   clangToolingCore
   clangToolingInclusions
   clangToolingRefactoring
+  clangToolingSyntax
   ${LLVM_PTHREAD_LIB}
   ${CLANGD_ATOMIC_LIB}
   )
diff --git a/clang-tools-extra/clangd/ClangdUnit.cpp b/clang-tools-extra/clangd/ClangdUnit.cpp
index 31476d99152..c82cc699e6c 100644
--- a/clang-tools-extra/clangd/ClangdUnit.cpp
+++ b/clang-tools-extra/clangd/ClangdUnit.cpp
@@ -36,6 +36,7 @@
 #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"
@@ -418,6 +419,9 @@ ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
       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());
 
@@ -446,8 +450,9 @@ ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
   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;
@@ -540,11 +545,13 @@ PreambleData::PreambleData(PrecompiledPreamble Preamble,
 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);
diff --git a/clang-tools-extra/clangd/ClangdUnit.h b/clang-tools-extra/clangd/ClangdUnit.h
index 16246eb2784..f5b18f97387 100644
--- a/clang-tools-extra/clangd/ClangdUnit.h
+++ b/clang-tools-extra/clangd/ClangdUnit.h
@@ -24,6 +24,7 @@
 #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>
@@ -115,10 +116,14 @@ public:
   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);
 
@@ -132,6 +137,11 @@ private:
   // 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;
diff --git a/clang-tools-extra/clangd/tool/CMakeLists.txt b/clang-tools-extra/clangd/tool/CMakeLists.txt
index 93d7f1459b1..0846b648db8 100644
--- a/clang-tools-extra/clangd/tool/CMakeLists.txt
+++ b/clang-tools-extra/clangd/tool/CMakeLists.txt
@@ -26,5 +26,6 @@ target_link_libraries(clangd
   clangSema
   clangTooling
   clangToolingCore
+  clangToolingSyntax
   ${CLANGD_XPC_LIBS}
   )
diff --git a/clang-tools-extra/clangd/unittests/CMakeLists.txt b/clang-tools-extra/clangd/unittests/CMakeLists.txt
index 11f028621ad..2b19b02a969 100644
--- a/clang-tools-extra/clangd/unittests/CMakeLists.txt
+++ b/clang-tools-extra/clangd/unittests/CMakeLists.txt
@@ -87,6 +87,7 @@ target_link_libraries(ClangdTests
   clangTooling
   clangToolingCore
   clangToolingInclusions
+  clangToolingSyntax
   LLVMSupport
   LLVMTestingSupport
   )
diff --git a/clang-tools-extra/clangd/unittests/ClangdUnitTests.cpp b/clang-tools-extra/clangd/unittests/ClangdUnitTests.cpp
index 100e92c3c65..bfeb1c83e8b 100644
--- a/clang-tools-extra/clangd/unittests/ClangdUnitTests.cpp
+++ b/clang-tools-extra/clangd/unittests/ClangdUnitTests.cpp
@@ -10,6 +10,8 @@
 #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"
@@ -81,6 +83,37 @@ TEST(ClangdUnitTest, TopLevelDecls) {
   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