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

[ilya-biryukov]

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.

[malaperle]

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

I'll be adding those soon for workspace symbols... And also for document symbols.

[ilya-biryukov]

In D47272#1109909, @malaperle wrote:

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

I'll be adding those soon for workspace symbols... And also for document symbols.

I can add extra code to build pieces for the AST later. This is not hard to do, but would require rearranging some code a bit more.
Will try to send the change tomorrow for review tomorrow. Does that SG?

[ilya-biryukov]

@malaperle, a slight offtopic. I was wondering which completion options do you use for clangd? Defaults?

[malaperle]

In D47272#1109914, @ilya-biryukov wrote:

In D47272#1109909, @malaperle wrote:

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

I'll be adding those soon for workspace symbols... And also for document symbols.

I can add extra code to build pieces for the AST later. This is not hard to do, but would require rearranging some code a bit more.
Will try to send the change tomorrow for review tomorrow. Does that SG?

Sounds good. Doesn't have to be tomorrow :) Just making sure we're not on incompatible paths.

[ilya-biryukov]

• Added forgotten renames

[ioeric]

Would it be possible to add some integration tests for file index plus preamble?

clangd/index/FileIndex.cpp

37

Would this give us decls in namespaces? It seems that AST.getTranslationUnitDecl()->decls() are only decls that are immediately inside the TU context.

[ilya-biryukov]

In D47272#1110958, @ioeric wrote:

Would it be possible to add some integration tests for file index plus preamble?

Sure, will do

clangd/index/FileIndex.cpp

37

I'll double check and add a test for that, but I think the indexer visits the namespaces that we provide, I'm not sure if we have tests though.

[ioeric]

clangd/index/FileIndex.cpp

37

Oh, right! In that case, it seems that a single TU decl which contains all interesting decls would be sufficient? Is it possible that preamble doesn't have all available top level decls in TU decl? This seems to be possible for ParsedAST, which is why we used getTopLevelDecls, I think?

[ilya-biryukov]

• Added a test for preamble rebuilds.

[ilya-biryukov]

Added the test.

clangd/index/FileIndex.cpp

37

In that case, it seems that a single TU decl which contains all interesting decls would be sufficient?

I've tried doing the TU decl only and it does not work, because indexTopLevelDecl checks that location of decl that you pass to it is valid xD. Obviously, that's not the case for the TU decl. But it does visit the namespaces recursively, new test also checks for that, so we should be fine with the current code.

This seems to be possible for ParsedAST, which is why we used getTopLevelDecls, I think?

getTopLevelDecls was copied from the ASTUnit without much thought.
IIRC, it allows less deserialization from preamble when traversing the AST for things like "go to def", etc. This is definitely not something we're optimizing for here, since we run on the preamble AST itself.
Now that I think about it, I don't think we need to store top-level decls of preamble at all for any of our features...

[ioeric]

lg. Thanks for adding the test!

clangd/index/FileIndex.cpp

37

Sounds good. Thanks for the clarification!

clangd/ClangdServer.cpp

 #include "FindSymbols.h"
 #include "Headers.h"
 #include "SourceCode.h"
 #include "XRefs.h"
 #include "index/Merge.h"
 #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"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
: CompileArgs(CDB, Opts.ResourceDir ? Opts.ResourceDir->str()
       DiagConsumer(DiagConsumer), FSProvider(FSProvider),
       FileIdx(Opts.BuildDynamicSymbolIndex ? new FileIndex() : nullptr),
       PCHs(std::make_shared<PCHContainerOperations>()),
       // Pass a callback into `WorkScheduler` to extract symbols from a newly
       // parsed file and rebuild the file index synchronously each time an AST
       // 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,
+              ? [this](PathRef Path, ASTContext &AST,
+                       std::shared_ptr<Preprocessor>
+                           PP) { FileIdx->update(Path, &AST, std::move(PP)); }
+              : PreambleParsedCallback(),
   if (FileIdx && Opts.StaticIndex) {
     MergedIndex = mergeIndex(FileIdx.get(), Opts.StaticIndex);
     Index = MergedIndex.get();
   } else if (FileIdx)
     Index = FileIdx.get();
   else if (Opts.StaticIndex)
     Index = Opts.StaticIndex;
   else

clangd/ClangdUnit.h

 
 #include "Diagnostics.h"
 #include "Function.h"
 #include "Headers.h"
 #include "Path.h"
 #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"
 #include <memory>
 #include <string>
 #include <vector>
 
 namespace llvm {
private:
 
   // Data, stored after parsing.
   std::vector<Diag> Diags;
   std::vector<const Decl *> TopLevelDecls;
   bool PreambleDeclsDeserialized;
   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.
 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.
   llvm::Optional<std::vector<Diag>> rebuild(ParseInputs &&Inputs);
   /// Returns the last built preamble.
   const std::shared_ptr<const PreambleData> &getPreamble() const;
   /// Returns the last built AST.
   ParsedAST *getAST() const;
private:
   /// The last parsed AST.
   llvm::Optional<ParsedAST> AST;
   /// The last built Preamble.
   std::shared_ptr<const PreambleData> Preamble;
   /// Utility class required by clang
   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.
 /// May be invalid if Pos is, or if there's no identifier.
 SourceLocation getBeginningOfIdentifier(ParsedAST &Unit, const Position &Pos,
                                         const FileID FID);
 
 /// For testing/debugging purposes. Note that this method deserializes all
 /// unserialized Decls, so use with care.
 void dumpAST(ParsedAST &AST, llvm::raw_ostream &OS);
 
 } // namespace clangd
 } // namespace clang
 #endif

