[ClangD] Refactor clangd into separate components

[ilya-biryukov]

Summary

Major refactoring to split LSP implementation, Clang API calls and threading(mostly synchronization)

[klimek]

clangd/ClangdLSPServer.h

24

I'd have expected something that's called LSP server to work on the LSP protocol level (that is, have a server(iostream) equivalent method).

25–26

I like to put the public interface first, as that's what people usually look at first.

49

AFAIK usually we don't duplicate access modifiers.

clangd/ClangdUnit.h

45

One question is how we can handle when we want the old one to stick around while we do a full reparse (when #includes change), so we still have fast code completion.

clangd/ClangdUnitStore.h

22

Tiny nit: please add empty lines between classes and namespace delcs :)

30–31

Not having read enough of the rest of the patch yet, I'd expect these things to be stored on the class level instead of being passed around all the time.
Similarly, we're missing the VFS?

32

Parameters after a function are unfortunate, as they make the code flow weird when you want to pass in a lambda, but still want to set the parameter after it :(
Given how different the semantics are, I'd perhaps make these 2 functions:

// Run on the translation unit.
// If there was a previous call of runOnUnitWithContent, the content provided there will be used.
// Otherwise, the file will be read from VFS.
runOnStoredUnit(PathRef File, Func Action);

// Run on the translation unit with the given content.
// Content will be stored as in-flight code for the given translation unit only:
// - if a different translation unit is parsed that #includes File, Content will not be used;
// - subsequent calls to runOnStoredUnit will re-use Content.
runOnUnitWithContent(PathRef File, StringRef Content, Func Action);

33

This is due to a restriction in the AST unit? If we don't need to reparse, why can't we run multiple things over an AST at once?

clangd/GlobalCompilationDatabase.cpp

56

Here and elsewhere: in Clang, we use FIXME instead of TODO :)

clangd/GlobalCompilationDatabase.h

29

What's global about this?

[ilya-biryukov]

Addressed review comments

[ilya-biryukov]

clangd/ClangdLSPServer.h

24

The idea is to have a class that wraps ClangdServer and is called by clangd's ProtocolHandlers.cpp to do all the work.
I.e. it's a replacement of what used to be called 'ASTManager'. Any suggestions for a better name? ClangdLSPHandler?

It's seems hard to implement the interface you suggest right away, lots of code will have to be moved. Happy to do that in further commits, but suggest we leave this class as is in this commit.

clangd/ClangdUnit.h

45

My idea was to have two ClangdUnits for that and handle that somewhere else(i.e. ClangdUnitStore seem like a good place to do that)

clangd/ClangdUnitStore.h

30–31

The VFS changes are not in this commit, they will follow later, that's an initial refactoring of clangd to allow more gradual changes later.

The idea is that this class is an implementation detail of ClangdServer that handles synchronized access to underlying ClangdUnits(which are not thread-safe) and doesn't do anything else (i.e. doesn't care about storing the contents of the files etc).

32

Moved the Func parameter to the end of the list.

The difference is semantics is different from what it seemed. Added comments and split function into two to clarify that.
We need Content in both cases, because if there's no ClangdUnit, then we'll have to create it and it requires initial text. And most actions need to call 'ClangdUnit::reparse' before they are run. But not code completion, which basically does reparse itself, hence the ReparseBeforeActionParameter.
See updated code for more details.

33

There are things like codeComplete that are actually mutable, but all other things could probably be run in parallel. (Don't have a great insight into that, though).
I really didn't want to change any of clangd behaviour in that commit, but I'd be glad to try adding that later.
What we could have in the future:
runOnUnitUnderReadLock(... [](const ClangdUnit &Unit) { auto diags = Unit.getLocalDiagnostics(); /* ... */} );
runOnUnitUnderWriteLock(... []( ClangdUnit &Unit) { auto completions = Unit.codeComplete(); /* ... */} );

clangd/GlobalCompilationDatabase.h

29

The fact that it's created only once(as opposed to tooling::CompilationDatabase, which can be recreated for each file) and handles all CompileCommand-related things(i.e. the idea is to add tracking of compile flags changes to this interface).
I called it ClangdCompilationDatabase first, but then I thought there's too much Clangd-prefixed names :-)
Any suggestions for a better name?

[krasimir]

clangd/CMakeLists.txt

12

nit: we might want to keep these sorted

clangd/ClangdMain.cpp

41

The combination of the name and the type of this variable makes no sense.

clangd/ClangdServer.h

41–43

I feel that shouldBuildDiagnosticsForFile doesn't belong to this interface.

49

If this models a request to the worker thread, why not call it WorkerRequest or something? The current name is confusing, because it implies that also other requests might be handled in this way.

62

Maybe explicitly mention Worker in the name of this class?

69

The constructor already gets a ClangdServer, why pass it here?

[ilya-biryukov]

Addressed style comments.
Moved ClangdScheduler to be the last field of the ClangdServer, so
that it's constructed last and destructed first.

[ilya-biryukov]

clangd/ClangdMain.cpp

41

Totally agree, renamed it.

clangd/ClangdServer.h

41–43

Sorry, this sneaked in from some code that didn't end up in this commit.
It's also not used anywhere now. I've removed it from this change.

62

I actually like the name ClangdScheduler. Albeit, it's a bit too generic and lives inside the clangd namespace, so I'm afraid of clashes later on, but until that happens I'd rather keep it.
Do you find ClangdScheduler confusing? Does the comment help?
But if you feel stongly that we shall change the name, let's do it.

69

Just to avoid introducing redundant fields.
(I.e. ClangdServer stores ClangdScheduler, which stores a reference to ClangdServer).

[klimek]

clangd/ClangdLSPServer.h

24

I think it's fine to have an intermediate step checked in; the important part is that we add comments explaining why this is the way it is, especially if we plan to change it in the future :) -> I'd add a class level comment explaining what this class does and why it exists.

clangd/ClangdUnit.h

44

Tiny nit: generally, use "." at the end of sentences :)

clangd/GlobalCompilationDatabase.h

29

Why don't we just use a tooling::CompilationDatabase and pass it around?

[ilya-biryukov]

Addressed new comments

clangd/ClangdLSPServer.h

24

Added comments to this class. And ClangdServer, it didn't have a comment too.

clangd/ClangdUnit.h

44

Fixed

clangd/GlobalCompilationDatabase.h

29

