2 files changed, 28 insertions, 4 deletions
diff --git a/clang-tools-extra/clangd/ClangdUnit.cpp b/clang-tools-extra/clangd/ClangdUnit.cpp
index 3beec0aa228..b337d851d0f 100644
--- a/clang-tools-extra/clangd/ClangdUnit.cpp
+++ b/clang-tools-extra/clangd/ClangdUnit.cpp
@@ -421,12 +421,16 @@ ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
   Clang->getPreprocessor().addCommentHandler(IWYUHandler.get());
 
   // Collect tokens of the main file.
-  syntax::TokenCollector Tokens(Clang->getPreprocessor());
+  syntax::TokenCollector CollectTokens(Clang->getPreprocessor());
 
   if (llvm::Error Err = Action->Execute())
     log("Execute() failed when building AST for {0}: {1}", MainInput.getFile(),
         toString(std::move(Err)));
 
+  // We have to consume the tokens before running clang-tidy to avoid collecting
+  // tokens from running the preprocessor inside the checks (only
+  // modernize-use-trailing-return-type does that today).
+  syntax::TokenBuffer Tokens = std::move(CollectTokens).consume();
   std::vector<Decl *> ParsedDecls = Action->takeTopLevelDecls();
   // AST traversals should exclude the preamble, to avoid performance cliffs.
   Clang->getASTContext().setTraversalScope(ParsedDecls);
@@ -452,9 +456,8 @@ 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(Tokens).consume(), std::move(ParsedDecls),
-                   std::move(Diags), std::move(Includes),
-                   std::move(CanonIncludes));
+                   std::move(Tokens), std::move(ParsedDecls), std::move(Diags),
+                   std::move(Includes), std::move(CanonIncludes));
 }
 
 ParsedAST::ParsedAST(ParsedAST &&Other) = default;
diff --git a/clang-tools-extra/clangd/unittests/ClangdUnitTests.cpp b/clang-tools-extra/clangd/unittests/ClangdUnitTests.cpp
index 3b5920e947d..47c353ee106 100644
--- a/clang-tools-extra/clangd/unittests/ClangdUnitTests.cpp
+++ b/clang-tools-extra/clangd/unittests/ClangdUnitTests.cpp
@@ -134,6 +134,27 @@ TEST(ClangdUnitTest, TokensAfterPreamble) {
   EXPECT_EQ(Spelled.back().text(SM), "last_token");
 }
 
+
+TEST(ClangdUnitTest, NoCrashOnTokensWithTidyCheck) {
+  TestTU TU;
+  // this check runs the preprocessor, we need to make sure it does not break
+  // our recording logic.
+  TU.ClangTidyChecks = "modernize-use-trailing-return-type";
+  TU.Code = "inline int foo() {}";
+
+  auto AST = TU.build();
+  const syntax::TokenBuffer &T = AST.getTokens();
+  const auto &SM = AST.getSourceManager();
+
+  ASSERT_GT(T.expandedTokens().size(), 7u);
+  // Check first token after the preamble.
+  EXPECT_EQ(T.expandedTokens().front().text(SM), "inline");
+  // 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), "}");
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang