[clangd] Replace ASTUnit with manual AST management.

[ilya-biryukov]

Summary

This refactoring does not aim to introduce any significant changes to
the behaviour of clangd to keep the change as simple as possible.

[ilya-biryukov]

Depends on change to cfe (https://reviews.llvm.org/D35405) to compile properly.

[ilya-biryukov]

• Replaced TODO with FIXME in a comment.

[malaperle]

Could you explain what the goal of this change is? It would help understand how it will impact the indexing work I am currently doing.

[ilya-biryukov]

The idea is to allows us changing the way we manage/rebuild PCHs and ASTs.
ASTUnit is used for many things and has a fairly complicated and verbose interface and does a lot of mutations all other the place inside it, so making changes to it is not particularly easy.
On the other hand, it doesn't add a lot of value for clangd specifically, since we use a very simple subset of it.

The next step I planned was to allow code completion to be run in parallel with recomputing preamble/computing diagnostics. (And in general, allow multiple TUs to be processed in parallel).
That would probably involve changes to the interface, so we should definitely sync with you on the work you've been doing.
And we were also planning to look into indexing at a later point, so discussing the design there might be interesting as well.

[malaperle]

In D35406#809609, @ilya-biryukov wrote:

The idea is to allows us changing the way we manage/rebuild PCHs and ASTs.
ASTUnit is used for many things and has a fairly complicated and verbose interface and does a lot of mutations all other the place inside it, so making changes to it is not particularly easy.
On the other hand, it doesn't add a lot of value for clangd specifically, since we use a very simple subset of it.

The next step I planned was to allow code completion to be run in parallel with recomputing preamble/computing diagnostics. (And in general, allow multiple TUs to be processed in parallel).
That would probably involve changes to the interface, so we should definitely sync with you on the work you've been doing.
And we were also planning to look into indexing at a later point, so discussing the design there might be interesting as well.

Thanks a lot for the additional information.

For synching with what I am doing. I am "only" looking right now at the modeling of the index and its on-disk storage, not really on the speeding up of the parsing of the TUs (i.e. the input of the index). I use index::IndexDataConsumer to feed the index and I see that in your change it is still pretty much used the same way as before so there is no problem. I will post on the indexing thread from before with a more detailed proposal from before once it is a bit further along.

[ilya-biryukov]

In D35406#809623, @malaperle wrote:

For synching with what I am doing. I am "only" looking right now at the modeling of the index and its on-disk storage, not really on the speeding up of the parsing of the TUs (i.e. the input of the index). I use index::IndexDataConsumer to feed the index and I see that in your change it is still pretty much used the same way as before so there is no problem. I will post on the indexing thread from before with a more detailed proposal from before once it is a bit further along.

Great that this change is not interfering. Eager to see your proposal for the indexing.

[krasimir]

clangd/ClangdUnit.cpp

167

Why DisableFree?

251

What's the purpose of this local scope here?
Wouldn't &CommandLineDiagsConsumer point to garbage after the end of this local scope?

262

We can not compute Bounds, if (Preamble).

268

Here PreambleDiagsEngine holds a pointer to the local variable PreambleDiagnosticsConsumer.
Next, BuiltPreamble holds a reference to PreambleDiagsEngine.
Next, Preamble is created by moving BuiltPreamble.
Doesn't then Preamble hold a transitive reference to junk after this local scope?

403

what's the purpose of this local scope?

690

Why do we need to ensure that Decls are deserialized?

clangd/ClangdUnit.h

77

Why is this a separate class and not just part of ClangdUnit?

[ilya-biryukov]

Added a few comments.

clangd/ClangdUnit.cpp

167

We rely on CompilerInstance to free the resources. It won't do that if DisableFree is set to true. Added a comment about that.

251

This DiagnosticsEngine is not stored in CompilerInvocation created here, it is used for reporting errors during command-line parsing.
Since we don't show those errors to the user, we discard them.

Local scope is there to avoid referencing CommandLineDiagsEngine later (as it will discard diagnostics, and we actually want to have them).

262

When Preamble is set, we need Bounds for Preamble.CanReuse.
When Preamble is not set, we need Bounds for PrecompiledPreamble::Build.

So we need them in both cases.

268

I don't think BuiltPreamble holds a reference to PreambleDiagsEngine. Here is a list of fields of PrecompiledPreamble class:

//...
TempPCHFile PCHFile;
llvm::StringMap<PreambleFileHash> FilesInPreamble;
std::vector<char> PreambleBytes;
bool PreambleEndsAtStartOfLine;

403

The DiagnosticsEngine used here uses default DiagnosticOptions.
A one that's properly configured after reading command-line arguments is created below, inside a prepareCompilerInstance call.
To avoid using DiagnosticsEngine configured by default, it is put into a local scope.

690

These Decls are used in findDefinitions (specifically, they are passed to indexTopLevelDecls). I've opted for mimicing the ASTUnit behaviour here.

It is possible that visiting only local decls (i.e. not from Preamble) will not break anything. But that requires thorough testing and digging into how PCHs works (one question is, if it is possible to have decls from the same file as part of the PCH). I actually want to investigate if we actually need all that deserializing logic here, but would prefer to do that with a separate change to keep semantic changes related to this commit as small as possible.

PS. Alternatively, we could just visit all decls in ASTContext whenever we need them for findDefinitions, but a comment in ASTUnit says it's slower.

clangd/ClangdUnit.h

77

For managing the lifetime of an AST. Specifically:

• there is a destructor here that calls EndSourceFile to free AST-related resources,
• all getters return references to objects that either have pointers to resources, managed by ParsedAST (getASTContext, getPreprocessor, etc) or semantically tied to this specific parsing result(getDiagnostics, PendingTopLevelDecls).

ClangdUnit is managing two things at the moment: a preamble and an AST. Since the lifetime of those two is different, it is arguably a good idea to separate them.

[krasimir]

Looks good!

[petarj]

Can you check if this change introduced clang-x86-windows-msvc2015 buildbot failure?

[ilya-biryukov]

In D35406#821567, @petarj wrote:

Can you check if this change introduced clang-x86-windows-msvc2015 buildbot failure?

It did. Was fixed in r308959.

clangd/ClangdServer.cpp

assert(FileContents.Draft &&
            "codeComplete is called for non-added document");
 
     DraftStorage = std::move(*FileContents.Draft);
     OverridenContents = DraftStorage;
   }
 
   std::vector<CompletionItem> Result;
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
-  // It would be nice to use runOnUnitWithoutReparse here, but we can't
-  // guarantee the correctness of code completion cache here if we don't do the
-  // reparse.
-  Units.runOnUnit(File, *OverridenContents, ResourceDir, CDB, PCHs,
-                  TaggedFS.Value, [&](ClangdUnit &Unit) {
-                    Result = Unit.codeComplete(*OverridenContents, Pos,
-                                               TaggedFS.Value);
+  Units.runOnUnitWithoutReparse(File, *OverridenContents, ResourceDir, CDB,
+                                PCHs, TaggedFS.Value, [&](ClangdUnit &Unit) {
+                                  Result = Unit.codeComplete(
+                                      *OverridenContents, Pos, TaggedFS.Value);
                                 });
   return make_tagged(std::move(Result), TaggedFS.Tag);
 }
 
 std::vector<tooling::Replacement> ClangdServer::formatRange(PathRef File,
                                                             Range Rng) {
   std::string Code = getDocument(File);
 
   size_t Begin = positionToOffset(Code, Rng.start);

clangd/ClangdUnit.h

 //===--- ClangdUnit.h -------------------------------------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===---------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNIT_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNIT_H
 
 #include "Path.h"
 #include "Protocol.h"
 #include "clang/Frontend/ASTUnit.h"
+#include "clang/Frontend/PrecompiledPreamble.h"
+#include "clang/Serialization/ASTBitCodes.h"
+#include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include <memory>
 
 namespace llvm {
 class raw_ostream;
 }
 
 namespace clang {
public:
   /// located in the current file.
   std::vector<DiagWithFixIts> getLocalDiagnostics() const;
 
   /// For testing/debugging purposes. Note that this method deserializes all
   /// unserialized Decls, so use with care.
   void dumpAST(llvm::raw_ostream &OS) const;
 
 private:
+  /// Stores and provides access to parsed AST.
+  class ParsedAST {
+  public:
+    /// Attempts to run Clang and store parsed AST. If \p Preamble is non-null
+    /// it is reused during parsing.
+    static llvm::Optional<ParsedAST>
+    Build(std::unique_ptr<clang::CompilerInvocation> CI,
+          const PrecompiledPreamble *Preamble,
+          ArrayRef<serialization::DeclID> PreambleDeclIDs,
+          std::unique_ptr<llvm::MemoryBuffer> Buffer,
+          std::shared_ptr<PCHContainerOperations> PCHs,
+          IntrusiveRefCntPtr<vfs::FileSystem> VFS);
+
+    ParsedAST(ParsedAST &&Other);
+    ParsedAST &operator=(ParsedAST &&Other);
+
+    ~ParsedAST();
+
+    ASTContext &getASTContext();
+    const ASTContext &getASTContext() const;
+
+    Preprocessor &getPreprocessor();
+    const Preprocessor &getPreprocessor() const;
+
+    ArrayRef<const Decl *> getTopLevelDecls();
+
+    const std::vector<DiagWithFixIts> &getDiagnostics() const;
+
+  private:
+    ParsedAST(std::unique_ptr<CompilerInstance> Clang,
+              std::unique_ptr<FrontendAction> Action,
+              std::vector<const Decl *> TopLevelDecls,
+              std::vector<serialization::DeclID> PendingTopLevelDecls,
+              std::vector<DiagWithFixIts> Diags);
+
+  private:
+    void ensurePreambleDeclsDeserialized();
+
+    // 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;
+
+    // Data, stored after parsing.
+    std::vector<DiagWithFixIts> Diags;
+    std::vector<const Decl *> TopLevelDecls;
+    std::vector<serialization::DeclID> PendingTopLevelDecls;
+  };
+
+  // Store Preamble and all associated data
+  struct PreambleData {
+    PreambleData(PrecompiledPreamble Preamble,
+                 std::vector<serialization::DeclID> TopLevelDeclIDs,
+                 std::vector<DiagWithFixIts> Diags);
+
+    PrecompiledPreamble Preamble;
+    std::vector<serialization::DeclID> TopLevelDeclIDs;
+    std::vector<DiagWithFixIts> Diags;
+  };
+
+  SourceLocation getBeginningOfIdentifier(const Position &Pos,
+                                          const FileEntry *FE) const;
+
   Path FileName;
-  std::unique_ptr<ASTUnit> Unit;
-  std::shared_ptr<PCHContainerOperations> PCHs;
+  tooling::CompileCommand Command;
 
-  SourceLocation getBeginningOfIdentifier(const Position& Pos, const FileEntry* FE) const;
+  llvm::Optional<PreambleData> Preamble;
+  llvm::Optional<ParsedAST> Unit;
+
+  std::shared_ptr<PCHContainerOperations> PCHs;
 };
 
 } // 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 "clang/Frontend/ASTUnit.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
+#include "clang/Frontend/FrontendActions.h"
 #include "clang/Frontend/Utils.h"
-#include "clang/Index/IndexingAction.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/Lex/PreprocessorOptions.h"
+#include "clang/Sema/Sema.h"
+#include "clang/Serialization/ASTWriter.h"
 #include "clang/Tooling/CompilationDatabase.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/CrashRecoveryContext.h"
 #include "llvm/Support/Format.h"
 
 #include <algorithm>
 
 using namespace clang::clangd;
 using namespace clang;
 
+namespace {
+
+class DeclTrackingASTConsumer : public ASTConsumer {
+public:
+  DeclTrackingASTConsumer(std::vector<const Decl *> &TopLevelDecls)
+      : TopLevelDecls(TopLevelDecls) {}
+
+  bool HandleTopLevelDecl(DeclGroupRef DG) override {
+    for (const Decl *D : DG) {
+      // ObjCMethodDecl are not actually top-level decls.
+      if (isa<ObjCMethodDecl>(D))
+        continue;
+
+      TopLevelDecls.push_back(D);
+    }
+    return true;
+  }
+
+private:
+  std::vector<const Decl *> &TopLevelDecls;
+};
+
+class ClangdFrontendAction : public SyntaxOnlyAction {
+public:
+  std::vector<const Decl *> takeTopLevelDecls() {
+    return std::move(TopLevelDecls);
+  }
+
+protected:
+  std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
+                                                 StringRef InFile) override {
+    return llvm::make_unique<DeclTrackingASTConsumer>(/*ref*/ TopLevelDecls);
+  }
+
+private:
+  std::vector<const Decl *> TopLevelDecls;
+};
+
+class ClangdUnitPreambleCallbacks : public PreambleCallbacks {
+public:
+  std::vector<serialization::DeclID> takeTopLevelDeclIDs() {
+    return std::move(TopLevelDeclIDs);
+  }
+
+  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 HandleTopLevelDecl(DeclGroupRef DG) override {
+    for (Decl *D : DG) {
+      if (isa<ObjCMethodDecl>(D))
+        continue;
+      TopLevelDecls.push_back(D);
+    }
+  }
+
+private:
+  std::vector<Decl *> TopLevelDecls;
+  std::vector<serialization::DeclID> TopLevelDeclIDs;
+};
+
+/// Convert from clang diagnostic level to LSP severity.
+static int getSeverity(DiagnosticsEngine::Level L) {
+  switch (L) {
+  case DiagnosticsEngine::Remark:
+    return 4;
+  case DiagnosticsEngine::Note:
+    return 3;
+  case DiagnosticsEngine::Warning:
+    return 2;
+  case DiagnosticsEngine::Fatal:
+  case DiagnosticsEngine::Error:
+    return 1;
+  case DiagnosticsEngine::Ignored:
+    return 0;
+  }
+  llvm_unreachable("Unknown diagnostic level!");
+}
+
+llvm::Optional<DiagWithFixIts> toClangdDiag(StoredDiagnostic D) {
+  auto Location = D.getLocation();
+  if (!Location.isValid() || !Location.getManager().isInMainFile(Location))
+    return llvm::None;
+
+  Position P;
+  P.line = Location.getSpellingLineNumber() - 1;
+  P.character = Location.getSpellingColumnNumber();
+  Range R = {P, P};
+  clangd::Diagnostic Diag = {R, getSeverity(D.getLevel()), D.getMessage()};
+
+  llvm::SmallVector<tooling::Replacement, 1> FixItsForDiagnostic;
+  for (const FixItHint &Fix : D.getFixIts()) {
+    FixItsForDiagnostic.push_back(clang::tooling::Replacement(
+        Location.getManager(), Fix.RemoveRange, Fix.CodeToInsert));
+  }
+  return DiagWithFixIts{Diag, std::move(FixItsForDiagnostic)};
+}
+
+class StoreDiagsConsumer : public DiagnosticConsumer {
+public:
+  StoreDiagsConsumer(std::vector<DiagWithFixIts> &Output) : Output(Output) {}
+
+  void HandleDiagnostic(DiagnosticsEngine::Level DiagLevel,
+                        const clang::Diagnostic &Info) override {
+    DiagnosticConsumer::HandleDiagnostic(DiagLevel, Info);
+
+    if (auto convertedDiag = toClangdDiag(StoredDiagnostic(DiagLevel, Info)))
+      Output.push_back(std::move(*convertedDiag));
+  }
+
+private:
+  std::vector<DiagWithFixIts> &Output;
+};
+
+class EmptyDiagsConsumer : public DiagnosticConsumer {
+public:
+  void HandleDiagnostic(DiagnosticsEngine::Level DiagLevel,
+                        const clang::Diagnostic &Info) override {}
+};
+
+std::unique_ptr<CompilerInvocation>
+createCompilerInvocation(ArrayRef<const char *> ArgList,
+                         IntrusiveRefCntPtr<DiagnosticsEngine> Diags,
+                         IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
+  auto CI = createInvocationFromCommandLine(ArgList, std::move(Diags),
+                                            std::move(VFS));
+  CI->getFrontendOpts().DisableFree = false;
+  return CI;
+}
+
+/// Creates a CompilerInstance from \p CI, with main buffer overriden to \p
+/// Buffer and arguments to read the PCH from \p Preamble, if \p Preamble is not
+/// null. Note that vfs::FileSystem inside returned instance may differ from \p
+/// VFS if additional file remapping were set in command-line arguments.
+/// On some errors, returns null. When non-null value is returned, it's expected
+/// to be consumed by the FrontendAction as it will have a pointer to the \p
+/// Buffer that will only be deleted if BeginSourceFile is called.
+std::unique_ptr<CompilerInstance>
+prepareCompilerInstance(std::unique_ptr<clang::CompilerInvocation> CI,
+                        const PrecompiledPreamble *Preamble,
+                        std::unique_ptr<llvm::MemoryBuffer> Buffer,
+                        std::shared_ptr<PCHContainerOperations> PCHs,
+                        IntrusiveRefCntPtr<vfs::FileSystem> VFS,
+                        DiagnosticConsumer &DiagsClient) {
+  assert(VFS && "VFS is null");
+  assert(!CI->getPreprocessorOpts().RetainRemappedFileBuffers &&
+         "Setting RetainRemappedFileBuffers to true will cause a memory leak "
+         "of ContentsBuffer");
+
+  // NOTE: we use Buffer.get() when adding remapped files, so we have to make
+  // sure it will be released if no error is emitted.
+  if (Preamble) {
+    Preamble->AddImplicitPreamble(*CI, Buffer.get());
+  } else {
+    CI->getPreprocessorOpts().addRemappedFile(
+        CI->getFrontendOpts().Inputs[0].getFile(), Buffer.get());
+  }
+
+  auto Clang = llvm::make_unique<CompilerInstance>(PCHs);
+  Clang->setInvocation(std::move(CI));
+  Clang->createDiagnostics(&DiagsClient, false);
+
+  if (auto VFSWithRemapping = createVFSFromCompilerInvocation(
+          Clang->getInvocation(), Clang->getDiagnostics(), VFS))
+    VFS = VFSWithRemapping;
+  Clang->setVirtualFileSystem(VFS);
+
+  Clang->setTarget(TargetInfo::CreateTargetInfo(
+      Clang->getDiagnostics(), Clang->getInvocation().TargetOpts));
+  if (!Clang->hasTarget())
+    return nullptr;
+
+  // RemappedFileBuffers will handle the lifetime of the Buffer pointer,
+  // release it.
+  Buffer.release();
+  return Clang;
+}
+
+} // namespace
+
 ClangdUnit::ClangdUnit(PathRef FileName, StringRef Contents,
                        StringRef ResourceDir,
                        std::shared_ptr<PCHContainerOperations> PCHs,
                        std::vector<tooling::CompileCommand> Commands,
                        IntrusiveRefCntPtr<vfs::FileSystem> VFS)
     : FileName(FileName), PCHs(PCHs) {
   assert(!Commands.empty() && "No compile commands provided");
 
   // Inject the resource dir.
   // FIXME: Don't overwrite it if it's already there.
   Commands.front().CommandLine.push_back("-resource-dir=" +
                                          std::string(ResourceDir));
 
-  IntrusiveRefCntPtr<DiagnosticsEngine> Diags =
-      CompilerInstance::createDiagnostics(new DiagnosticOptions);
-
-  std::vector<const char *> ArgStrs;
-  for (const auto &S : Commands.front().CommandLine)
-    ArgStrs.push_back(S.c_str());
-
-  ASTUnit::RemappedFile RemappedSource(
-      FileName,
-      llvm::MemoryBuffer::getMemBufferCopy(Contents, FileName).release());
-
-  auto ArgP = &*ArgStrs.begin();
-  Unit = std::unique_ptr<ASTUnit>(ASTUnit::LoadFromCommandLine(
-      ArgP, ArgP + ArgStrs.size(), PCHs, Diags, ResourceDir,
-      /*OnlyLocalDecls=*/false, /*CaptureDiagnostics=*/true, RemappedSource,
-      /*RemappedFilesKeepOriginalName=*/true,
-      /*PrecompilePreambleAfterNParses=*/1, /*TUKind=*/TU_Prefix,
-      /*CacheCodeCompletionResults=*/true,
-      /*IncludeBriefCommentsInCodeCompletion=*/true,
-      /*AllowPCHWithCompilerErrors=*/true,
-      /*SkipFunctionBodies=*/false,
-      /*SingleFileParse=*/false,
-      /*UserFilesAreVolatile=*/false, /*ForSerialization=*/false,
-      /*ModuleFormat=*/llvm::None,
-      /*ErrAST=*/nullptr, VFS));
-  assert(Unit && "Unit wasn't created");
+  Command = std::move(Commands.front());
+  reparse(Contents, VFS);
 }
 
 void ClangdUnit::reparse(StringRef Contents,
                          IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
-  // Do a reparse if this wasn't the first parse.
-  // FIXME: This might have the wrong working directory if it changed in the
-  // meantime.
-  ASTUnit::RemappedFile RemappedSource(
-      FileName,
-      llvm::MemoryBuffer::getMemBufferCopy(Contents, FileName).release());
+  std::vector<const char *> ArgStrs;
+  for (const auto &S : Command.CommandLine)
+    ArgStrs.push_back(S.c_str());
 
-  Unit->Reparse(PCHs, RemappedSource, VFS);
+  std::unique_ptr<CompilerInvocation> CI;
+  {
+    // TOOD: store diagnostics from CommandLine when we start reporting them.
+    EmptyDiagsConsumer CommandLineDiagsConsumer;
+    IntrusiveRefCntPtr<DiagnosticsEngine> CommandLineDiagsEngine =
+        CompilerInstance::createDiagnostics(new DiagnosticOptions,
+                                            &CommandLineDiagsConsumer, false);
+    CI = createCompilerInvocation(ArgStrs, CommandLineDiagsEngine, VFS);
+  }
+  assert(CI && "Couldn't create CompilerInvocation");
+
+  std::unique_ptr<llvm::MemoryBuffer> ContentsBuffer =
+      llvm::MemoryBuffer::getMemBufferCopy(Contents, FileName);
+
+  // Rebuild the preamble if it is missing or can not be reused.
+  auto Bounds =
+      ComputePreambleBounds(*CI->getLangOpts(), ContentsBuffer.get(), 0);
+  if (!Preamble || !Preamble->Preamble.CanReuse(*CI, ContentsBuffer.get(),
+                                                Bounds, VFS.get())) {
+    std::vector<DiagWithFixIts> PreambleDiags;
+    StoreDiagsConsumer PreambleDiagnosticsConsumer(/*ref*/ PreambleDiags);
+    IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
+        CompilerInstance::createDiagnostics(
+            &CI->getDiagnosticOpts(), &PreambleDiagnosticsConsumer, false);
+    ClangdUnitPreambleCallbacks SerializedDeclsCollector;
+    auto BuiltPreamble = PrecompiledPreamble::Build(
+        *CI, ContentsBuffer.get(), Bounds, *PreambleDiagsEngine, VFS, PCHs,
+        SerializedDeclsCollector);
+    if (BuiltPreamble)
+      Preamble = PreambleData(std::move(*BuiltPreamble),
+                              SerializedDeclsCollector.takeTopLevelDeclIDs(),
+                              std::move(PreambleDiags));
+  }
+  Unit = ParsedAST::Build(
+      std::move(CI), Preamble ? &Preamble->Preamble : nullptr,
+      Preamble ? llvm::makeArrayRef(Preamble->TopLevelDeclIDs) : llvm::None,
+      std::move(ContentsBuffer), PCHs, VFS);
 }
 
 namespace {
 
 CompletionItemKind getKind(CXCursorKind K) {
   switch (K) {
   case CXCursor_MacroInstantiation:
   case CXCursor_MacroDefinition:
public:
 
   CodeCompletionTUInfo &getCodeCompletionTUInfo() override { return CCTUInfo; }
 };
 } // namespace
 
 std::vector<CompletionItem>
 ClangdUnit::codeComplete(StringRef Contents, Position Pos,
                          IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
-  CodeCompleteOptions CCO;
-  CCO.IncludeBriefComments = 1;
-  // This is where code completion stores dirty buffers. Need to free after
-  // completion.
-  SmallVector<const llvm::MemoryBuffer *, 4> OwnedBuffers;
-  SmallVector<StoredDiagnostic, 4> StoredDiagnostics;
-  IntrusiveRefCntPtr<DiagnosticsEngine> DiagEngine(
-      new DiagnosticsEngine(new DiagnosticIDs, new DiagnosticOptions));
+  std::vector<const char *> ArgStrs;
+  for (const auto &S : Command.CommandLine)
+    ArgStrs.push_back(S.c_str());
+
+  std::unique_ptr<CompilerInvocation> CI;
+  EmptyDiagsConsumer DummyDiagsConsumer;
+  {
+    IntrusiveRefCntPtr<DiagnosticsEngine> CommandLineDiagsEngine =
+        CompilerInstance::createDiagnostics(new DiagnosticOptions,
+                                            &DummyDiagsConsumer, false);
+    CI = createCompilerInvocation(ArgStrs, CommandLineDiagsEngine, VFS);
+  }
+  assert(CI && "Couldn't create CompilerInvocation");
+
+  std::unique_ptr<llvm::MemoryBuffer> ContentsBuffer =
+      llvm::MemoryBuffer::getMemBufferCopy(Contents, FileName);
+
+  // Attempt to reuse the PCH from precompiled preamble, if it was built.
+  const PrecompiledPreamble *PreambleForCompletion = nullptr;
+  if (Preamble) {
+    auto Bounds =
+        ComputePreambleBounds(*CI->getLangOpts(), ContentsBuffer.get(), 0);
+    if (Preamble->Preamble.CanReuse(*CI, ContentsBuffer.get(), Bounds,
+                                    VFS.get()))
+      PreambleForCompletion = &Preamble->Preamble;
+  }
+
+  auto Clang = prepareCompilerInstance(std::move(CI), PreambleForCompletion,
+                                       std::move(ContentsBuffer), PCHs, VFS,
+                                       DummyDiagsConsumer);
+  auto &DiagOpts = Clang->getDiagnosticOpts();
+  DiagOpts.IgnoreWarnings = true;
+
+  auto &FrontendOpts = Clang->getFrontendOpts();
+  FrontendOpts.SkipFunctionBodies = true;
+
+  FrontendOpts.CodeCompleteOpts.IncludeGlobals = true;
+  FrontendOpts.CodeCompletionAt.FileName = FileName;
+  FrontendOpts.CodeCompletionAt.Line = Pos.line + 1;
+  FrontendOpts.CodeCompletionAt.Column = Pos.character + 1;
+
   std::vector<CompletionItem> Items;
-  CompletionItemsCollector Collector(&Items, CCO);
+  Clang->setCodeCompletionConsumer(
+      new CompletionItemsCollector(&Items, FrontendOpts.CodeCompleteOpts));
 
-  ASTUnit::RemappedFile RemappedSource(
-      FileName,
-      llvm::MemoryBuffer::getMemBufferCopy(Contents, FileName).release());
-
-  IntrusiveRefCntPtr<FileManager> FileMgr(
-      new FileManager(Unit->getFileSystemOpts(), VFS));
-  IntrusiveRefCntPtr<SourceManager> SourceMgr(
-      new SourceManager(*DiagEngine, *FileMgr));
-  // CodeComplete seems to require fresh LangOptions.
-  LangOptions LangOpts = Unit->getLangOpts();
-  // The language server protocol uses zero-based line and column numbers.
-  // The clang code completion uses one-based numbers.
-  Unit->CodeComplete(FileName, Pos.line + 1, Pos.character + 1, RemappedSource,
-                     CCO.IncludeMacros, CCO.IncludeCodePatterns,
-                     CCO.IncludeBriefComments, Collector, PCHs, *DiagEngine,
-                     LangOpts, *SourceMgr, *FileMgr, StoredDiagnostics,
-                     OwnedBuffers);
-  for (const llvm::MemoryBuffer *Buffer : OwnedBuffers)
-    delete Buffer;
+  SyntaxOnlyAction Action;
+  if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0])) {
+    // FIXME(ibiryukov): log errors
     return Items;
   }
-
-namespace {
-/// Convert from clang diagnostic level to LSP severity.
-static int getSeverity(DiagnosticsEngine::Level L) {
-  switch (L) {
-  case DiagnosticsEngine::Remark:
-    return 4;
-  case DiagnosticsEngine::Note:
-    return 3;
-  case DiagnosticsEngine::Warning:
-    return 2;
-  case DiagnosticsEngine::Fatal:
-  case DiagnosticsEngine::Error:
-    return 1;
-  case DiagnosticsEngine::Ignored:
-    return 0;
+  if (!Action.Execute()) {
+    // FIXME(ibiryukov): log errors
   }
-  llvm_unreachable("Unknown diagnostic level!");
+  Action.EndSourceFile();
+
+  return Items;
 }
-} // namespace
 
 std::vector<DiagWithFixIts> ClangdUnit::getLocalDiagnostics() const {
-  std::vector<DiagWithFixIts> Result;
-  for (ASTUnit::stored_diag_iterator D = Unit->stored_diag_begin(),
-                                     DEnd = Unit->stored_diag_end();
-       D != DEnd; ++D) {
-    if (!D->getLocation().isValid() ||
-        !D->getLocation().getManager().isInMainFile(D->getLocation()))
-      continue;
-    Position P;
-    P.line = D->getLocation().getSpellingLineNumber() - 1;
-    P.character = D->getLocation().getSpellingColumnNumber();
-    Range R = {P, P};
-    clangd::Diagnostic Diag = {R, getSeverity(D->getLevel()), D->getMessage()};
+  if (!Unit)
+    return {}; // Parsing failed.
 
-    llvm::SmallVector<tooling::Replacement, 1> FixItsForDiagnostic;
-    for (const FixItHint &Fix : D->getFixIts()) {
-      FixItsForDiagnostic.push_back(clang::tooling::Replacement(
-          Unit->getSourceManager(), Fix.RemoveRange, Fix.CodeToInsert));
-    }
-    Result.push_back({Diag, std::move(FixItsForDiagnostic)});
-  }
+  std::vector<DiagWithFixIts> Result;
+  auto PreambleDiagsSize = Preamble ? Preamble->Diags.size() : 0;
+  const auto &Diags = Unit->getDiagnostics();
+  Result.reserve(PreambleDiagsSize + Diags.size());
+
+  if (Preamble)
+    Result.insert(Result.end(), Preamble->Diags.begin(), Preamble->Diags.end());
+  Result.insert(Result.end(), Diags.begin(), Diags.end());
   return Result;
 }
 
 void ClangdUnit::dumpAST(llvm::raw_ostream &OS) const {
+  if (!Unit) {
+    OS << "<no-ast-in-clang>";
+    return; // Parsing failed.
+  }
   Unit->getASTContext().getTranslationUnitDecl()->dump(OS, true);
 }
 
+llvm::Optional<ClangdUnit::ParsedAST>
+ClangdUnit::ParsedAST::Build(std::unique_ptr<clang::CompilerInvocation> CI,
+                             const PrecompiledPreamble *Preamble,
+                             ArrayRef<serialization::DeclID> PreambleDeclIDs,
+                             std::unique_ptr<llvm::MemoryBuffer> Buffer,
+                             std::shared_ptr<PCHContainerOperations> PCHs,
+                             IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
+
+  std::vector<DiagWithFixIts> ASTDiags;
+  StoreDiagsConsumer UnitDiagsConsumer(/*ref*/ ASTDiags);
+
+  auto Clang =
+      prepareCompilerInstance(std::move(CI), Preamble, std::move(Buffer), PCHs,
+                              VFS, /*ref*/ UnitDiagsConsumer);
+
+  // Recover resources if we crash before exiting this method.
+  llvm::CrashRecoveryContextCleanupRegistrar<CompilerInstance> CICleanup(
+      Clang.get());
+
+  auto Action = llvm::make_unique<ClangdFrontendAction>();
+  if (!Action->BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0])) {
+    // FIXME(ibiryukov): log error
+    return llvm::None;
+  }
+  if (!Action->Execute()) {
+    // FIXME(ibiryukov): log error
+  }
+
+  // UnitDiagsConsumer is local, we can not store it in CompilerInstance that
+  // has a longer lifetime.
+  Clang->getDiagnostics().setClient(new EmptyDiagsConsumer);
+
+  std::vector<const Decl *> ParsedDecls = Action->takeTopLevelDecls();
+  std::vector<serialization::DeclID> PendingDecls;
+  if (Preamble) {
+    PendingDecls.reserve(PreambleDeclIDs.size());
+    PendingDecls.insert(PendingDecls.begin(), PreambleDeclIDs.begin(),
+                        PreambleDeclIDs.end());
+  }
+
+  return ParsedAST(std::move(Clang), std::move(Action), std::move(ParsedDecls),
+                   std::move(PendingDecls), std::move(ASTDiags));
+}
+
 namespace {
+
+SourceLocation getMacroArgExpandedLocation(const SourceManager &Mgr,
+                                           const FileEntry *FE,
+                                           unsigned Offset) {
+  SourceLocation FileLoc = Mgr.translateFileLineCol(FE, 1, 1);
+  return Mgr.getMacroArgExpandedLocation(FileLoc.getLocWithOffset(Offset));
+}
+
+SourceLocation getMacroArgExpandedLocation(const SourceManager &Mgr,
+                                           const FileEntry *FE, Position Pos) {
+  SourceLocation InputLoc =
+      Mgr.translateFileLineCol(FE, Pos.line + 1, Pos.character + 1);
+  return Mgr.getMacroArgExpandedLocation(InputLoc);
+}
+
 /// Finds declarations locations that a given source location refers to.
 class DeclarationLocationsFinder : public index::IndexDataConsumer {
   std::vector<Location> DeclarationLocations;
   const SourceLocation &SearchedLocation;
-  ASTUnit &Unit;
+  const ASTContext &AST;
+  Preprocessor &PP;
+
 public:
   DeclarationLocationsFinder(raw_ostream &OS,
-      const SourceLocation &SearchedLocation, ASTUnit &Unit) :
-      SearchedLocation(SearchedLocation), Unit(Unit) {}
+                             const SourceLocation &SearchedLocation,
+                             ASTContext &AST, Preprocessor &PP)
+      : SearchedLocation(SearchedLocation), AST(AST), PP(PP) {}
 
   std::vector<Location> takeLocations() {
     // Don't keep the same location multiple times.
     // This can happen when nodes in the AST are visited twice.
     std::sort(DeclarationLocations.begin(), DeclarationLocations.end());
     auto last =
         std::unique(DeclarationLocations.begin(), DeclarationLocations.end());
     DeclarationLocations.erase(last, DeclarationLocations.end());
     return std::move(DeclarationLocations);
   }
 
   bool handleDeclOccurence(const Decl* D, index::SymbolRoleSet Roles,
       ArrayRef<index::SymbolRelation> Relations, FileID FID, unsigned Offset,
       index::IndexDataConsumer::ASTNodeInfo ASTNode) override
       {
     if (isSearchedLocation(FID, Offset)) {
       addDeclarationLocation(D->getSourceRange());
     }
     return true;
   }
 
 private:
   bool isSearchedLocation(FileID FID, unsigned Offset) const {
-    const SourceManager &SourceMgr = Unit.getSourceManager();
-    return SourceMgr.getFileOffset(SearchedLocation) == Offset
-        && SourceMgr.getFileID(SearchedLocation) == FID;
+    const SourceManager &SourceMgr = AST.getSourceManager();
+    return SourceMgr.getFileOffset(SearchedLocation) == Offset &&
+           SourceMgr.getFileID(SearchedLocation) == FID;
   }
 
-  void addDeclarationLocation(const SourceRange& ValSourceRange) {
-    const SourceManager& SourceMgr = Unit.getSourceManager();
-    const LangOptions& LangOpts = Unit.getLangOpts();
+  void addDeclarationLocation(const SourceRange &ValSourceRange) {
+    const SourceManager &SourceMgr = AST.getSourceManager();
+    const LangOptions &LangOpts = AST.getLangOpts();
     SourceLocation LocStart = ValSourceRange.getBegin();
     SourceLocation LocEnd = Lexer::getLocForEndOfToken(ValSourceRange.getEnd(),
                                                        0, SourceMgr, LangOpts);
     Position Begin;
     Begin.line = SourceMgr.getSpellingLineNumber(LocStart) - 1;
     Begin.character = SourceMgr.getSpellingColumnNumber(LocStart) - 1;
     Position End;
     End.line = SourceMgr.getSpellingLineNumber(LocEnd) - 1;
     End.character = SourceMgr.getSpellingColumnNumber(LocEnd) - 1;
     Range R = {Begin, End};
     Location L;
     L.uri = URI::fromFile(
         SourceMgr.getFilename(SourceMgr.getSpellingLoc(LocStart)));
     L.range = R;
     DeclarationLocations.push_back(L);
   }
 
   void finish() override {
     // Also handle possible macro at the searched location.
     Token Result;
-    if (!Lexer::getRawToken(SearchedLocation, Result, Unit.getSourceManager(),
-        Unit.getASTContext().getLangOpts(), false)) {
+    if (!Lexer::getRawToken(SearchedLocation, Result, AST.getSourceManager(),
+                            AST.getLangOpts(), false)) {
       if (Result.is(tok::raw_identifier)) {
-        Unit.getPreprocessor().LookUpIdentifierInfo(Result);
+        PP.LookUpIdentifierInfo(Result);
       }
-      IdentifierInfo* IdentifierInfo = Result.getIdentifierInfo();
+      IdentifierInfo *IdentifierInfo = Result.getIdentifierInfo();
       if (IdentifierInfo && IdentifierInfo->hadMacroDefinition()) {
         std::pair<FileID, unsigned int> DecLoc =
-            Unit.getSourceManager().getDecomposedExpansionLoc(SearchedLocation);
+            AST.getSourceManager().getDecomposedExpansionLoc(SearchedLocation);
         // Get the definition just before the searched location so that a macro
         // referenced in a '#undef MACRO' can still be found.
-        SourceLocation BeforeSearchedLocation = Unit.getLocation(
-            Unit.getSourceManager().getFileEntryForID(DecLoc.first),
+        SourceLocation BeforeSearchedLocation = getMacroArgExpandedLocation(
+            AST.getSourceManager(),
+            AST.getSourceManager().getFileEntryForID(DecLoc.first),
             DecLoc.second - 1);
         MacroDefinition MacroDef =
-            Unit.getPreprocessor().getMacroDefinitionAtLoc(IdentifierInfo,
-                BeforeSearchedLocation);
-        MacroInfo* MacroInf = MacroDef.getMacroInfo();
+            PP.getMacroDefinitionAtLoc(IdentifierInfo, BeforeSearchedLocation);
+        MacroInfo *MacroInf = MacroDef.getMacroInfo();
         if (MacroInf) {
           addDeclarationLocation(
               SourceRange(MacroInf->getDefinitionLoc(),
                   MacroInf->getDefinitionEndLoc()));
         }
       }
     }
   }
 };
 } // namespace
 
 std::vector<Location> ClangdUnit::findDefinitions(Position Pos) {
-  const FileEntry *FE = Unit->getFileManager().getFile(Unit->getMainFileName());
+  if (!Unit)
+    return {}; // Parsing failed.
+
+  const SourceManager &SourceMgr = Unit->getASTContext().getSourceManager();
+  const FileEntry *FE = SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
   if (!FE)
     return {};
 
   SourceLocation SourceLocationBeg = getBeginningOfIdentifier(Pos, FE);
 
   auto DeclLocationsFinder = std::make_shared<DeclarationLocationsFinder>(
-      llvm::errs(), SourceLocationBeg, *Unit);
+      llvm::errs(), SourceLocationBeg, Unit->getASTContext(),
+      Unit->getPreprocessor());
   index::IndexingOptions IndexOpts;
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::All;
   IndexOpts.IndexFunctionLocals = true;
-  index::indexASTUnit(*Unit, DeclLocationsFinder, IndexOpts);
+
+  indexTopLevelDecls(Unit->getASTContext(), Unit->getTopLevelDecls(),
+                     DeclLocationsFinder, IndexOpts);
 
   return DeclLocationsFinder->takeLocations();
 }
 
 SourceLocation ClangdUnit::getBeginningOfIdentifier(const Position &Pos,
                                                     const FileEntry *FE) const {
+
   // The language server protocol uses zero-based line and column numbers.
   // Clang uses one-based numbers.
-  SourceLocation InputLocation = Unit->getLocation(FE, Pos.line + 1,
-      Pos.character + 1);
 
+  const ASTContext &AST = Unit->getASTContext();
+  const SourceManager &SourceMgr = AST.getSourceManager();
+
+  SourceLocation InputLocation =
+      getMacroArgExpandedLocation(SourceMgr, FE, Pos);
   if (Pos.character == 0) {
     return InputLocation;
   }
 
   // This handle cases where the position is in the middle of a token or right
   // after the end of a token. In theory we could just use GetBeginningOfToken
   // to find the start of the token at the input position, but this doesn't
   // work when right after the end, i.e. foo|.
   // So try to go back by one and see if we're still inside the an identifier
   // token. If so, Take the beginning of this token.
   // (It should be the same identifier because you can't have two adjacent
   // identifiers without another token in between.)
-  SourceLocation PeekBeforeLocation = Unit->getLocation(FE, Pos.line + 1,
-      Pos.character);
-  const SourceManager &SourceMgr = Unit->getSourceManager();
+  SourceLocation PeekBeforeLocation = getMacroArgExpandedLocation(
+      SourceMgr, FE, Position{Pos.line, Pos.character - 1});
   Token Result;
   if (Lexer::getRawToken(PeekBeforeLocation, Result, SourceMgr,
-                         Unit->getASTContext().getLangOpts(), false)) {
+                         AST.getLangOpts(), false)) {
     // getRawToken failed, just use InputLocation.
     return InputLocation;
   }
 
   if (Result.is(tok::raw_identifier)) {
     return Lexer::GetBeginningOfToken(PeekBeforeLocation, SourceMgr,
                                       Unit->getASTContext().getLangOpts());
   }
 
   return InputLocation;
 }