It has getAllFiles and getAllCompileCommands which we don't need for now.
It doesn't have capabilities to track changes to flags, which we'll want to add at some point. (Added a comment that we'll be adding those).

Also, currently implementations of CompilationDatabase actually require us to recreate them for each file, this logic is currently moved into DirectoryBasedGlobalCompilationDatabase for convenience.

[bkramer]

I believe this is good enough now.

clangd/ASTManager.h

-//===--- ASTManager.h - Clang AST manager -----------------------*- 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_ASTMANAGER_H
-#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_ASTMANAGER_H
-
-#include "DocumentStore.h"
-#include "JSONRPCDispatcher.h"
-#include "Protocol.h"
-#include "clang/Tooling/Core/Replacement.h"
-#include "llvm/ADT/IntrusiveRefCntPtr.h"
-#include <condition_variable>
-#include <deque>
-#include <thread>
-
-namespace clang {
-class ASTUnit;
-class DiagnosticsEngine;
-class PCHContainerOperations;
-namespace tooling {
-class CompilationDatabase;
-} // namespace tooling
-
-namespace clangd {
-
-/// Using 'unsigned' here to avoid undefined behaviour on overflow.
-typedef unsigned DocVersion;
-
-/// Stores ASTUnit and FixIts map for an opened document
-class DocData {
-public:
-  typedef std::map<clangd::Diagnostic, std::vector<clang::tooling::Replacement>>
-      DiagnosticToReplacementMap;
-
-public:
-  void setAST(std::unique_ptr<ASTUnit> AST);
-  ASTUnit *getAST() const;
-
-  void cacheFixIts(DiagnosticToReplacementMap FixIts);
-  std::vector<clang::tooling::Replacement>
-  getFixIts(const clangd::Diagnostic &D) const;
-
-private:
-  std::unique_ptr<ASTUnit> AST;
-  DiagnosticToReplacementMap FixIts;
-};
-
-enum class ASTManagerRequestType { ParseAndPublishDiagnostics, RemoveDocData };
-
-/// A request to the worker thread
-class ASTManagerRequest {
-public:
-  ASTManagerRequest() = default;
-  ASTManagerRequest(ASTManagerRequestType Type, std::string File,
-                    DocVersion Version);
-
-  ASTManagerRequestType Type;
-  std::string File;
-  DocVersion Version;
-};
-
-class ASTManager : public DocumentStoreListener {
-public:
-  ASTManager(JSONOutput &Output, DocumentStore &Store, bool RunSynchronously);
-  ~ASTManager() override;
-
-  void onDocumentAdd(StringRef File) override;
-  void onDocumentRemove(StringRef File) override;
-
-  /// Get code completions at a specified \p Line and \p Column in \p File.
-  ///
-  /// This function is thread-safe and returns completion items that own the
-  /// data they contain.
-  std::vector<CompletionItem> codeComplete(StringRef File, unsigned Line,
-                                           unsigned Column);
-
-  /// Get the fixes associated with a certain diagnostic in a specified file as
-  /// replacements.
-  ///
-  /// This function is thread-safe. It returns a copy to avoid handing out
-  /// references to unguarded data.
-  std::vector<clang::tooling::Replacement>
-  getFixIts(StringRef File, const clangd::Diagnostic &D);
-
-  DocumentStore &getStore() const { return Store; }
-
-private:
-  JSONOutput &Output;
-  DocumentStore &Store;
-
-  // Set to true if requests should block instead of being processed
-  // asynchronously.
-  bool RunSynchronously;
-
-  /// Loads a compilation database for File. May return nullptr if it fails. The
-  /// database is cached for subsequent accesses.
-  clang::tooling::CompilationDatabase *
-  getOrCreateCompilationDatabaseForFile(StringRef File);
-  // Creates a new ASTUnit for the document at File.
-  // FIXME: This calls chdir internally, which is thread unsafe.
-  std::unique_ptr<clang::ASTUnit>
-  createASTUnitForFile(StringRef File, const DocumentStore &Docs);
-
-  /// If RunSynchronously is false, queues the request to be run on the worker
-  /// thread.
-  /// If RunSynchronously is true, runs the request handler immediately on the
-  /// main thread.
-  void queueOrRun(ASTManagerRequestType RequestType, StringRef File);
-
-  void runWorker();
-  void handleRequest(ASTManagerRequestType RequestType, StringRef File);
-
-  /// Parses files and publishes diagnostics.
-  /// This function is called on the worker thread in asynchronous mode and
-  /// on the main thread in synchronous mode.
-  void parseFileAndPublishDiagnostics(StringRef File);
-
-  /// Caches compilation databases loaded from directories(keys are directories).
-  llvm::StringMap<std::unique_ptr<clang::tooling::CompilationDatabase>>
-      CompilationDatabases;
-
-  /// Clang objects.
-  /// A map from filenames to DocData structures that store ASTUnit and Fixits for
-  /// the files. The ASTUnits are used for generating diagnostics and fix-it-s
-  /// asynchronously by the worker thread and synchronously for code completion.
-  llvm::StringMap<DocData> DocDatas;
-  std::shared_ptr<clang::PCHContainerOperations> PCHs;
-  /// A lock for access to the DocDatas, CompilationDatabases and PCHs.
-  std::mutex ClangObjectLock;
-
-  /// Stores latest versions of the tracked documents to discard outdated requests.
-  /// Guarded by RequestLock.
-  /// TODO(ibiryukov): the entries are neved deleted from this map.
-  llvm::StringMap<DocVersion> DocVersions;
-
-  /// A LIFO queue of requests. Note that requests are discarded if the `version`
-  /// field is not equal to the one stored inside DocVersions.
-  /// TODO(krasimir): code completion should always have priority over parsing
-  /// for diagnostics.
-  std::deque<ASTManagerRequest> RequestQueue;
-  /// Setting Done to true will make the worker thread terminate.
-  bool Done = false;
-  /// Condition variable to wake up the worker thread.
-  std::condition_variable ClangRequestCV;
-  /// Lock for accesses to RequestQueue, DocVersions and Done.
-  std::mutex RequestLock;
-
-  /// We run parsing on a separate thread. This thread looks into RequestQueue to
-  /// find requests to handle and terminates when Done is set to true.
-  std::thread ClangWorker;
-};
-
-} // namespace clangd
-} // namespace clang
-
-#endif

clangd/CMakeLists.txt

 add_clang_executable(clangd
-  ASTManager.cpp
+  ClangdLSPServer.cpp
   ClangdMain.cpp
+  ClangdServer.cpp
+  ClangdUnit.cpp
+  ClangdUnitStore.cpp
+  DraftStore.cpp
+  GlobalCompilationDatabase.cpp
   JSONRPCDispatcher.cpp
   Protocol.cpp
   ProtocolHandlers.cpp
   )
 
 install(TARGETS clangd RUNTIME DESTINATION bin)
 
 target_link_libraries(clangd
   clangBasic
   clangFormat
   clangFrontend
   clangSema
   clangTooling
   clangToolingCore
   LLVMSupport
   )

clangd/ClangdLSPServer.h

+//===--- ClangdLSPServer.h - LSP server --------------------------*- 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_CLANGDLSPSERVER_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDLSPSERVER_H
+
+#include "ClangdServer.h"
+#include "Path.h"
+#include "Protocol.h"
+#include "clang/Tooling/Core/Replacement.h"
+
+namespace clang {
+namespace clangd {
+
+class JSONOutput;
+
+/// This class serves as an intermediate layer of LSP server implementation,
+/// glueing the JSON LSP protocol layer and ClangdServer together. It doesn't
+/// directly handle input from LSP client.
+/// Most methods are synchronous and return their result directly, but
+/// diagnostics are provided asynchronously when ready via
+/// JSONOutput::writeMessage.
+class ClangdLSPServer {
+public:
+  ClangdLSPServer(JSONOutput &Out, bool RunSynchronously);
+
+  /// Update the document text for \p File with \p Contents, schedule update of
+  /// diagnostics. Out.writeMessage will called to push diagnostics to LSP
+  /// client asynchronously when they are ready.
+  void openDocument(PathRef File, StringRef Contents);
+  /// Stop tracking the document for \p File.
+  void closeDocument(PathRef File);
+
+  /// Run code completion synchronously.
+  std::vector<CompletionItem> codeComplete(PathRef File, Position Pos);
+
+  /// Get the fixes associated with a certain diagnostic in a specified file as
+  /// replacements.
+  ///
+  /// This function is thread-safe. It returns a copy to avoid handing out
+  /// references to unguarded data.
+  std::vector<clang::tooling::Replacement>
+  getFixIts(StringRef File, const clangd::Diagnostic &D);
+
+  /// Get the current document contents stored for \p File.
+  /// FIXME(ibiryukov): This function is here to allow implementation of
+  /// formatCode from ProtocolHandlers.cpp. We should move formatCode to
+  /// ClangdServer class and remove this function from public interface.
+  std::string getDocument(PathRef File);
+
+private:
+  class LSPDiagnosticsConsumer;
+
+  /// Function that will be called on a separate thread when diagnostics are
+  /// ready. Sends the Dianostics to LSP client via Out.writeMessage and caches
+  /// corresponding fixits in the FixItsMap.
+  void consumeDiagnostics(PathRef File,
+                          std::vector<DiagWithFixIts> Diagnostics);
+
+  JSONOutput &Out;
+  ClangdServer Server;
+
+  std::mutex FixItsMutex;
+  typedef std::map<clangd::Diagnostic, std::vector<clang::tooling::Replacement>>
+      DiagnosticToReplacementMap;
+  /// Caches FixIts per file and diagnostics
+  llvm::StringMap<DiagnosticToReplacementMap> FixItsMap;
+};
+
+} // namespace clangd
+} // namespace clang
+
+#endif

clangd/ClangdLSPServer.cpp

+//===--- ClangdLSPServer.cpp - LSP server ------------------------*- C++-*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===---------------------------------------------------------------------===//
+
+#include "ClangdLSPServer.h"
+#include "JSONRPCDispatcher.h"
+
+using namespace clang::clangd;
+using namespace clang;
+
+class ClangdLSPServer::LSPDiagnosticsConsumer : public DiagnosticsConsumer {
+public:
+  LSPDiagnosticsConsumer(ClangdLSPServer &Server) : Server(Server) {}
+
+  virtual void onDiagnosticsReady(PathRef File,
+                                  std::vector<DiagWithFixIts> Diagnostics) {
+    Server.consumeDiagnostics(File, Diagnostics);
+  }
+
+private:
+  ClangdLSPServer &Server;
+};
+
+ClangdLSPServer::ClangdLSPServer(JSONOutput &Out, bool RunSynchronously)
+    : Out(Out),
+      Server(llvm::make_unique<DirectoryBasedGlobalCompilationDatabase>(),
+             llvm::make_unique<LSPDiagnosticsConsumer>(*this),
+             RunSynchronously) {}
+
+void ClangdLSPServer::openDocument(StringRef File, StringRef Contents) {
+  Server.addDocument(File, Contents);
+}
+
+void ClangdLSPServer::closeDocument(StringRef File) {
+  Server.removeDocument(File);
+}
+
+std::vector<CompletionItem> ClangdLSPServer::codeComplete(PathRef File,
+                                                          Position Pos) {
+  return Server.codeComplete(File, Pos);
+}
+
+std::vector<clang::tooling::Replacement>
+ClangdLSPServer::getFixIts(StringRef File, const clangd::Diagnostic &D) {
+  std::lock_guard<std::mutex> Lock(FixItsMutex);
+  auto DiagToFixItsIter = FixItsMap.find(File);
+  if (DiagToFixItsIter == FixItsMap.end())
+    return {};
+
+  const auto &DiagToFixItsMap = DiagToFixItsIter->second;
+  auto FixItsIter = DiagToFixItsMap.find(D);
+  if (FixItsIter == DiagToFixItsMap.end())
+    return {};
+
+  return FixItsIter->second;
+}
+
+std::string ClangdLSPServer::getDocument(PathRef File) {
+  return Server.getDocument(File);
+}
+
+void ClangdLSPServer::consumeDiagnostics(
+    PathRef File, std::vector<DiagWithFixIts> Diagnostics) {
+  std::string DiagnosticsJSON;
+
+  DiagnosticToReplacementMap LocalFixIts; // Temporary storage
+  for (auto &DiagWithFixes : Diagnostics) {
+    auto Diag = DiagWithFixes.Diag;
+    DiagnosticsJSON +=
+        R"({"range":)" + Range::unparse(Diag.range) +
+        R"(,"severity":)" + std::to_string(Diag.severity) +
+        R"(,"message":")" + llvm::yaml::escape(Diag.message) +
+        R"("},)";
+
+    // We convert to Replacements to become independent of the SourceManager.
+    auto &FixItsForDiagnostic = LocalFixIts[Diag];
+    std::copy(DiagWithFixes.FixIts.begin(), DiagWithFixes.FixIts.end(),
+              std::back_inserter(FixItsForDiagnostic));
+  }
+
+  // Cache FixIts
+  {
+    // FIXME(ibiryukov): should be deleted when documents are removed
+    std::lock_guard<std::mutex> Lock(FixItsMutex);
+    FixItsMap[File] = LocalFixIts;
+  }
+
+  // Publish diagnostics.
+  if (!DiagnosticsJSON.empty())
+    DiagnosticsJSON.pop_back(); // Drop trailing comma.
+  Out.writeMessage(
+      R"({"jsonrpc":"2.0","method":"textDocument/publishDiagnostics","params":{"uri":")" +
+      URI::fromFile(File).uri + R"(","diagnostics":[)" + DiagnosticsJSON +
+      R"(]}})");
+}

clangd/ClangdMain.cpp

 //===--- ClangdMain.cpp - clangd server loop ------------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
-#include "ASTManager.h"
-#include "DocumentStore.h"
 #include "JSONRPCDispatcher.h"
+#include "ClangdLSPServer.h"
+#include "Protocol.h"
 #include "ProtocolHandlers.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Program.h"
+
 #include <iostream>
+#include <memory>
 #include <string>
