[clangd] Build index on preamble changes instead of the AST changes

Summary:
This is more efficient and avoids data races when reading files that
come from the preamble. The staleness can occur when reading a file
from disk that changed after the preamble was built. This can lead to
crashes, e.g. when parsing comments.

We do not to rely on symbols from the main file anyway, since any info
that those provide can always be taken from the AST.

Reviewers: ioeric, sammccall

Reviewed By: ioeric

Subscribers: malaperle, klimek, javed.absar, MaskRay, jkorous, cfe-commits

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

10 files changed, 119 insertions, 40 deletions

diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index 6d22bcdbeea..0512bca61de 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -17,6 +17,7 @@
 #include "clang/Format/Format.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
+#include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Refactoring/RefactoringResultConsumer.h"
 #include "clang/Tooling/Refactoring/Rename/RenamingAction.h"
@@ -92,12 +93,14 @@ ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
       // is parsed.
       // FIXME(ioeric): this can be slow and we may be able to index on less
       // critical paths.
-      WorkScheduler(Opts.AsyncThreadsCount, Opts.StorePreamblesInMemory,
-                    FileIdx
-                        ? [this](PathRef Path,
-                                 ParsedAST *AST) { FileIdx->update(Path, AST); }
-                        : ASTParsedCallback(),
-                    Opts.UpdateDebounce) {
+      WorkScheduler(
+          Opts.AsyncThreadsCount, Opts.StorePreamblesInMemory,
+          FileIdx
+              ? [this](PathRef Path, ASTContext &AST,
+                       std::shared_ptr<Preprocessor>
+                           PP) { FileIdx->update(Path, &AST, std::move(PP)); }
+              : PreambleParsedCallback(),
+          Opts.UpdateDebounce) {
   if (FileIdx && Opts.StaticIndex) {
     MergedIndex = mergeIndex(FileIdx.get(), Opts.StaticIndex);
     Index = MergedIndex.get();
diff --git a/clang-tools-extra/clangd/ClangdUnit.cpp b/clang-tools-extra/clangd/ClangdUnit.cpp
index eb2fc184ccf..41f2df3669e 100644
--- a/clang-tools-extra/clangd/ClangdUnit.cpp
+++ b/clang-tools-extra/clangd/ClangdUnit.cpp
@@ -13,6 +13,8 @@
 #include "Logger.h"
 #include "SourceCode.h"
 #include "Trace.h"
+#include "clang/AST/ASTContext.h"
+#include "clang/Basic/LangOptions.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
@@ -25,7 +27,6 @@
 #include "clang/Sema/Sema.h"
 #include "clang/Serialization/ASTWriter.h"
 #include "clang/Tooling/CompilationDatabase.h"
-#include "clang/Basic/LangOptions.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/CrashRecoveryContext.h"
@@ -86,6 +87,9 @@ private:
 
 class CppFilePreambleCallbacks : public PreambleCallbacks {
 public:
+  CppFilePreambleCallbacks(PathRef File, PreambleParsedCallback ParsedCallback)
+      : File(File), ParsedCallback(ParsedCallback) {}
+
   std::vector<serialization::DeclID> takeTopLevelDeclIDs() {
     return std::move(TopLevelDeclIDs);
   }
@@ -102,6 +106,13 @@ public:
     }
   }
 
+  void AfterExecute(CompilerInstance &CI) override {
+    if (!ParsedCallback)
+      return;
+    trace::Span Tracer("Running PreambleCallback");
+    ParsedCallback(File, CI.getASTContext(), CI.getPreprocessorPtr());
+  }
+
   void HandleTopLevelDecl(DeclGroupRef DG) override {
     for (Decl *D : DG) {
       if (isa<ObjCMethodDecl>(D))
@@ -122,6 +133,8 @@ public:
   }
 
 private:
+  PathRef File;
+  PreambleParsedCallback ParsedCallback;
   std::vector<Decl *> TopLevelDecls;
   std::vector<serialization::DeclID> TopLevelDeclIDs;
   std::vector<Inclusion> Inclusions;
@@ -277,9 +290,9 @@ ParsedAST::ParsedAST(std::shared_ptr<const PreambleData> Preamble,
 
 CppFile::CppFile(PathRef FileName, bool StorePreamblesInMemory,
                  std::shared_ptr<PCHContainerOperations> PCHs,
-                 ASTParsedCallback ASTCallback)
+                 PreambleParsedCallback PreambleCallback)
     : FileName(FileName), StorePreamblesInMemory(StorePreamblesInMemory),
-      PCHs(std::move(PCHs)), ASTCallback(std::move(ASTCallback)) {
+      PCHs(std::move(PCHs)), PreambleCallback(std::move(PreambleCallback)) {
   log("Created CppFile for " + FileName);
 }
 
@@ -346,10 +359,6 @@ llvm::Optional<std::vector<Diag>> CppFile::rebuild(ParseInputs &&Inputs) {
     Diagnostics.insert(Diagnostics.end(), NewAST->getDiagnostics().begin(),
                        NewAST->getDiagnostics().end());
   }
-  if (ASTCallback && NewAST) {
-    trace::Span Tracer("Running ASTCallback");
-    ASTCallback(FileName, NewAST.getPointer());
-  }
 
   // Write the results of rebuild into class fields.
   Command = std::move(Inputs.CompileCommand);
@@ -404,7 +413,7 @@ CppFile::rebuildPreamble(CompilerInvocation &CI,
   assert(!CI.getFrontendOpts().SkipFunctionBodies);
   CI.getFrontendOpts().SkipFunctionBodies = true;
 
-  CppFilePreambleCallbacks SerializedDeclsCollector;
+  CppFilePreambleCallbacks SerializedDeclsCollector(FileName, PreambleCallback);
   auto BuiltPreamble = PrecompiledPreamble::Build(
       CI, &ContentsBuffer, Bounds, *PreambleDiagsEngine, FS, PCHs,
       /*StoreInMemory=*/StorePreamblesInMemory, SerializedDeclsCollector);
diff --git a/clang-tools-extra/clangd/ClangdUnit.h b/clang-tools-extra/clangd/ClangdUnit.h
index 64a19c0e61a..3b3d891f1d1 100644
--- a/clang-tools-extra/clangd/ClangdUnit.h
+++ b/clang-tools-extra/clangd/ClangdUnit.h
@@ -17,6 +17,7 @@
 #include "Protocol.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"
@@ -126,7 +127,8 @@ private:
   std::vector<Inclusion> Inclusions;
 };
 
-using ASTParsedCallback = std::function<void(PathRef Path, ParsedAST *)>;
+using PreambleParsedCallback = std::function<void(
+    PathRef Path, ASTContext &, std::shared_ptr<clang::Preprocessor>)>;
 
 /// Manages resources, required by clangd. Allows to rebuild file with new
 /// contents, and provides AST and Preamble for it.
@@ -134,7 +136,7 @@ class CppFile {
 public:
   CppFile(PathRef FileName, bool StorePreamblesInMemory,
           std::shared_ptr<PCHContainerOperations> PCHs,
-          ASTParsedCallback ASTCallback);
+          PreambleParsedCallback PreambleCallback);
 
   /// Rebuild the AST and the preamble.
   /// Returns a list of diagnostics or llvm::None, if an error occured.
@@ -170,7 +172,7 @@ private:
   std::shared_ptr<PCHContainerOperations> PCHs;
   /// This is called after the file is parsed. This can be nullptr if there is
   /// no callback.
-  ASTParsedCallback ASTCallback;
+  PreambleParsedCallback PreambleCallback;
 };
 
 /// Get the beginning SourceLocation at a specified \p Pos.
diff --git a/clang-tools-extra/clangd/TUScheduler.cpp b/clang-tools-extra/clangd/TUScheduler.cpp
index 51b13e2b7ba..588d257bfa0 100644
--- a/clang-tools-extra/clangd/TUScheduler.cpp
+++ b/clang-tools-extra/clangd/TUScheduler.cpp
@@ -414,11 +414,11 @@ struct TUScheduler::FileData {
 
 TUScheduler::TUScheduler(unsigned AsyncThreadsCount,
                          bool StorePreamblesInMemory,
-                         ASTParsedCallback ASTCallback,
+                         PreambleParsedCallback PreambleCallback,
                          steady_clock::duration UpdateDebounce)
     : StorePreamblesInMemory(StorePreamblesInMemory),
       PCHOps(std::make_shared<PCHContainerOperations>()),
-      ASTCallback(std::move(ASTCallback)), Barrier(AsyncThreadsCount),
+      PreambleCallback(std::move(PreambleCallback)), Barrier(AsyncThreadsCount),
       UpdateDebounce(UpdateDebounce) {
   if (0 < AsyncThreadsCount) {
     PreambleTasks.emplace();
@@ -455,7 +455,7 @@ void TUScheduler::update(PathRef File, ParseInputs Inputs,
     // Create a new worker to process the AST-related tasks.
     ASTWorkerHandle Worker = ASTWorker::Create(
         File, WorkerThreads ? WorkerThreads.getPointer() : nullptr, Barrier,
-        CppFile(File, StorePreamblesInMemory, PCHOps, ASTCallback),
+        CppFile(File, StorePreamblesInMemory, PCHOps, PreambleCallback),
         UpdateDebounce);
     FD = std::unique_ptr<FileData>(new FileData{
         Inputs.Contents, Inputs.CompileCommand, std::move(Worker)});
diff --git a/clang-tools-extra/clangd/TUScheduler.h b/clang-tools-extra/clangd/TUScheduler.h
index aea669d9c73..99748427443 100644
--- a/clang-tools-extra/clangd/TUScheduler.h
+++ b/clang-tools-extra/clangd/TUScheduler.h
@@ -52,7 +52,7 @@ enum class WantDiagnostics {
 class TUScheduler {
 public:
   TUScheduler(unsigned AsyncThreadsCount, bool StorePreamblesInMemory,
-              ASTParsedCallback ASTCallback,
+              PreambleParsedCallback PreambleCallback,
               std::chrono::steady_clock::duration UpdateDebounce);
   ~TUScheduler();
 
@@ -101,7 +101,7 @@ private:
 
   const bool StorePreamblesInMemory;
   const std::shared_ptr<PCHContainerOperations> PCHOps;
-  const ASTParsedCallback ASTCallback;
+  const PreambleParsedCallback PreambleCallback;
   Semaphore Barrier;
   llvm::StringMap<std::unique_ptr<FileData>> Files;
   // None when running tasks synchronously and non-None when running tasks
diff --git a/clang-tools-extra/clangd/index/FileIndex.cpp b/clang-tools-extra/clangd/index/FileIndex.cpp
index a6e80f83db8..7e19b00458f 100644
--- a/clang-tools-extra/clangd/index/FileIndex.cpp
+++ b/clang-tools-extra/clangd/index/FileIndex.cpp
@@ -10,12 +10,12 @@
 #include "FileIndex.h"
 #include "SymbolCollector.h"
 #include "clang/Index/IndexingAction.h"
+#include "clang/Lex/Preprocessor.h"
 
 namespace clang {
 namespace clangd {
 
-SymbolSlab indexAST(ParsedAST *AST) {
-  assert(AST && "AST must not be nullptr!");
+SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP) {
   SymbolCollector::Options CollectorOpts;
   // FIXME(ioeric): we might also want to collect include headers. We would need
   // to make sure all includes are canonicalized (with CanonicalIncludes), which
@@ -26,15 +26,18 @@ SymbolSlab indexAST(ParsedAST *AST) {
   CollectorOpts.CountReferences = false;
 
   SymbolCollector Collector(std::move(CollectorOpts));
-  Collector.setPreprocessor(AST->getPreprocessorPtr());
+  Collector.setPreprocessor(PP);
   index::IndexingOptions IndexOpts;
   // We only need declarations, because we don't count references.
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::DeclarationsOnly;
   IndexOpts.IndexFunctionLocals = false;
 
-  index::indexTopLevelDecls(AST->getASTContext(), AST->getTopLevelDecls(),
-                            Collector, IndexOpts);
+  std::vector<const Decl *> TopLevelDecls(
+      AST.getTranslationUnitDecl()->decls().begin(),
+      AST.getTranslationUnitDecl()->decls().end());
+  index::indexTopLevelDecls(AST, TopLevelDecls, Collector, IndexOpts);
+
   return Collector.takeSymbols();
 }
 
@@ -69,12 +72,14 @@ std::shared_ptr<std::vector<const Symbol *>> FileSymbols::allSymbols() {
   return {std::move(Snap), Pointers};
 }
 
-void FileIndex::update(PathRef Path, ParsedAST *AST) {
+void FileIndex::update(PathRef Path, ASTContext *AST,
+                       std::shared_ptr<Preprocessor> PP) {
   if (!AST) {
     FSymbols.update(Path, nullptr);
   } else {
+    assert(PP);
     auto Slab = llvm::make_unique<SymbolSlab>();
-    *Slab = indexAST(AST);
+    *Slab = indexAST(*AST, PP);
     FSymbols.update(Path, std::move(Slab));
   }
   auto Symbols = FSymbols.allSymbols();
diff --git a/clang-tools-extra/clangd/index/FileIndex.h b/clang-tools-extra/clangd/index/FileIndex.h
index a0afccea5be..f25ed1749bc 100644
--- a/clang-tools-extra/clangd/index/FileIndex.h
+++ b/clang-tools-extra/clangd/index/FileIndex.h
@@ -19,6 +19,7 @@
 #include "../ClangdUnit.h"
 #include "Index.h"
 #include "MemIndex.h"
+#include "clang/Lex/Preprocessor.h"
 
 namespace clang {
 namespace clangd {
@@ -56,8 +57,10 @@ private:
 class FileIndex : public SymbolIndex {
 public:
   /// \brief Update symbols in \p Path with symbols in \p AST. If \p AST is
-  /// nullptr, this removes all symbols in the file
-  void update(PathRef Path, ParsedAST *AST);
+  /// nullptr, this removes all symbols in the file.
+  /// If \p AST is not null, \p PP cannot be null and it should be the
+  /// preprocessor that was used to build \p AST.
+  void update(PathRef Path, ASTContext *AST, std::shared_ptr<Preprocessor> PP);
 
   bool
   fuzzyFind(const FuzzyFindRequest &Req,
@@ -73,7 +76,7 @@ private:
 
 /// Retrieves namespace and class level symbols in \p AST.
 /// Exposed to assist in unit tests.
-SymbolSlab indexAST(ParsedAST *AST);
+SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP);
 
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/unittests/clangd/FileIndexTests.cpp b/clang-tools-extra/unittests/clangd/FileIndexTests.cpp
index b1107b98d3f..1068a4b2f4d 100644
--- a/clang-tools-extra/unittests/clangd/FileIndexTests.cpp
+++ b/clang-tools-extra/unittests/clangd/FileIndexTests.cpp
@@ -7,8 +7,13 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "ClangdUnit.h"
+#include "TestFS.h"
 #include "TestTU.h"
 #include "index/FileIndex.h"
+#include "clang/Frontend/PCHContainerOperations.h"
+#include "clang/Lex/Preprocessor.h"
+#include "clang/Tooling/CompilationDatabase.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
@@ -87,7 +92,7 @@ void update(FileIndex &M, llvm::StringRef Basename, llvm::StringRef Code) {
   File.HeaderFilename = (Basename + ".h").str();
   File.HeaderCode = Code;
   auto AST = File.build();
-  M.update(File.Filename, &AST);
+  M.update(File.Filename, &AST.getASTContext(), AST.getPreprocessorPtr());
 }
 
 TEST(FileIndexTest, IndexAST) {
@@ -129,13 +134,13 @@ TEST(FileIndexTest, RemoveAST) {
   Req.Scopes = {"ns::"};
   EXPECT_THAT(match(M, Req), UnorderedElementsAre("ns::f", "ns::X"));
 
-  M.update("f1.cpp", nullptr);
+  M.update("f1.cpp", nullptr, nullptr);
   EXPECT_THAT(match(M, Req), UnorderedElementsAre());
 }
 
 TEST(FileIndexTest, RemoveNonExisting) {
   FileIndex M;
-  M.update("no.cpp", nullptr);
+  M.update("no.cpp", nullptr, nullptr);
   EXPECT_THAT(match(M, FuzzyFindRequest()), UnorderedElementsAre());
 }
 
@@ -200,6 +205,58 @@ vector<Ty> make_vector(Arg A) {}
   EXPECT_TRUE(SeenMakeVector);
 }
 
+TEST(FileIndexTest, RebuildWithPreamble) {
+  auto FooCpp = testPath("foo.cpp");
+  auto FooH = testPath("foo.h");
+  FileIndex Index;
+  bool IndexUpdated = false;
+  CppFile File("foo.cpp", /*StorePreambleInMemory=*/true,
+               std::make_shared<PCHContainerOperations>(),
+               [&Index, &IndexUpdated](PathRef FilePath, ASTContext &Ctx,
+                                       std::shared_ptr<Preprocessor> PP) {
+                 EXPECT_FALSE(IndexUpdated)
+                     << "Expected only a single index update";
+                 IndexUpdated = true;
+                 Index.update(FilePath, &Ctx, std::move(PP));
+               });
+
+  // Preparse ParseInputs.
+  ParseInputs PI;
+  PI.CompileCommand.Directory = testRoot();
+  PI.CompileCommand.Filename = FooCpp;
+  PI.CompileCommand.CommandLine = {"clang", "-xc++", FooCpp};
+
+  llvm::StringMap<std::string> Files;
+  Files[FooCpp] = "";
+  Files[FooH] = R"cpp(
+    namespace ns_in_header {
+      int func_in_header();
+    }
+  )cpp";
+  PI.FS = buildTestFS(std::move(Files));
+
+  PI.Contents = R"cpp(
+    #include "foo.h"
+    namespace ns_in_source {
+      int func_in_source();
+    }
+  )cpp";
+
+  // Rebuild the file.
+  File.rebuild(std::move(PI));
+  ASSERT_TRUE(IndexUpdated);
+
+  // Check the index contains symbols from the preamble, but not from the main
+  // file.
+  FuzzyFindRequest Req;
+  Req.Query = "";
+  Req.Scopes = {"", "ns_in_header::"};
+
+  EXPECT_THAT(
+      match(Index, Req),
+      UnorderedElementsAre("ns_in_header", "ns_in_header::func_in_header"));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/unittests/clangd/TUSchedulerTests.cpp b/clang-tools-extra/unittests/clangd/TUSchedulerTests.cpp
index 2235e8722c2..879bc5828a3 100644
--- a/clang-tools-extra/unittests/clangd/TUSchedulerTests.cpp
+++ b/clang-tools-extra/unittests/clangd/TUSchedulerTests.cpp
@@ -42,7 +42,7 @@ private:
 TEST_F(TUSchedulerTests, MissingFiles) {
   TUScheduler S(getDefaultAsyncThreadsCount(),
                 /*StorePreamblesInMemory=*/true,
-                /*ASTParsedCallback=*/nullptr,
+                /*PreambleParsedCallback=*/nullptr,
                 /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero());
 
   auto Added = testPath("added.cpp");
@@ -98,7 +98,7 @@ TEST_F(TUSchedulerTests, WantDiagnostics) {
     TUScheduler S(
         getDefaultAsyncThreadsCount(),
         /*StorePreamblesInMemory=*/true,
-        /*ASTParsedCallback=*/nullptr,
+        /*PreambleParsedCallback=*/nullptr,
         /*UpdateDebounce=*/std::chrono::steady_clock::duration::zero());
     auto Path = testPath("foo.cpp");
     S.update(Path, getInputs(Path, ""), WantDiagnostics::Yes,
@@ -126,7 +126,7 @@ TEST_F(TUSchedulerTests, Debounce) {
   {
     TUScheduler S(getDefaultAsyncThreadsCount(),
                   /*StorePreamblesInMemory=*/true,
-                  /*ASTParsedCallback=*/nullptr,
+                  /*PreambleParsedCallback=*/nullptr,
                   /*UpdateDebounce=*/std::chrono::seconds(1));
     // FIXME: we could probably use timeouts lower than 1 second here.
     auto Path = testPath("foo.cpp");
@@ -157,7 +157,7 @@ TEST_F(TUSchedulerTests, ManyUpdates) {
   {
     TUScheduler S(getDefaultAsyncThreadsCount(),
                   /*StorePreamblesInMemory=*/true,
-                  /*ASTParsedCallback=*/nullptr,
+                  /*PreambleParsedCallback=*/nullptr,
                   /*UpdateDebounce=*/std::chrono::milliseconds(50));
 
     std::vector<std::string> Files;
diff --git a/clang-tools-extra/unittests/clangd/TestTU.cpp b/clang-tools-extra/unittests/clangd/TestTU.cpp
index 66a9ba1e72f..864a6906375 100644
--- a/clang-tools-extra/unittests/clangd/TestTU.cpp
+++ b/clang-tools-extra/unittests/clangd/TestTU.cpp
@@ -43,7 +43,7 @@ ParsedAST TestTU::build() const {
 
 SymbolSlab TestTU::headerSymbols() const {
   auto AST = build();
-  return indexAST(&AST);
+  return indexAST(AST.getASTContext(), AST.getPreprocessorPtr());
 }
 
 std::unique_ptr<SymbolIndex> TestTU::index() const {