+
+void ClangdUnit::ParsedAST::ensurePreambleDeclsDeserialized() {
+  if (PendingTopLevelDecls.empty())
+    return;
+
+  std::vector<const Decl *> Resolved;
+  Resolved.reserve(PendingTopLevelDecls.size());
+
+  ExternalASTSource &Source = *getASTContext().getExternalSource();
+  for (serialization::DeclID TopLevelDecl : PendingTopLevelDecls) {
+    // Resolve the declaration ID to an actual declaration, possibly
+    // deserializing the declaration in the process.
+    if (Decl *D = Source.GetExternalDecl(TopLevelDecl))
+      Resolved.push_back(D);
+  }
+
+  TopLevelDecls.reserve(TopLevelDecls.size() + PendingTopLevelDecls.size());
+  TopLevelDecls.insert(TopLevelDecls.begin(), Resolved.begin(), Resolved.end());
+
+  PendingTopLevelDecls.clear();
+}
+
+ClangdUnit::ParsedAST::ParsedAST(ParsedAST &&Other) = default;
+
+ClangdUnit::ParsedAST &ClangdUnit::ParsedAST::
+operator=(ParsedAST &&Other) = default;
+
+ClangdUnit::ParsedAST::~ParsedAST() {
+  if (Action) {
+    Action->EndSourceFile();
+  }
+}
+
+ASTContext &ClangdUnit::ParsedAST::getASTContext() {
+  return Clang->getASTContext();
+}
+
+const ASTContext &ClangdUnit::ParsedAST::getASTContext() const {
+  return Clang->getASTContext();
+}
+
+Preprocessor &ClangdUnit::ParsedAST::getPreprocessor() {
+  return Clang->getPreprocessor();
+}
+
+const Preprocessor &ClangdUnit::ParsedAST::getPreprocessor() const {
+  return Clang->getPreprocessor();
+}
+
+ArrayRef<const Decl *> ClangdUnit::ParsedAST::getTopLevelDecls() {
+  ensurePreambleDeclsDeserialized();
+  return TopLevelDecls;
+}
+
+const std::vector<DiagWithFixIts> &
+ClangdUnit::ParsedAST::getDiagnostics() const {
+  return Diags;
+}
+
+ClangdUnit::ParsedAST::ParsedAST(
+    std::unique_ptr<CompilerInstance> Clang,
+    std::unique_ptr<FrontendAction> Action,
+    std::vector<const Decl *> TopLevelDecls,
+    std::vector<serialization::DeclID> PendingTopLevelDecls,
+    std::vector<DiagWithFixIts> Diags)
+    : Clang(std::move(Clang)), Action(std::move(Action)),
+      Diags(std::move(Diags)), TopLevelDecls(std::move(TopLevelDecls)),
+      PendingTopLevelDecls(std::move(PendingTopLevelDecls)) {
+  assert(this->Clang);
+  assert(this->Action);
+}
+
+ClangdUnit::PreambleData::PreambleData(
+    PrecompiledPreamble Preamble,
+    std::vector<serialization::DeclID> TopLevelDeclIDs,
+    std::vector<DiagWithFixIts> Diags)
+    : Preamble(std::move(Preamble)),
+      TopLevelDeclIDs(std::move(TopLevelDeclIDs)), Diags(std::move(Diags)) {}