+
+using namespace clang;
 using namespace clang::clangd;
 
 static llvm::cl::opt<bool>
     RunSynchronously("run-synchronously",
                      llvm::cl::desc("parse on main thread"),
                      llvm::cl::init(false), llvm::cl::Hidden);
 
 int main(int argc, char *argv[]) {
   llvm::cl::ParseCommandLineOptions(argc, argv, "clangd");
   llvm::raw_ostream &Outs = llvm::outs();
   llvm::raw_ostream &Logs = llvm::errs();
   JSONOutput Out(Outs, Logs);
 
   // Change stdin to binary to not lose \r\n on windows.
   llvm::sys::ChangeStdinToBinary();
 
   // Set up a document store and intialize all the method handlers for JSONRPC
   // dispatching.
-  DocumentStore Store;
-  ASTManager AST(Out, Store, RunSynchronously);
-  Store.addListener(&AST);
+  ClangdLSPServer LSPServer(Out, RunSynchronously);
   JSONRPCDispatcher Dispatcher(llvm::make_unique<Handler>(Out));
   Dispatcher.registerHandler("initialize",
                              llvm::make_unique<InitializeHandler>(Out));
   auto ShutdownPtr = llvm::make_unique<ShutdownHandler>(Out);
   auto *ShutdownHandler = ShutdownPtr.get();
   Dispatcher.registerHandler("shutdown", std::move(ShutdownPtr));
   Dispatcher.registerHandler(
       "textDocument/didOpen",
-      llvm::make_unique<TextDocumentDidOpenHandler>(Out, Store));
+      llvm::make_unique<TextDocumentDidOpenHandler>(Out, LSPServer));
   Dispatcher.registerHandler(
       "textDocument/didClose",
-      llvm::make_unique<TextDocumentDidCloseHandler>(Out, Store));
+      llvm::make_unique<TextDocumentDidCloseHandler>(Out, LSPServer));
   Dispatcher.registerHandler(
       "textDocument/didChange",
-      llvm::make_unique<TextDocumentDidChangeHandler>(Out, Store));
+      llvm::make_unique<TextDocumentDidChangeHandler>(Out, LSPServer));
   Dispatcher.registerHandler(
       "textDocument/rangeFormatting",
-      llvm::make_unique<TextDocumentRangeFormattingHandler>(Out, Store));
+      llvm::make_unique<TextDocumentRangeFormattingHandler>(Out, LSPServer));
   Dispatcher.registerHandler(
       "textDocument/onTypeFormatting",
-      llvm::make_unique<TextDocumentOnTypeFormattingHandler>(Out, Store));
+      llvm::make_unique<TextDocumentOnTypeFormattingHandler>(Out, LSPServer));
   Dispatcher.registerHandler(
       "textDocument/formatting",
-      llvm::make_unique<TextDocumentFormattingHandler>(Out, Store));
+      llvm::make_unique<TextDocumentFormattingHandler>(Out, LSPServer));
   Dispatcher.registerHandler("textDocument/codeAction",
-                             llvm::make_unique<CodeActionHandler>(Out, AST));
+                             llvm::make_unique<CodeActionHandler>(Out, LSPServer));
   Dispatcher.registerHandler("textDocument/completion",
-                             llvm::make_unique<CompletionHandler>(Out, AST));
+                             llvm::make_unique<CompletionHandler>(Out, LSPServer));
 
   while (std::cin.good()) {
     // A Language Server Protocol message starts with a HTTP header, delimited
     // by \r\n.
     std::string Line;
     std::getline(std::cin, Line);
     if (!std::cin.good() && errno == EINTR) {
       std::cin.clear();

clangd/ClangdServer.h

-//===--- ASTManager.h - Clang AST manager -----------------------*- C++ -*-===//
+//===--- ClangdServer.h - Main clangd server code ----------------*- 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_ASTMANAGER_H
-#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_ASTMANAGER_H
+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDSERVER_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDSERVER_H
 
-#include "DocumentStore.h"
-#include "JSONRPCDispatcher.h"
-#include "Protocol.h"
+#include "ClangdUnitStore.h"
+#include "DraftStore.h"
+#include "GlobalCompilationDatabase.h"
+#include "clang/Frontend/ASTUnit.h"
+#include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/StringRef.h"
+
+#include "ClangdUnit.h"
+#include "Protocol.h"
+
 #include <condition_variable>
-#include <deque>
+#include <mutex>
+#include <string>
 #include <thread>
+#include <utility>
 
 namespace clang {
-class ASTUnit;
-class DiagnosticsEngine;
 class PCHContainerOperations;
-namespace tooling {
-class CompilationDatabase;
-} // namespace tooling
 
 namespace clangd {
 
-/// Using 'unsigned' here to avoid undefined behaviour on overflow.
-typedef unsigned DocVersion;
-
-/// Stores ASTUnit and FixIts map for an opened document
-class DocData {
-public:
-  typedef std::map<clangd::Diagnostic, std::vector<clang::tooling::Replacement>>
-      DiagnosticToReplacementMap;
-
+class DiagnosticsConsumer {
 public:
-  void setAST(std::unique_ptr<ASTUnit> AST);
-  ASTUnit *getAST() const;
-
-  void cacheFixIts(DiagnosticToReplacementMap FixIts);
-  std::vector<clang::tooling::Replacement>
-  getFixIts(const clangd::Diagnostic &D) const;
+  virtual ~DiagnosticsConsumer() = default;
 
-private:
-  std::unique_ptr<ASTUnit> AST;
-  DiagnosticToReplacementMap FixIts;
+  /// Called by ClangdServer when \p Diagnostics for \p File are ready.
+  virtual void onDiagnosticsReady(PathRef File,
+                                  std::vector<DiagWithFixIts> Diagnostics) = 0;
 };
 
-enum class ASTManagerRequestType { ParseAndPublishDiagnostics, RemoveDocData };
+enum class WorkerRequestKind { ParseAndPublishDiagnostics, RemoveDocData };
 
 /// A request to the worker thread
-class ASTManagerRequest {
+class WorkerRequest {
 public:
-  ASTManagerRequest() = default;
-  ASTManagerRequest(ASTManagerRequestType Type, std::string File,
-                    DocVersion Version);
+  WorkerRequest() = default;
+  WorkerRequest(WorkerRequestKind Kind, Path File, DocVersion Version);
 
-  ASTManagerRequestType Type;
-  std::string File;
+  WorkerRequestKind Kind;
+  Path File;
   DocVersion Version;
 };
 
-class ASTManager : public DocumentStoreListener {
-public:
-  ASTManager(JSONOutput &Output, DocumentStore &Store, bool RunSynchronously);
-  ~ASTManager() override;
-
-  void onDocumentAdd(StringRef File) override;
-  void onDocumentRemove(StringRef File) override;
+class ClangdServer;
 
-  /// Get code completions at a specified \p Line and \p Column in \p File.
-  ///
-  /// This function is thread-safe and returns completion items that own the
-  /// data they contain.
-  std::vector<CompletionItem> codeComplete(StringRef File, unsigned Line,
-                                           unsigned Column);
-
-  /// Get the fixes associated with a certain diagnostic in a specified file as
-  /// replacements.
-  ///
-  /// This function is thread-safe. It returns a copy to avoid handing out
-  /// references to unguarded data.
-  std::vector<clang::tooling::Replacement>
-  getFixIts(StringRef File, const clangd::Diagnostic &D);
+/// Handles running WorkerRequests of ClangdServer on a separate threads.
+/// Currently runs only one worker thread.
+class ClangdScheduler {
+public:
+  ClangdScheduler(ClangdServer &Server, bool RunSynchronously);
+  ~ClangdScheduler();
 
-  DocumentStore &getStore() const { return Store; }
+  /// Enqueue WorkerRequest to be run on a worker thread
+  void enqueue(ClangdServer &Server, WorkerRequest Request);
 
 private:
-  JSONOutput &Output;
-  DocumentStore &Store;
-
-  // Set to true if requests should block instead of being processed
-  // asynchronously.
   bool RunSynchronously;
-
-  /// Loads a compilation database for File. May return nullptr if it fails. The
-  /// database is cached for subsequent accesses.
-  clang::tooling::CompilationDatabase *
-  getOrCreateCompilationDatabaseForFile(StringRef File);
-  // Creates a new ASTUnit for the document at File.
-  // FIXME: This calls chdir internally, which is thread unsafe.
-  std::unique_ptr<clang::ASTUnit>
-  createASTUnitForFile(StringRef File, const DocumentStore &Docs);
-
-  /// If RunSynchronously is false, queues the request to be run on the worker
-  /// thread.
-  /// If RunSynchronously is true, runs the request handler immediately on the
-  /// main thread.
-  void queueOrRun(ASTManagerRequestType RequestType, StringRef File);
-
-  void runWorker();
-  void handleRequest(ASTManagerRequestType RequestType, StringRef File);
-
-  /// Parses files and publishes diagnostics.
-  /// This function is called on the worker thread in asynchronous mode and
-  /// on the main thread in synchronous mode.
-  void parseFileAndPublishDiagnostics(StringRef File);
-
-  /// Caches compilation databases loaded from directories(keys are directories).
-  llvm::StringMap<std::unique_ptr<clang::tooling::CompilationDatabase>>
-      CompilationDatabases;
-
-  /// Clang objects.
-  /// A map from filenames to DocData structures that store ASTUnit and Fixits for
-  /// the files. The ASTUnits are used for generating diagnostics and fix-it-s
-  /// asynchronously by the worker thread and synchronously for code completion.
-  llvm::StringMap<DocData> DocDatas;
-  std::shared_ptr<clang::PCHContainerOperations> PCHs;
-  /// A lock for access to the DocDatas, CompilationDatabases and PCHs.
-  std::mutex ClangObjectLock;
-
-  /// Stores latest versions of the tracked documents to discard outdated requests.
-  /// Guarded by RequestLock.
-  /// TODO(ibiryukov): the entries are neved deleted from this map.
-  llvm::StringMap<DocVersion> DocVersions;
-
-  /// A LIFO queue of requests. Note that requests are discarded if the `version`
-  /// field is not equal to the one stored inside DocVersions.
-  /// TODO(krasimir): code completion should always have priority over parsing
-  /// for diagnostics.
-  std::deque<ASTManagerRequest> RequestQueue;
+  std::mutex Mutex;
+  /// We run some tasks on a separate thread(parsing, ClangdUnit cleanup).
+  /// This thread looks into RequestQueue to find requests to handle and
+  /// terminates when Done is set to true.
+  std::thread Worker;
   /// Setting Done to true will make the worker thread terminate.
   bool Done = false;
+  /// A LIFO queue of requests. Note that requests are discarded if the
+  /// `version` field is not equal to the one stored inside DraftStore.
+  /// FIXME(krasimir): code completion should always have priority over parsing
+  /// for diagnostics.
+  std::deque<WorkerRequest> RequestQueue;
   /// Condition variable to wake up the worker thread.
-  std::condition_variable ClangRequestCV;
-  /// Lock for accesses to RequestQueue, DocVersions and Done.
-  std::mutex RequestLock;
-
-  /// We run parsing on a separate thread. This thread looks into RequestQueue to
-  /// find requests to handle and terminates when Done is set to true.
-  std::thread ClangWorker;
+  std::condition_variable RequestCV;
+};
+
+/// Provides API to manage ASTs for a collection of C++ files and request
+/// various language features(currently, only codeCompletion and async
+/// diagnostics for tracked files).
+class ClangdServer {
+public:
+  ClangdServer(std::unique_ptr<GlobalCompilationDatabase> CDB,
+               std::unique_ptr<DiagnosticsConsumer> DiagConsumer,
+               bool RunSynchronously);
+
+  /// Add a \p File to the list of tracked C++ files or update the contents if
+  /// \p File is already tracked. Also schedules parsing of the AST for it on a
+  /// separate thread. When the parsing is complete, DiagConsumer passed in
+  /// constructor will receive onDiagnosticsReady callback.
+  void addDocument(PathRef File, StringRef Contents);
+
+  /// Remove \p File from list of tracked files, schedule a request to free
+  /// resources associated with it.
+  void removeDocument(PathRef File);
+
+  /// Run code completion for \p File at \p Pos.
+  std::vector<CompletionItem> codeComplete(PathRef File, Position Pos);
+
+  /// Gets current document contents for \p File. \p File must point to a
+  /// currently tracked file.
+  /// FIXME(ibiryukov): This function is here to allow implementation of
+  /// formatCode from ProtocolHandlers.cpp. We should move formatCode to this
+  /// class and remove this function from public interface.
+  std::string getDocument(PathRef File);
+
+private:
+  friend class ClangdScheduler;
+
+  /// This function is called on a worker thread.
+  void handleRequest(WorkerRequest Request);
+
+  std::unique_ptr<GlobalCompilationDatabase> CDB;
+  std::unique_ptr<DiagnosticsConsumer> DiagConsumer;
+  DraftStore DraftMgr;
+  ClangdUnitStore Units;
+  std::shared_ptr<PCHContainerOperations> PCHs;
+  // WorkScheduler has to be the last member, because its destructor has to be
+  // called before all other members to stop the worker thread that references
+  // ClangdServer
+  ClangdScheduler WorkScheduler;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/ClangdServer.cpp

-//===--- ASTManager.cpp - Clang AST manager -------------------------------===//
+//===--- ClangdServer.cpp - Main clangd server code --------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
-//===----------------------------------------------------------------------===//
+//===-------------------------------------------------------------------===//
 
-#include "ASTManager.h"
-#include "JSONRPCDispatcher.h"
-#include "Protocol.h"
+#include "ClangdServer.h"
 #include "clang/Frontend/ASTUnit.h"
 #include "clang/Frontend/CompilerInstance.h"
+#include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Tooling/CompilationDatabase.h"
-#include "llvm/Support/Format.h"
-#include "llvm/Support/Path.h"
-#include <mutex>
-#include <thread>
-using namespace clang;
-using namespace clangd;
+#include "llvm/Support/FileSystem.h"
 
-void DocData::setAST(std::unique_ptr<ASTUnit> AST) {
-  this->AST = std::move(AST);
-}
-
-ASTUnit *DocData::getAST() const { return AST.get(); }
-
-void DocData::cacheFixIts(DiagnosticToReplacementMap FixIts) {
-  this->FixIts = std::move(FixIts);
-}
-
-std::vector<clang::tooling::Replacement>
-DocData::getFixIts(const clangd::Diagnostic &D) const {
-  auto it = FixIts.find(D);
-  if (it != FixIts.end())
-    return it->second;
-  return {};
-}
+using namespace clang::clangd;
 
-ASTManagerRequest::ASTManagerRequest(ASTManagerRequestType Type,
-                                     std::string File,
-                                     DocVersion Version)
-    : Type(Type), File(File), Version(Version) {}
+WorkerRequest::WorkerRequest(WorkerRequestKind Kind, Path File,
+    : Kind(Kind), File(File), Version(Version) {}
 
-/// Retrieve a copy of the contents of every file in the store, for feeding into
-/// ASTUnit.
-static std::vector<ASTUnit::RemappedFile>
-getRemappedFiles(const DocumentStore &Docs) {
-  // FIXME: Use VFS instead. This would allow us to get rid of the chdir below.
-  std::vector<ASTUnit::RemappedFile> RemappedFiles;
-  for (const auto &P : Docs.getAllDocuments()) {
-    StringRef FileName = P.first;
-    RemappedFiles.push_back(ASTUnit::RemappedFile(
-        FileName,
-        llvm::MemoryBuffer::getMemBufferCopy(P.second, FileName).release()));
-  }
-  return RemappedFiles;
-}
-
-/// 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!");
-}
-
-static CompletionItemKind getKind(CXCursorKind K) {
-  switch (K) {
-  case CXCursor_MacroInstantiation:
-  case CXCursor_MacroDefinition:
-    return CompletionItemKind::Text;
-  case CXCursor_CXXMethod:
-    return CompletionItemKind::Method;
-  case CXCursor_FunctionDecl:
-  case CXCursor_FunctionTemplate:
-    return CompletionItemKind::Function;
-  case CXCursor_Constructor:
-  case CXCursor_Destructor:
-    return CompletionItemKind::Constructor;
-  case CXCursor_FieldDecl:
-    return CompletionItemKind::Field;
-  case CXCursor_VarDecl:
-  case CXCursor_ParmDecl:
-    return CompletionItemKind::Variable;
-  case CXCursor_ClassDecl:
-  case CXCursor_StructDecl:
-  case CXCursor_UnionDecl:
-  case CXCursor_ClassTemplate:
-  case CXCursor_ClassTemplatePartialSpecialization:
-    return CompletionItemKind::Class;
-  case CXCursor_Namespace:
-  case CXCursor_NamespaceAlias:
-  case CXCursor_NamespaceRef:
-    return CompletionItemKind::Module;
-  case CXCursor_EnumConstantDecl:
-    return CompletionItemKind::Value;
-  case CXCursor_EnumDecl:
-    return CompletionItemKind::Enum;
-  case CXCursor_TypeAliasDecl:
-  case CXCursor_TypeAliasTemplateDecl:
-  case CXCursor_TypedefDecl:
-  case CXCursor_MemberRef:
-  case CXCursor_TypeRef:
-    return CompletionItemKind::Reference;
-  default:
-    return CompletionItemKind::Missing;
-  }
+ClangdScheduler::ClangdScheduler(ClangdServer &Server, bool RunSynchronously)
+    : RunSynchronously(RunSynchronously) {
+  if (RunSynchronously) {
+    // Don't start the worker thread if we're running synchronously
+    return;
   }
 
-ASTManager::ASTManager(JSONOutput &Output, DocumentStore &Store,
-                       bool RunSynchronously)
-    : Output(Output), Store(Store), RunSynchronously(RunSynchronously),
-      PCHs(std::make_shared<PCHContainerOperations>()),
-      ClangWorker([this]() { runWorker(); }) {}
-
-void ASTManager::runWorker() {
+  // Initialize Worker in ctor body, rather than init list to avoid potentially
+  // using not-yet-initialized members
+  Worker = std::thread([&Server, this]() {
     while (true) {
-    ASTManagerRequest Request;
+      WorkerRequest Request;
 
       // Pick request from the queue
       {
-      std::unique_lock<std::mutex> Lock(RequestLock);
+        std::unique_lock<std::mutex> Lock(Mutex);
         // Wait for more requests.
-      ClangRequestCV.wait(Lock,
-                          [this] { return !RequestQueue.empty() || Done; });
+        RequestCV.wait(Lock, [this] { return !RequestQueue.empty() || Done; });
         if (Done)
           return;
+
         assert(!RequestQueue.empty() && "RequestQueue was empty");
 
         Request = std::move(RequestQueue.back());
         RequestQueue.pop_back();
 
         // Skip outdated requests
-      if (Request.Version != DocVersions.find(Request.File)->second) {
-        Output.log("Version for " + Twine(Request.File) +
-                   " in request is outdated, skipping request\n");
+        if (Request.Version != Server.DraftMgr.getVersion(Request.File)) {
+          // FIXME(ibiryukov): Logging
+          // Output.log("Version for " + Twine(Request.File) +
+          //            " in request is outdated, skipping request\n");
           continue;
         }
-    } // unlock RequestLock
-
-    handleRequest(Request.Type, Request.File);
-  }
-}
+      } // unlock Mutex
 
-void ASTManager::queueOrRun(ASTManagerRequestType RequestType, StringRef File) {
-  if (RunSynchronously) {
-    handleRequest(RequestType, File);
-    return;
+      Server.handleRequest(std::move(Request));
     }
-
-  std::lock_guard<std::mutex> Guard(RequestLock);
-  // We increment the version of the added document immediately and schedule
-  // the requested operation to be run on a worker thread
-  DocVersion version = ++DocVersions[File];
-  RequestQueue.push_back(ASTManagerRequest(RequestType, File, version));
-  ClangRequestCV.notify_one();
+  });
 }
 
-void ASTManager::handleRequest(ASTManagerRequestType RequestType,
-                               StringRef File) {
-  switch (RequestType) {
-  case ASTManagerRequestType::ParseAndPublishDiagnostics:
-    parseFileAndPublishDiagnostics(File);
-    break;
-  case ASTManagerRequestType::RemoveDocData: {
-    std::lock_guard<std::mutex> Lock(ClangObjectLock);
-    auto DocDataIt = DocDatas.find(File);
-    // We could get the remove request before parsing for the document is
-    // started, just do nothing in that case, parsing request will be discarded
-    // because it has a lower version value
-    if (DocDataIt == DocDatas.end())
-      return;
-    DocDatas.erase(DocDataIt);
-    break;
-  } // unlock ClangObjectLock
-  }
-}
+ClangdScheduler::~ClangdScheduler() {
+  if (RunSynchronously)
+    return; // no worker thread is running in that case
 
-void ASTManager::parseFileAndPublishDiagnostics(StringRef File) {
-  std::unique_lock<std::mutex> ClangObjectLockGuard(ClangObjectLock);
-
-  auto &DocData = DocDatas[File];
-  ASTUnit *Unit = DocData.getAST();
-  if (!Unit) {
-    auto newAST = createASTUnitForFile(File, this->Store);
-    Unit = newAST.get();
-
-    DocData.setAST(std::move(newAST));
-  } else {
-    // Do a reparse if this wasn't the first parse.
-    // FIXME: This might have the wrong working directory if it changed in the
-    // meantime.
-    Unit->Reparse(PCHs, getRemappedFiles(this->Store));
-  }
-
-  if (!Unit)
-    return;
-
-  // Send the diagnotics to the editor.
-  // FIXME: If the diagnostic comes from a different file, do we want to
-  // show them all? Right now we drop everything not coming from the
-  // main file.
-  std::string Diagnostics;
-  DocData::DiagnosticToReplacementMap LocalFixIts; // Temporary storage
-  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};
-    Diagnostics +=
-        R"({"range":)" + Range::unparse(R) +
-        R"(,"severity":)" + std::to_string(getSeverity(D->getLevel())) +
-        R"(,"message":")" + llvm::yaml::escape(D->getMessage()) +
-        R"("},)";
-
-    // We convert to Replacements to become independent of the SourceManager.
-    clangd::Diagnostic Diag = {R, getSeverity(D->getLevel()), D->getMessage()};
-    auto &FixItsForDiagnostic = LocalFixIts[Diag];
-    for (const FixItHint &Fix : D->getFixIts()) {
-      FixItsForDiagnostic.push_back(clang::tooling::Replacement(
-          Unit->getSourceManager(), Fix.RemoveRange, Fix.CodeToInsert));
-    }
-  }
-
-  // Put FixIts into place.
-  DocData.cacheFixIts(std::move(LocalFixIts));
-
-  ClangObjectLockGuard.unlock();
-  // No accesses to clang objects are allowed after this point.
-
-  // Publish diagnostics.
-  if (!Diagnostics.empty())
-    Diagnostics.pop_back(); // Drop trailing comma.
-  Output.writeMessage(
-      R"({"jsonrpc":"2.0","method":"textDocument/publishDiagnostics","params":{"uri":")" +
-      URI::fromFile(File).uri + R"(","diagnostics":[)" + Diagnostics + R"(]}})");
-}
-
-ASTManager::~ASTManager() {
   {
-    std::lock_guard<std::mutex> Guard(RequestLock);
-    // Wake up the clang worker thread, then exit.
+    std::lock_guard<std::mutex> Lock(Mutex);
+    // Wake up the worker thread
     Done = true;
-    ClangRequestCV.notify_one();
-  } // unlock DocDataLock
-  ClangWorker.join();
+    RequestCV.notify_one();
+  } // unlock Mutex
+  Worker.join();
 }
 
-void ASTManager::onDocumentAdd(StringRef File) {
-  queueOrRun(ASTManagerRequestType::ParseAndPublishDiagnostics, File);
+void ClangdScheduler::enqueue(ClangdServer &Server, WorkerRequest Request) {
+  if (RunSynchronously) {
+    Server.handleRequest(Request);
+    return;
   }
 
-void ASTManager::onDocumentRemove(StringRef File) {
-  queueOrRun(ASTManagerRequestType::RemoveDocData, File);
+  std::lock_guard<std::mutex> Lock(Mutex);
+  RequestQueue.push_back(Request);
+  RequestCV.notify_one();
 }
 
-tooling::CompilationDatabase *
-ASTManager::getOrCreateCompilationDatabaseForFile(StringRef File) {
-  namespace path = llvm::sys::path;
-
-  assert((path::is_absolute(File, path::Style::posix) ||
-          path::is_absolute(File, path::Style::windows)) &&
-         "path must be absolute");
-
-  for (auto Path = path::parent_path(File); !Path.empty();
-       Path = path::parent_path(Path)) {
+ClangdServer::ClangdServer(std::unique_ptr<GlobalCompilationDatabase> CDB,
+                           std::unique_ptr<DiagnosticsConsumer> DiagConsumer,
+                           bool RunSynchronously)
+    : CDB(std::move(CDB)), DiagConsumer(std::move(DiagConsumer)),
+      PCHs(std::make_shared<PCHContainerOperations>()),
+      WorkScheduler(*this, RunSynchronously) {}
 
-    auto CachedIt = CompilationDatabases.find(Path);
-    if (CachedIt != CompilationDatabases.end())
-      return CachedIt->second.get();
-    std::string Error;
-    auto CDB = tooling::CompilationDatabase::loadFromDirectory(Path, Error);
-    if (!CDB) {
-      if (!Error.empty()) {
-        Output.log("Error when trying to load compilation database from " +
-                   Twine(Path) + ": " + Twine(Error) + "\n");
-      }
-      continue;
-    }
+void ClangdServer::addDocument(PathRef File, StringRef Contents) {
+  DocVersion NewVersion = DraftMgr.updateDraft(File, Contents);
+  WorkScheduler.enqueue(
+      *this, WorkerRequest(WorkerRequestKind::ParseAndPublishDiagnostics, File,
+                           NewVersion));
+}
+
+void ClangdServer::removeDocument(PathRef File) {
+  auto NewVersion = DraftMgr.removeDraft(File);
+  WorkScheduler.enqueue(
+      *this, WorkerRequest(WorkerRequestKind::RemoveDocData, File, NewVersion));
+}
+
+std::vector<CompletionItem> ClangdServer::codeComplete(PathRef File,
+                                                       Position Pos) {
+  auto FileContents = DraftMgr.getDraft(File);
+  assert(FileContents.Draft && "codeComplete is called for non-added document");
+
+  std::vector<CompletionItem> Result;
+  Units.runOnUnitWithoutReparse(
+      File, *FileContents.Draft, *CDB, PCHs, [&](ClangdUnit &Unit) {
+        Result = Unit.codeComplete(*FileContents.Draft, Pos);
+      });
+  return Result;
+}
+
+std::string ClangdServer::getDocument(PathRef File) {
+  auto draft = DraftMgr.getDraft(File);
+  assert(draft.Draft && "File is not tracked, cannot get contents");
+  return *draft.Draft;
+}
+
+void ClangdServer::handleRequest(WorkerRequest Request) {
+  switch (Request.Kind) {
+  case WorkerRequestKind::ParseAndPublishDiagnostics: {
+    auto FileContents = DraftMgr.getDraft(Request.File);
+    if (FileContents.Version != Request.Version)
+      return; // This request is outdated, do nothing
+
+    assert(FileContents.Draft &&
+           "No contents inside a file that was scheduled for reparse");
+    Units.runOnUnit(Request.File, *FileContents.Draft, *CDB, PCHs,
+                    [&](ClangdUnit const &Unit) {
+                      DiagConsumer->onDiagnosticsReady(
+                          Request.File, Unit.getLocalDiagnostics());
+                    });
+    break;
+  case WorkerRequestKind::RemoveDocData:
+    if (Request.Version != DraftMgr.getVersion(Request.File))
+      return; // This request is outdated, do nothing
 
-    // TODO(ibiryukov): Invalidate cached compilation databases on changes
-    auto result = CDB.get();
-    CompilationDatabases.insert(std::make_pair(Path, std::move(CDB)));
-    return result;
-  }
-
-  Output.log("Failed to find compilation database for " + Twine(File) + "\n");
-  return nullptr;
-}
-
-std::unique_ptr<clang::ASTUnit>
-ASTManager::createASTUnitForFile(StringRef File, const DocumentStore &Docs) {
-  tooling::CompilationDatabase *CDB =
-      getOrCreateCompilationDatabaseForFile(File);
-
-  std::vector<tooling::CompileCommand> Commands;
-
-  if (CDB) {
-    Commands = CDB->getCompileCommands(File);
-    // chdir. This is thread hostile.
-    if (!Commands.empty())
-      llvm::sys::fs::set_current_path(Commands.front().Directory);
-  }
-  if (Commands.empty()) {
-    // Add a fake command line if we know nothing.
-    Commands.push_back(tooling::CompileCommand(
-        llvm::sys::path::parent_path(File), llvm::sys::path::filename(File),
-        {"clang", "-fsyntax-only", File.str()}, ""));
-  }
-
-  // Inject the resource dir.
-  // FIXME: Don't overwrite it if it's already there.
-  static int Dummy; // Just an address in this process.
-  std::string ResourceDir =
-      CompilerInvocation::GetResourcesPath("clangd", (void *)&Dummy);
-  Commands.front().CommandLine.push_back("-resource-dir=" + 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());
-
-  auto ArgP = &*ArgStrs.begin();
-  return std::unique_ptr<clang::ASTUnit>(ASTUnit::LoadFromCommandLine(
-      ArgP, ArgP + ArgStrs.size(), PCHs, Diags, ResourceDir,
-      /*OnlyLocalDecls=*/false, /*CaptureDiagnostics=*/true,
-      getRemappedFiles(Docs),
-      /*RemappedFilesKeepOriginalName=*/true,
-      /*PrecompilePreambleAfterNParses=*/1, /*TUKind=*/TU_Complete,
-      /*CacheCodeCompletionResults=*/true,
-      /*IncludeBriefCommentsInCodeCompletion=*/true,
-      /*AllowPCHWithCompilerErrors=*/true));
-}
-
-std::vector<clang::tooling::Replacement>
-ASTManager::getFixIts(StringRef File, const clangd::Diagnostic &D) {
-  // TODO(ibiryukov): the FixIts should be available immediately
-  // even when parsing is being run on a worker thread
-  std::lock_guard<std::mutex> Guard(ClangObjectLock);
-  return DocDatas[File].getFixIts(D);
-}
-
-namespace {
-class CompletionItemsCollector : public CodeCompleteConsumer {
-  std::vector<CompletionItem> *Items;
-  std::shared_ptr<clang::GlobalCodeCompletionAllocator> Allocator;
-  CodeCompletionTUInfo CCTUInfo;
-
-public:
-  CompletionItemsCollector(std::vector<CompletionItem> *Items,
-                           const CodeCompleteOptions &CodeCompleteOpts)
-      : CodeCompleteConsumer(CodeCompleteOpts, /*OutputIsBinary=*/false),
-        Items(Items),
-        Allocator(std::make_shared<clang::GlobalCodeCompletionAllocator>()),
-        CCTUInfo(Allocator) {}
-
-  void ProcessCodeCompleteResults(Sema &S, CodeCompletionContext Context,
-                                  CodeCompletionResult *Results,
-                                  unsigned NumResults) override {
-    for (unsigned I = 0; I != NumResults; ++I) {
-      CodeCompletionResult &Result = Results[I];
-      CodeCompletionString *CCS = Result.CreateCodeCompletionString(
-          S, Context, *Allocator, CCTUInfo,
-          CodeCompleteOpts.IncludeBriefComments);
-      if (CCS) {
-        CompletionItem Item;
-        assert(CCS->getTypedText());
-        Item.label = CCS->getTypedText();
-        Item.kind = getKind(Result.CursorKind);
-        if (CCS->getBriefComment())
-          Item.documentation = CCS->getBriefComment();
-        Items->push_back(std::move(Item));
-      }
-    }
-  }
-
-  GlobalCodeCompletionAllocator &getAllocator() override { return *Allocator; }
-
-  CodeCompletionTUInfo &getCodeCompletionTUInfo() override { return CCTUInfo; }
-};
-
-} // namespace
-
-std::vector<CompletionItem>
-ASTManager::codeComplete(StringRef File, unsigned Line, unsigned Column) {
-  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<CompletionItem> Items;
-  CompletionItemsCollector Collector(&Items, CCO);
-
-  std::lock_guard<std::mutex> Guard(ClangObjectLock);
-  auto &DocData = DocDatas[File];
-  auto Unit = DocData.getAST();
-  if (!Unit) {
-    auto newAST = createASTUnitForFile(File, this->Store);
-    Unit = newAST.get();
-    DocData.setAST(std::move(newAST));
-  }
-  if (!Unit)
-    return {};
-  IntrusiveRefCntPtr<SourceManager> SourceMgr(
-      new SourceManager(*DiagEngine, Unit->getFileManager()));
-  // 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(File, Line + 1, Column + 1, getRemappedFiles(this->Store),
-                     CCO.IncludeMacros, CCO.IncludeCodePatterns,
-                     CCO.IncludeBriefComments, Collector, PCHs, *DiagEngine,
-                     LangOpts, *SourceMgr, Unit->getFileManager(),
-                     StoredDiagnostics, OwnedBuffers);
-  for (const llvm::MemoryBuffer *Buffer : OwnedBuffers)
-    delete Buffer;
-  return Items;
+    Units.removeUnitIfPresent(Request.File);
+    break;
+  }
 }

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 "Protocol.h"
+#include "Path.h"
+#include "clang/Frontend/ASTUnit.h"
+#include "clang/Tooling/Core/Replacement.h"
+#include <memory>
+
+namespace clang {
+class ASTUnit;
+class PCHContainerOperations;
+
+namespace tooling {
+struct CompileCommand;
+}
+
+namespace clangd {
+
+/// A diagnostic with its FixIts.
+struct DiagWithFixIts {
+  clangd::Diagnostic Diag;
+  llvm::SmallVector<tooling::Replacement, 1> FixIts;
+};
+
+/// Stores parsed C++ AST and provides implementations of all operations clangd
+/// would want to perform on parsed C++ files.
+class ClangdUnit {
+public:
+  ClangdUnit(PathRef FileName, StringRef Contents,
+             std::shared_ptr<PCHContainerOperations> PCHs,
+             std::vector<tooling::CompileCommand> Commands);
+
+  /// Reparse with new contents.
+  void reparse(StringRef Contents);
+
+  /// Get code completions at a specified \p Line and \p Column in \p File.
+  ///
+  /// This function is thread-safe and returns completion items that own the
+  /// data they contain.
+  std::vector<CompletionItem> codeComplete(StringRef Contents, Position Pos);
+  /// Returns diagnostics and corresponding FixIts for each diagnostic that are
+  /// located in the current file.
+  std::vector<DiagWithFixIts> getLocalDiagnostics() const;
+
+private:
+  Path FileName;
+  std::unique_ptr<ASTUnit> 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/Tooling/CompilationDatabase.h"
+
+using namespace clang::clangd;
+using namespace clang;
+
+ClangdUnit::ClangdUnit(PathRef FileName, StringRef Contents,
+                       std::shared_ptr<PCHContainerOperations> PCHs,
+                       std::vector<tooling::CompileCommand> Commands)
+    : 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.
+  static int Dummy; // Just an address in this process.
+  std::string ResourceDir =
+      CompilerInvocation::GetResourcesPath("clangd", (void *)&Dummy);
+  Commands.front().CommandLine.push_back("-resource-dir=" + 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_Complete,
+      /*CacheCodeCompletionResults=*/true,
+      /*IncludeBriefCommentsInCodeCompletion=*/true,
+      /*AllowPCHWithCompilerErrors=*/true));
+}
+
+void ClangdUnit::reparse(StringRef Contents) {
+  // 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());
+
+  Unit->Reparse(PCHs, RemappedSource);
+}
+
+namespace {
+
+CompletionItemKind getKind(CXCursorKind K) {
+  switch (K) {
+  case CXCursor_MacroInstantiation:
+  case CXCursor_MacroDefinition:
+    return CompletionItemKind::Text;
+  case CXCursor_CXXMethod:
+    return CompletionItemKind::Method;
+  case CXCursor_FunctionDecl:
+  case CXCursor_FunctionTemplate:
+    return CompletionItemKind::Function;
+  case CXCursor_Constructor:
+  case CXCursor_Destructor:
+    return CompletionItemKind::Constructor;
+  case CXCursor_FieldDecl:
+    return CompletionItemKind::Field;
+  case CXCursor_VarDecl:
+  case CXCursor_ParmDecl:
+    return CompletionItemKind::Variable;
+  case CXCursor_ClassDecl:
+  case CXCursor_StructDecl:
+  case CXCursor_UnionDecl:
+  case CXCursor_ClassTemplate:
+  case CXCursor_ClassTemplatePartialSpecialization:
+    return CompletionItemKind::Class;
+  case CXCursor_Namespace:
+  case CXCursor_NamespaceAlias:
+  case CXCursor_NamespaceRef:
+    return CompletionItemKind::Module;
+  case CXCursor_EnumConstantDecl:
+    return CompletionItemKind::Value;
+  case CXCursor_EnumDecl:
+    return CompletionItemKind::Enum;
+  case CXCursor_TypeAliasDecl:
+  case CXCursor_TypeAliasTemplateDecl:
+  case CXCursor_TypedefDecl:
+  case CXCursor_MemberRef:
+  case CXCursor_TypeRef:
+    return CompletionItemKind::Reference;
+  default:
+    return CompletionItemKind::Missing;
+  }
+}
+
+class CompletionItemsCollector : public CodeCompleteConsumer {
+  std::vector<CompletionItem> *Items;
+  std::shared_ptr<clang::GlobalCodeCompletionAllocator> Allocator;
+  CodeCompletionTUInfo CCTUInfo;
+
+public:
+  CompletionItemsCollector(std::vector<CompletionItem> *Items,
+                           const CodeCompleteOptions &CodeCompleteOpts)
+      : CodeCompleteConsumer(CodeCompleteOpts, /*OutputIsBinary=*/false),
+        Items(Items),
+        Allocator(std::make_shared<clang::GlobalCodeCompletionAllocator>()),
+        CCTUInfo(Allocator) {}
+
+  void ProcessCodeCompleteResults(Sema &S, CodeCompletionContext Context,
+                                  CodeCompletionResult *Results,
+                                  unsigned NumResults) override {
+    for (unsigned I = 0; I != NumResults; ++I) {
+      CodeCompletionResult &Result = Results[I];
+      CodeCompletionString *CCS = Result.CreateCodeCompletionString(
+          S, Context, *Allocator, CCTUInfo,
+          CodeCompleteOpts.IncludeBriefComments);
+      if (CCS) {
+        CompletionItem Item;
+        assert(CCS->getTypedText());
+        Item.label = CCS->getTypedText();
+        Item.kind = getKind(Result.CursorKind);
+        if (CCS->getBriefComment())
+          Item.documentation = CCS->getBriefComment();
+        Items->push_back(std::move(Item));
+      }
+    }
+  }
+
+  GlobalCodeCompletionAllocator &getAllocator() override { return *Allocator; }
+
+  CodeCompletionTUInfo &getCodeCompletionTUInfo() override { return CCTUInfo; }
+};
+} // namespace
+
+std::vector<CompletionItem> ClangdUnit::codeComplete(StringRef Contents,
+                                                     Position Pos) {
+  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<CompletionItem> Items;
+  CompletionItemsCollector Collector(&Items, CCO);
+
+  ASTUnit::RemappedFile RemappedSource(
+      FileName,
+      llvm::MemoryBuffer::getMemBufferCopy(Contents, FileName).release());
+
+  IntrusiveRefCntPtr<SourceManager> SourceMgr(
+      new SourceManager(*DiagEngine, Unit->getFileManager()));
+  // 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, Unit->getFileManager(),
+                     StoredDiagnostics, OwnedBuffers);
+  for (const llvm::MemoryBuffer *Buffer : OwnedBuffers)
+    delete Buffer;
+  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;
+  }
+  llvm_unreachable("Unknown diagnostic level!");
+}
+} // 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()};
+
+    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)});
+  }
+  return Result;
+}