clangd/ClangdUnit.cpp

 //===--- ClangdUnit.cpp -----------------------------------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===---------------------------------------------------------------------===//
 
 #include "ClangdUnit.h"
 #include "Compiler.h"
 #include "Diagnostics.h"
 #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"
 #include "clang/Frontend/Utils.h"
 #include "clang/Index/IndexDataConsumer.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/Preprocessor.h"
 #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"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 
 using namespace clang::clangd;
 using namespace clang;
protected:
   }
 
 private:
   std::vector<const Decl *> TopLevelDecls;
 };
 
 class CppFilePreambleCallbacks : public PreambleCallbacks {
 public:
+  CppFilePreambleCallbacks(PathRef File, PreambleParsedCallback ParsedCallback)
+      : File(File), ParsedCallback(ParsedCallback) {}
+
   std::vector<serialization::DeclID> takeTopLevelDeclIDs() {
     return std::move(TopLevelDeclIDs);
   }
 
   std::vector<Inclusion> takeInclusions() { return std::move(Inclusions); }
 
   void AfterPCHEmitted(ASTWriter &Writer) override {
     TopLevelDeclIDs.reserve(TopLevelDecls.size());
     for (Decl *D : TopLevelDecls) {
       // Invalid top-level decls may not have been serialized.
       if (D->isInvalidDecl())
         continue;
       TopLevelDeclIDs.push_back(Writer.getDeclID(D));
     }
   }
 
+  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))
         continue;
       TopLevelDecls.push_back(D);
     }
   }
 
   void BeforeExecute(CompilerInstance &CI) override {
     SourceMgr = &CI.getSourceManager();
   }
 
   std::unique_ptr<PPCallbacks> createPPCallbacks() override {
     assert(SourceMgr && "SourceMgr must be set at this point");
     return collectInclusionsInMainFileCallback(
         *SourceMgr,
         [this](Inclusion Inc) { Inclusions.push_back(std::move(Inc)); });
   }
 
 private:
+  PathRef File;
+  PreambleParsedCallback ParsedCallback;
   std::vector<Decl *> TopLevelDecls;
   std::vector<serialization::DeclID> TopLevelDeclIDs;
   std::vector<Inclusion> Inclusions;
   SourceManager *SourceMgr = nullptr;
 };
 
 } // namespace
 
: Preamble(std::move(Preamble)), Clang(std::move(Clang)),
       TopLevelDecls(std::move(TopLevelDecls)), PreambleDeclsDeserialized(false),
       Inclusions(std::move(Inclusions)) {
   assert(this->Clang);
   assert(this->Action);
 }
 
 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);
 }
 
 llvm::Optional<std::vector<Diag>> CppFile::rebuild(ParseInputs &&Inputs) {
   log("Rebuilding file " + FileName + " with command [" +
       Inputs.CompileCommand.Directory + "] " +
       llvm::join(Inputs.CompileCommand.CommandLine, " "));
 
llvm::Optional<std::vector<Diag>> CppFile::rebuild(ParseInputs &&Inputs) {
   std::vector<Diag> Diagnostics;
   if (NewAST) {
     // Collect diagnostics from both the preamble and the AST.
     if (NewPreamble)
       Diagnostics = NewPreamble->Diags;
     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);
   Preamble = std::move(NewPreamble);
   AST = std::move(NewAST);
   return Diagnostics;
 }
 
IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
       CompilerInstance::createDiagnostics(&CI.getDiagnosticOpts(),
                                           &PreambleDiagnostics, false);
 
   // Skip function bodies when building the preamble to speed up building
   // the preamble and make it smaller.
   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);
 
   // When building the AST for the main file, we do want the function
   // bodies.
   CI.getFrontendOpts().SkipFunctionBodies = false;
 