clangd/ClangdUnitStore.h

+//===--- ClangdUnitStore.h - A ClangdUnits container -------------*-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_CLANGDUNITSTORE_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNITSTORE_H
+
+#include <mutex>
+
+#include "ClangdUnit.h"
+#include "GlobalCompilationDatabase.h"
+#include "Path.h"
+#include "clang/Tooling/CompilationDatabase.h"
+
+namespace clang {
+namespace clangd {
+
+/// Thread-safe collection of ASTs built for specific files. Provides
+/// synchronized access to ASTs.
+class ClangdUnitStore {
+public:
+  /// Run specified \p Action on the ClangdUnit for \p File.
+  /// If the file is not present in ClangdUnitStore, a new ClangdUnit will be
+  /// created from the \p FileContents. If the file is already present in the
+  /// store, ClangdUnit::reparse will be called with the new contents before
+  /// running \p Action.
+  template <class Func>
+  void runOnUnit(PathRef File, StringRef FileContents,
+                 GlobalCompilationDatabase &CDB,
+                 std::shared_ptr<PCHContainerOperations> PCHs, Func Action) {
+    runOnUnitImpl(File, FileContents, CDB, PCHs, /*ReparseBeforeAction=*/true,
+                  std::forward<Func>(Action));
+  }
+
+  /// Run specified \p Action on the ClangdUnit for \p File.
+  /// If the file is not present in ClangdUnitStore, a new ClangdUnit will be
+  /// created from the \p FileContents. If the file is already present in the
+  /// store, the \p Action will be run directly on it.
+  template <class Func>
+  void runOnUnitWithoutReparse(PathRef File, StringRef FileContents,
+                               GlobalCompilationDatabase &CDB,
+                               std::shared_ptr<PCHContainerOperations> PCHs,
+                               Func Action) {
+    runOnUnitImpl(File, FileContents, CDB, PCHs, /*ReparseBeforeAction=*/false,
+                  std::forward<Func>(Action));
+  }
+
+  /// Remove ClangdUnit for \p File, if any
+  void removeUnitIfPresent(PathRef File);
+
+private:
+  /// Run specified \p Action on the ClangdUnit for \p File.
+  template <class Func>
+  void runOnUnitImpl(PathRef File, StringRef FileContents,
+                     GlobalCompilationDatabase &CDB,
+                     std::shared_ptr<PCHContainerOperations> PCHs,
+                     bool ReparseBeforeAction, Func Action) {
+    std::lock_guard<std::mutex> Lock(Mutex);
+
+    auto Commands = getCompileCommands(CDB, File);
+    assert(!Commands.empty() &&
+           "getCompileCommands should add default command");
+    // chdir. This is thread hostile.
+    // FIXME(ibiryukov): get rid of this
+    llvm::sys::fs::set_current_path(Commands.front().Directory);
+
+    auto It = OpenedFiles.find(File);
+    if (It == OpenedFiles.end()) {
+      It = OpenedFiles
+               .insert(std::make_pair(
+                   File, ClangdUnit(File, FileContents, PCHs, Commands)))
+               .first;
+    } else if (ReparseBeforeAction) {
+      It->second.reparse(FileContents);
+    }
+    return Action(It->second);
+  }
+
+  std::vector<tooling::CompileCommand>
+  getCompileCommands(GlobalCompilationDatabase &CDB, PathRef File);
+
+  std::mutex Mutex;
+  llvm::StringMap<ClangdUnit> OpenedFiles;
+};
+} // namespace clangd
+} // namespace clang
+
+#endif

clangd/ClangdUnitStore.cpp

+//===--- ClangdUnitStore.cpp - A ClangdUnits container -----------*-C++-*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "ClangdUnitStore.h"
+#include "llvm/Support/Path.h"
+
+using namespace clang::clangd;
+using namespace clang;
+
+void ClangdUnitStore::removeUnitIfPresent(PathRef File) {
+  std::lock_guard<std::mutex> Lock(Mutex);
+
+  auto It = OpenedFiles.find(File);
+  if (It == OpenedFiles.end())
+    return;
+  OpenedFiles.erase(It);
+}
+
+std::vector<tooling::CompileCommand> ClangdUnitStore::getCompileCommands(GlobalCompilationDatabase &CDB, PathRef File) {
+  std::vector<tooling::CompileCommand> Commands = CDB.getCompileCommands(File);
+  if (Commands.empty()) {
+    // Add a fake command line if we know nothing.
+    Commands.push_back(tooling::CompileCommand(
+        llvm::sys::path::parent_path(File), llvm::sys::path::filename(File),
+        {"clang", "-fsyntax-only", File.str()}, ""));
+  }
+  return Commands;
+}

clangd/DraftStore.h

-//===--- DocumentStore.h - File contents container --------------*- C++ -*-===//
+//===--- DraftStore.h - File contents container -----------------*- 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_DOCUMENTSTORE_H
-#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_DOCUMENTSTORE_H
+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_DRAFTSTORE_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_DRAFTSTORE_H
 