clangd/TUScheduler.h

 /// TUScheduler is not thread-safe, only one thread should be providing updates
 /// and scheduling tasks.
 /// Callbacks are run on a threadpool and it's appropriate to do slow work in
 /// them. Each task has a name, used for tracing (should be UpperCamelCase).
 /// FIXME(sammccall): pull out a scheduler options struct.
 class TUScheduler {
 public:
   TUScheduler(unsigned AsyncThreadsCount, bool StorePreamblesInMemory,
-              ASTParsedCallback ASTCallback,
+              PreambleParsedCallback PreambleCallback,
               std::chrono::steady_clock::duration UpdateDebounce);
   ~TUScheduler();
 
   /// Returns estimated memory usage for each of the currently open files.
   /// The order of results is unspecified.
   std::vector<std::pair<Path, std::size_t>> getUsedBytesPerFile() const;
 
   /// Schedule an update for \p File. Adds \p File to a list of tracked files if
public:
   bool blockUntilIdle(Deadline D) const;
 
 private:
   /// This class stores per-file data in the Files map.
   struct FileData;
 
   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
   // asynchronously.
   llvm::Optional<AsyncTaskRunner> PreambleTasks;
   llvm::Optional<AsyncTaskRunner> WorkerThreads;
   std::chrono::steady_clock::duration UpdateDebounce;
 };
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/TUScheduler.cpp