+#include "Path.h"
 #include "clang/Basic/LLVM.h"
 #include "llvm/ADT/StringMap.h"
 #include <mutex>
 #include <string>
 #include <vector>
 
 namespace clang {
 namespace clangd {
-class DocumentStore;
 
-struct DocumentStoreListener {
-  virtual ~DocumentStoreListener() = default;
-  virtual void onDocumentAdd(StringRef File) {}
-  virtual void onDocumentRemove(StringRef File) {}
+/// Using 'unsigned' here to avoid undefined behaviour on overflow.
+typedef unsigned DocVersion;
+
+/// Document draft with a version of this draft.
+struct VersionedDraft {
+  DocVersion Version;
+  /// If the value of the field is None, draft is now deleted
+  llvm::Optional<std::string> Draft;
 };
 
-/// A container for files opened in a workspace, addressed by File. The contents
-/// are owned by the DocumentStore.
-class DocumentStore {
+/// A thread-safe container for files opened in a workspace, addressed by
+/// filenames. The contents are owned by the DraftStore. Versions are mantained
+/// for the all added documents, including removed ones. The document version is
+/// incremented on each update and removal of the document.
+class DraftStore {
 public:
-  /// Add a document to the store. Overwrites existing contents.
-  void addDocument(StringRef File, StringRef Text) {
-    {
-      std::lock_guard<std::mutex> Guard(DocsMutex);
-      Docs[File] = Text;
-    }
-    for (const auto &Listener : Listeners)
-      Listener->onDocumentAdd(File);
-  }
-  /// Delete a document from the store.
-  void removeDocument(StringRef File) {
-    {
-      std::lock_guard<std::mutex> Guard(DocsMutex);
-      Docs.erase(File);
-    }
-    for (const auto &Listener : Listeners)
-      Listener->onDocumentRemove(File);
-  }
-  /// Retrieve a document from the store. Empty string if it's unknown.
-  ///
-  /// This function is thread-safe. It returns a copy to avoid handing out
-  /// references to unguarded data.
-  std::string getDocument(StringRef File) const {
-    // FIXME: This could be a reader lock.
-    std::lock_guard<std::mutex> Guard(DocsMutex);
-    return Docs.lookup(File);
-  }
-
-  /// Add a listener. Does not take ownership.
-  void addListener(DocumentStoreListener *DSL) { Listeners.push_back(DSL); }
-
-  /// Get name and constents of all documents in this store.
-  ///
-  /// This function is thread-safe. It returns a copies to avoid handing out
-  /// references to unguarded data.
-  std::vector<std::pair<std::string, std::string>> getAllDocuments() const {
-    std::vector<std::pair<std::string, std::string>> AllDocs;
-    std::lock_guard<std::mutex> Guard(DocsMutex);
-    for (const auto &P : Docs)
-      AllDocs.emplace_back(P.first(), P.second);
-    return AllDocs;
-  }
+  /// \return version and contents of the stored document.
+  /// For untracked files, a (0, None) pair is returned.
+  VersionedDraft getDraft(PathRef File) const;
+  /// \return version of the tracked document.
+  /// For untracked files, 0 is returned.
+  DocVersion getVersion(PathRef File) const;
+
+  /// Replace contents of the draft for \p File with \p Contents.
+  /// \return The new version of the draft for \p File.
+  DocVersion updateDraft(PathRef File, StringRef Contents);
+  /// Remove the contents of the draft
+  /// \return The new version of the draft for \p File.
+  DocVersion removeDraft(PathRef File);
 
 private:
-  llvm::StringMap<std::string> Docs;
-  std::vector<DocumentStoreListener *> Listeners;
-
-  mutable std::mutex DocsMutex;
+  mutable std::mutex Mutex;
+  llvm::StringMap<VersionedDraft> Drafts;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/DraftStore.cpp

+//===--- DraftStore.cpp - File contents container ---------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "DraftStore.h"
+
+using namespace clang::clangd;
+
+VersionedDraft DraftStore::getDraft(PathRef File) const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+
+  auto It = Drafts.find(File);
+  if (It == Drafts.end())
+    return {0, llvm::None};
+  return It->second;
+}
+
+DocVersion DraftStore::getVersion(PathRef File) const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+
+  auto It = Drafts.find(File);
+  if (It == Drafts.end())
+    return 0;
+  return It->second.Version;
+}
+
+DocVersion DraftStore::updateDraft(PathRef File, StringRef Contents) {
+  std::lock_guard<std::mutex> Lock(Mutex);
+
+  auto &Entry = Drafts[File];
+  DocVersion NewVersion = ++Entry.Version;
+  Entry.Draft = Contents;
+  return NewVersion;
+}
+
+DocVersion DraftStore::removeDraft(PathRef File) {
+  std::lock_guard<std::mutex> Lock(Mutex);
+
+  auto &Entry = Drafts[File];
+  DocVersion NewVersion = ++Entry.Version;
+  Entry.Draft = llvm::None;
+  return NewVersion;
+}

clangd/GlobalCompilationDatabase.h

+//===--- GlobalCompilationDatabase.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_GLOBALCOMPILATIONDATABASE_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_GLOBALCOMPILATIONDATABASE_H
+
+#include "Path.h"
+#include "llvm/ADT/StringMap.h"
+#include <memory>
+#include <mutex>
+
+namespace clang {
+
+namespace tooling {
+class CompilationDatabase;
+struct CompileCommand;
+} // namespace tooling
+
+namespace clangd {
+
+/// Provides compilation arguments used for building ClangdUnit.
+class GlobalCompilationDatabase {
+public:
+  virtual ~GlobalCompilationDatabase() = default;
+
+  virtual std::vector<tooling::CompileCommand>
+  getCompileCommands(PathRef File) = 0;
+
+  /// FIXME(ibiryukov): add facilities to track changes to compilation flags of
+  /// existing targets.
+};
+
+/// Gets compile args from tooling::CompilationDatabases built for parent
+/// directories.
+class DirectoryBasedGlobalCompilationDatabase
+    : public GlobalCompilationDatabase {
+public:
+  std::vector<tooling::CompileCommand>
+  getCompileCommands(PathRef File) override;
+
+private:
+  tooling::CompilationDatabase *getCompilationDatabase(PathRef File);
+
+  std::mutex Mutex;
+  /// Caches compilation databases loaded from directories(keys are
+  /// directories).
+  llvm::StringMap<std::unique_ptr<clang::tooling::CompilationDatabase>>
+      CompilationDatabases;
+};
+} // namespace clangd
+} // namespace clang
+
+#endif

clangd/GlobalCompilationDatabase.cpp

+//===--- GlobalCompilationDatabase.cpp --------------------------*- C++-*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===---------------------------------------------------------------------===//
+
+#include "GlobalCompilationDatabase.h"
+#include "clang/Tooling/CompilationDatabase.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
+
+using namespace clang::clangd;
+using namespace clang;
+
+std::vector<tooling::CompileCommand>
+DirectoryBasedGlobalCompilationDatabase::getCompileCommands(PathRef File) {
+  std::vector<tooling::CompileCommand> Commands;
+
+  auto CDB = getCompilationDatabase(File);
+  if (!CDB)
+    return {};
+  return CDB->getCompileCommands(File);
+}
+
+tooling::CompilationDatabase *
+DirectoryBasedGlobalCompilationDatabase::getCompilationDatabase(PathRef File) {
+  std::lock_guard<std::mutex> Lock(Mutex);
+
+  namespace path = llvm::sys::path;
+
+  assert((path::is_absolute(File, path::Style::posix) ||
+          path::is_absolute(File, path::Style::windows)) &&
+         "path must be absolute");
+
+  for (auto Path = path::parent_path(File); !Path.empty();
+       Path = path::parent_path(Path)) {
+
+    auto CachedIt = CompilationDatabases.find(Path);
+    if (CachedIt != CompilationDatabases.end())
+      return CachedIt->second.get();
+    std::string Error;
+    auto CDB = tooling::CompilationDatabase::loadFromDirectory(Path, Error);
+    if (!CDB) {
+      if (!Error.empty()) {
+        // FIXME(ibiryukov): logging
+        // Output.log("Error when trying to load compilation database from " +
+        //            Twine(Path) + ": " + Twine(Error) + "\n");
+      }
+      continue;
+    }
+
+    // FIXME(ibiryukov): Invalidate cached compilation databases on changes
+    auto result = CDB.get();
+    CompilationDatabases.insert(std::make_pair(Path, std::move(CDB)));
+    return result;
+  }
+
+  // FIXME(ibiryukov): logging
+  // Output.log("Failed to find compilation database for " + Twine(File) +
+  // "\n");
+  return nullptr;
+}

clangd/Path.h

+//===--- Path.h - Helper typedefs --------------------------------*- 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_PATH_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_PATH_H
+
+#include "llvm/ADT/StringRef.h"
+#include <string>
+
+namespace clang {
+namespace clangd {
+
+/// A typedef to represent a file path. Used solely for more descriptive
+/// signatures.
+using Path = std::string;
+/// A typedef to represent a ref to file path. Used solely for more descriptive
+/// signatures.
+using PathRef = llvm::StringRef;
+
+} // namespace clangd
+} // namespace clang
+
+#endif

clangd/ProtocolHandlers.h

 
 #include "JSONRPCDispatcher.h"
 #include "Protocol.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace clang {
 namespace clangd {
-class ASTManager;
-class DocumentStore;
+class ClangdLSPServer;
+class ClangdLSPServer;
 
 struct InitializeHandler : Handler {
   InitializeHandler(JSONOutput &Output) : Handler(Output) {}
 
   void handleMethod(llvm::yaml::MappingNode *Params, StringRef ID) override {
     writeMessage(
         R"({"jsonrpc":"2.0","id":)" + ID +
         R"(,"result":{"capabilities":{
struct ShutdownHandler : Handler {
 
   bool isDone() const { return IsDone; }
 
 private:
   bool IsDone = false;
 };
 
 struct TextDocumentDidOpenHandler : Handler {
-  TextDocumentDidOpenHandler(JSONOutput &Output, DocumentStore &Store)
-      : Handler(Output), Store(Store) {}
+  TextDocumentDidOpenHandler(JSONOutput &Output, ClangdLSPServer &AST)
+      : Handler(Output), AST(AST) {}
 
   void handleNotification(llvm::yaml::MappingNode *Params) override;
 
 private:
-  DocumentStore &Store;
+  ClangdLSPServer *
 };
 
 struct TextDocumentDidChangeHandler : Handler {
-  TextDocumentDidChangeHandler(JSONOutput &Output, DocumentStore &Store)
-      : Handler(Output), Store(Store) {}
+  TextDocumentDidChangeHandler(JSONOutput &Output, ClangdLSPServer &AST)
+      : Handler(Output), AST(AST) {}
 
   void handleNotification(llvm::yaml::MappingNode *Params) override;
 
 private:
-  DocumentStore &Store;
+  ClangdLSPServer *
 };
 
 struct TextDocumentDidCloseHandler : Handler {
-  TextDocumentDidCloseHandler(JSONOutput &Output, DocumentStore &Store)
-      : Handler(Output), Store(Store) {}
+  TextDocumentDidCloseHandler(JSONOutput &Output, ClangdLSPServer &AST)
+      : Handler(Output), AST(AST) {}
 
   void handleNotification(llvm::yaml::MappingNode *Params) override;
 
 private:
-  DocumentStore &Store;
+  ClangdLSPServer *
 };
 
 struct TextDocumentOnTypeFormattingHandler : Handler {
-  TextDocumentOnTypeFormattingHandler(JSONOutput &Output, DocumentStore &Store)
-      : Handler(Output), Store(Store) {}
+  TextDocumentOnTypeFormattingHandler(JSONOutput &Output, ClangdLSPServer &AST)
+      : Handler(Output), AST(AST) {}
 
   void handleMethod(llvm::yaml::MappingNode *Params, StringRef ID) override;
 
 private:
-  DocumentStore &Store;
+  ClangdLSPServer *
 };
 
 struct TextDocumentRangeFormattingHandler : Handler {
-  TextDocumentRangeFormattingHandler(JSONOutput &Output, DocumentStore &Store)
-      : Handler(Output), Store(Store) {}
+  TextDocumentRangeFormattingHandler(JSONOutput &Output, ClangdLSPServer &AST)
+      : Handler(Output), AST(AST) {}
 
   void handleMethod(llvm::yaml::MappingNode *Params, StringRef ID) override;
 
 private:
-  DocumentStore &Store;
+  ClangdLSPServer *
 };
 
 struct TextDocumentFormattingHandler : Handler {
-  TextDocumentFormattingHandler(JSONOutput &Output, DocumentStore &Store)
-      : Handler(Output), Store(Store) {}
+  TextDocumentFormattingHandler(JSONOutput &Output, ClangdLSPServer &AST)
+      : Handler(Output), AST(AST) {}
 
   void handleMethod(llvm::yaml::MappingNode *Params, StringRef ID) override;
 
 private:
-  DocumentStore &Store;
+  ClangdLSPServer *
 };
 
 struct CodeActionHandler : Handler {
-  CodeActionHandler(JSONOutput &Output, ASTManager &AST)
+  CodeActionHandler(JSONOutput &Output, ClangdLSPServer &AST)
       : Handler(Output), AST(AST) {}
 
   void handleMethod(llvm::yaml::MappingNode *Params, StringRef ID) override;
 
 private:
-  ASTManager *
+  ClangdLSPServer *
 };
 
 struct CompletionHandler : Handler {
-  CompletionHandler(JSONOutput &Output, ASTManager &AST)
+  CompletionHandler(JSONOutput &Output, ClangdLSPServer &AST)
       : Handler(Output), AST(AST) {}
 
   void handleMethod(llvm::yaml::MappingNode *Params, StringRef ID) override;
 
  private:
-  ASTManager *
+  ClangdLSPServer *
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/ProtocolHandlers.cpp

 //===--- ProtocolHandlers.cpp - LSP callbacks -----------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "ProtocolHandlers.h"
-#include "ASTManager.h"
-#include "DocumentStore.h"
+#include "ClangdServer.h"
+#include "DraftStore.h"
 #include "clang/Format/Format.h"
+#include "ClangdLSPServer.h"
 using namespace clang;
 using namespace clangd;
 
 void TextDocumentDidOpenHandler::handleNotification(
     llvm::yaml::MappingNode *Params) {
   auto DOTDP = DidOpenTextDocumentParams::parse(Params);
   if (!DOTDP) {
     Output.log("Failed to decode DidOpenTextDocumentParams!\n");
     return;
   }
-  Store.addDocument(DOTDP->textDocument.uri.file, DOTDP->textDocument.text);
+  AST.openDocument(DOTDP->textDocument.uri.file, DOTDP->textDocument.text);
 }
 
 void TextDocumentDidCloseHandler::handleNotification(
     llvm::yaml::MappingNode *Params) {
   auto DCTDP = DidCloseTextDocumentParams::parse(Params);
   if (!DCTDP) {
     Output.log("Failed to decode DidCloseTextDocumentParams!\n");
     return;
   }
 
-  Store.removeDocument(DCTDP->textDocument.uri.file);
+  AST.closeDocument(DCTDP->textDocument.uri.file);
 }
 
 void TextDocumentDidChangeHandler::handleNotification(
     llvm::yaml::MappingNode *Params) {
   auto DCTDP = DidChangeTextDocumentParams::parse(Params);
   if (!DCTDP || DCTDP->contentChanges.size() != 1) {
     Output.log("Failed to decode DidChangeTextDocumentParams!\n");
     return;
   }
   // We only support full syncing right now.
-  Store.addDocument(DCTDP->textDocument.uri.file, DCTDP->contentChanges[0].text);
+  AST.openDocument(DCTDP->textDocument.uri.file, DCTDP->contentChanges[0].text);
 }
 
 /// Turn a [line, column] pair into an offset in Code.
 static size_t positionToOffset(StringRef Code, Position P) {
   size_t Offset = 0;
   for (int I = 0; I != P.line; ++I) {
     // FIXME: \r\n
     // FIXME: UTF-8

 void TextDocumentRangeFormattingHandler::handleMethod(
     llvm::yaml::MappingNode *Params, StringRef ID) {
   auto DRFP = DocumentRangeFormattingParams::parse(Params);
   if (!DRFP) {
     Output.log("Failed to decode DocumentRangeFormattingParams!\n");
     return;
   }
 
-  std::string Code = Store.getDocument(DRFP->textDocument.uri.file);
+  std::string Code = AST.getDocument(DRFP->textDocument.uri.file);
 
   size_t Begin = positionToOffset(Code, DRFP->range.start);
   size_t Len = positionToOffset(Code, DRFP->range.end) - Begin;
 
   writeMessage(formatCode(Code, DRFP->textDocument.uri.file,
                           {clang::tooling::Range(Begin, Len)}, ID));
 }
 
 void TextDocumentOnTypeFormattingHandler::handleMethod(
     llvm::yaml::MappingNode *Params, StringRef ID) {
   auto DOTFP = DocumentOnTypeFormattingParams::parse(Params);
   if (!DOTFP) {
     Output.log("Failed to decode DocumentOnTypeFormattingParams!\n");
     return;
   }
 
   // Look for the previous opening brace from the character position and format
   // starting from there.
-  std::string Code = Store.getDocument(DOTFP->textDocument.uri.file);
+  std::string Code = AST.getDocument(DOTFP->textDocument.uri.file);
   size_t CursorPos = positionToOffset(Code, DOTFP->position);
   size_t PreviousLBracePos = StringRef(Code).find_last_of('{', CursorPos);
   if (PreviousLBracePos == StringRef::npos)
     PreviousLBracePos = CursorPos;
   size_t Len = 1 + CursorPos - PreviousLBracePos;
 
   writeMessage(formatCode(Code, DOTFP->textDocument.uri.file,
                           {clang::tooling::Range(PreviousLBracePos, Len)}, ID));
 }
 
 void TextDocumentFormattingHandler::handleMethod(
     llvm::yaml::MappingNode *Params, StringRef ID) {
   auto DFP = DocumentFormattingParams::parse(Params);
   if (!DFP) {
     Output.log("Failed to decode DocumentFormattingParams!\n");
     return;
   }
 
   // Format everything.
-  std::string Code = Store.getDocument(DFP->textDocument.uri.file);
+  std::string Code = AST.getDocument(DFP->textDocument.uri.file);
   writeMessage(formatCode(Code, DFP->textDocument.uri.file,
                           {clang::tooling::Range(0, Code.size())}, ID));
 }
 
 void CodeActionHandler::handleMethod(llvm::yaml::MappingNode *Params,
                                      StringRef ID) {
   auto CAP = CodeActionParams::parse(Params);
   if (!CAP) {
     Output.log("Failed to decode CodeActionParams!\n");
     return;
   }
 
   // We provide a code action for each diagnostic at the requested location
   // which has FixIts available.
-  std::string Code = AST.getStore().getDocument(CAP->textDocument.uri.file);
+  std::string Code = AST.getDocument(CAP->textDocument.uri.file);
   std::string Commands;
   for (Diagnostic &D : CAP->context.diagnostics) {
     std::vector<clang::tooling::Replacement> Fixes = AST.getFixIts(CAP->textDocument.uri.file, D);
     std::string Edits = replacementsToEdits(Code, Fixes);
 
     if (!Edits.empty())
       Commands +=
           R"({"title":"Apply FixIt ')" + llvm::yaml::escape(D.message) +

 void CompletionHandler::handleMethod(llvm::yaml::MappingNode *Params,
                                      StringRef ID) {
   auto TDPP = TextDocumentPositionParams::parse(Params);
   if (!TDPP) {
     Output.log("Failed to decode TextDocumentPositionParams!\n");
     return;
   }
 
-  auto Items = AST.codeComplete(TDPP->textDocument.uri.file, TDPP->position.line,
-                                TDPP->position.character);
+  auto Items = AST.codeComplete(TDPP->textDocument.uri.file, Position{TDPP->position.line,
+          TDPP->position.character});
   std::string Completions;
   for (const auto &Item : Items) {
     Completions += CompletionItem::unparse(Item);
     Completions += ",";
   }
   if (!Completions.empty())
     Completions.pop_back();
   writeMessage(
       R"({"jsonrpc":"2.0","id":)" + ID.str() +
       R"(,"result":[)" + Completions + R"(]})");
 }