struct TUScheduler::FileData {
   /// Latest inputs, passed to TUScheduler::update().
   std::string Contents;
   tooling::CompileCommand Command;
   ASTWorkerHandle Worker;
 };
 
 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();
     WorkerThreads.emplace();
   }
 }
 
 TUScheduler::~TUScheduler() {

 void TUScheduler::update(PathRef File, ParseInputs Inputs,
                          WantDiagnostics WantDiags,
                          UniqueFunction<void(std::vector<Diag>)> OnUpdated) {
   std::unique_ptr<FileData> &FD = Files[File];
   if (!FD) {
     // 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)});
   } else {
     FD->Contents = Inputs.Contents;
     FD->Command = Inputs.CompileCommand;
   }
   FD->Worker->update(std::move(Inputs), WantDiags, std::move(OnUpdated));

clangd/index/FileIndex.h

 //===---------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H
 
 #include "../ClangdUnit.h"
 #include "Index.h"
 #include "MemIndex.h"
+#include "clang/Lex/Preprocessor.h"
 
 namespace clang {
 namespace clangd {
 
 /// \brief A container of Symbols from several source files. It can be updated
 /// at source-file granularity, replacing all symbols from one file with a new
 /// set.
 ///
private:
   /// \brief Stores the latest snapshots for all active files.
   llvm::StringMap<std::shared_ptr<SymbolSlab>> FileToSlabs;
 };
 
 /// \brief This manages symbls from files and an in-memory index on all symbols.
 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,
             llvm::function_ref<void(const Symbol &)> Callback) const override;
 
   void lookup(const LookupRequest &Req,
               llvm::function_ref<void(const Symbol &)> Callback) const override;
 
 private:
   FileSymbols FSymbols;
   MemIndex Index;
 };
 
 /// 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
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H

clangd/index/FileIndex.cpp

 //===--- FileIndex.cpp - Indexes for files. ------------------------ C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #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
   // is not trivial given the current way of collecting symbols: we only have
   // AST at this point, but we also need preprocessor callbacks (e.g.
   // CommentHandler for IWYU pragma) to canonicalize includes.
   CollectorOpts.CollectIncludePath = false;
   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();
 }
 
 void FileSymbols::update(PathRef Path, std::unique_ptr<SymbolSlab> Slab) {
   std::lock_guard<std::mutex> Lock(Mutex);
   if (!Slab)
     FileToSlabs.erase(Path);
   else
auto Snap = std::make_shared<Snapshot>();
     }
   }
   auto *Pointers = &Snap->Pointers;
   // Use aliasing constructor to keep the snapshot alive along with the
   // pointers.
   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();
   Index.build(std::move(Symbols));
 }
 
 bool FileIndex::fuzzyFind(
     const FuzzyFindRequest &Req,

unittests/clangd/FileIndexTests.cpp

 //===-- FileIndexTests.cpp  ---------------------------*- C++ -*-----------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
+#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"
 
 using testing::UnorderedElementsAre;
 
 namespace clang {
 namespace clangd {
 

 
 // Adds Basename.cpp, which includes Basename.h, which contains Code.
 void update(FileIndex &M, llvm::StringRef Basename, llvm::StringRef Code) {
   TestTU File;
   File.Filename = (Basename + ".cpp").str();
   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) {
   FileIndex M;
   update(M, "f1", "namespace ns { void f() {} class X {}; }");
 
   FuzzyFindRequest Req;
   Req.Query = "";
TEST(FileIndexTest, RemoveAST) {
   FileIndex M;
   update(M, "f1", "namespace ns { void f() {} class X {}; }");
 
   FuzzyFindRequest Req;
   Req.Query = "";
   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());
 }
 
 TEST(FileIndexTest, IgnoreClassMembers) {
   FileIndex M;
   update(M, "f1", "class X { static int m1; int m2; static void f(); };");
 
   FuzzyFindRequest Req;
if (Sym.Name == "make_vector") {
       EXPECT_EQ(Sym.CompletionPlainInsertText, "make_vector");
       SeenMakeVector = true;
     }
   });
   EXPECT_TRUE(SeenVector);
   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

unittests/clangd/TUSchedulerTests.cpp

 
 private:
   MockCompilationDatabase CDB;
 };
 
 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");
   Files[Added] = "";
 
   auto Missing = testPath("missing.cpp");
   Files[Missing] = "";
 
TEST_F(TUSchedulerTests, WantDiagnostics) {
   std::atomic<int> CallbackCount(0);
   {
     // To avoid a racy test, don't allow tasks to actualy run on the worker
     // thread until we've scheduled them all.
     Notification Ready;
     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,
              [&](std::vector<Diag>) { Ready.wait(); });
 
     S.update(Path, getInputs(Path, "request diags"), WantDiagnostics::Yes,
              [&](std::vector<Diag> Diags) { ++CallbackCount; });
     S.update(Path, getInputs(Path, "auto (clobbered)"), WantDiagnostics::Auto,
TEST_F(TUSchedulerTests, WantDiagnostics) {
   EXPECT_EQ(2, CallbackCount);
 }
 
 TEST_F(TUSchedulerTests, Debounce) {
   std::atomic<int> CallbackCount(0);
   {
     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");
     S.update(Path, getInputs(Path, "auto (debounced)"), WantDiagnostics::Auto,
              [&](std::vector<Diag> Diags) {
                ADD_FAILURE() << "auto should have been debounced and canceled";
              });
     std::this_thread::sleep_for(std::chrono::milliseconds(200));
TEST_F(TUSchedulerTests, ManyUpdates) {
   int TotalASTReads = 0;
   int TotalPreambleReads = 0;
   int TotalUpdates = 0;
 
   // Run TUScheduler and collect some stats.
   {
     TUScheduler S(getDefaultAsyncThreadsCount(),
                   /*StorePreamblesInMemory=*/true,
-                  /*ASTParsedCallback=*/nullptr,
+                  /*PreambleParsedCallback=*/nullptr,
                   /*UpdateDebounce=*/std::chrono::milliseconds(50));
 
     std::vector<std::string> Files;
     for (int I = 0; I < FilesCount; ++I) {
       std::string Name = "foo" + std::to_string(I) + ".cpp";
       Files.push_back(testPath(Name));
       this->Files[Files.back()] = "";
     }

unittests/clangd/TestTU.cpp

if (!AST.hasValue()) {
     ADD_FAILURE() << "Failed to build code:\n" << Code;
     llvm_unreachable("Failed to build TestTU!");
   }
   return std::move(*AST);
 }
 
 SymbolSlab TestTU::headerSymbols() const {
   auto AST = build();
-  return indexAST(&AST);
+  return indexAST(AST.getASTContext(), AST.getPreprocessorPtr());
 }
 
 std::unique_ptr<SymbolIndex> TestTU::index() const {
   return MemIndex::build(headerSymbols());
 }
 
 // Look up a symbol by qualified name, which must be unique.
 const Symbol &findSymbol(const SymbolSlab &Slab, llvm::StringRef QName) {