Remove ASTUnits for closed documents and cache CompilationDatabase per directory in clangd.

[ilya-biryukov]

Summary

This commit adds the textDocument/didClose handler and adds code to remove ASTUnits
for those documents.
It also changes caching strategy of CompilationDatabase to per-directory rather than per-file.
As part of the change, FixIts were moved to per-file cache.
As a result, they are now synchronized with the mutex used during parsing,
I'm planning to restore the previous behaviour(fixits are immideately availble during parsing)
in a later commit, see TODOs in the code

[krasimir]

Looks great! I'm wondering, can you think of ways to test the didClose method similarly to how it's done for other handlers?

clangd/ASTManager.cpp

203

Why not make the locking explicit here and don't require handleRequest and parseFileAndPublishDiagnostics to be called under a lock?

clangd/ASTManager.h

60

I'd call the second parameter just Uri.

64

I'd call this Uri.

65

I'd capitalize the member variables. In general, this is the llvm/clang style and the lowercase member variables in Protocol.h are for consistency with the Language Server Protocol.

118

nit: s/immideately/immediately and add a full stop.

121

It's sad that handle request should be called under the lock :(

123

Maybe it's worth noting whether the main or the worker thread calls these during asynchronous execution.

134

nit: missing trailing full stop.

135

I'd rename this lock to ClangObjectLock or something else; its name currently gives the wrong feeling that it's used only for DocDatas.

[ilya-biryukov]

Minor fixes.
Fixed variable name issues and comment spelling errors.

[ilya-biryukov]

Moved locking of ClangObjectLock into request handlers.

[ilya-biryukov]

Addressed the locking comments. Locking inside the request handlers looks much nicer indeed.

clangd/ASTManager.cpp

203

Good point, my initial attempt was aimed at avoiding locking when RunSynchronously is true. But that's not a very useful constraint.

[ilya-biryukov]

Regarding the testing. It looks like we have two ways of testing this:

1. Add clangd-specific protocol handlers that output useful stats(i.e. currently opened documents), use those in tests.
2. Add unit tests that emulate the protocol actions and run the checks in C++.

Option 1 seems preferable, we could probably use the 'workspace/executeCommand' for that.

[ilya-biryukov]

Fixed conflicts with the upstream

[krasimir]

Looks good!

clangd/ASTManager.h

 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 Uri,
+                    DocVersion Version);
+
+  ASTManagerRequestType Type;
+  std::string Uri;
+  DocVersion Version;
+};
+
 class ASTManager : public DocumentStoreListener {
 public:
   ASTManager(JSONOutput &Output, DocumentStore &Store, bool RunSynchronously);
   ~ASTManager() override;
 
   void onDocumentAdd(StringRef Uri) override;
-  // FIXME: Implement onDocumentRemove
+  void onDocumentRemove(StringRef Uri) 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 as replacements.
+  /// 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(const clangd::Diagnostic &D);
+  getFixIts(StringRef Uri, 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 URI. May return nullptr if it fails. The
   /// database is cached for subsequent accesses.
   clang::tooling::CompilationDatabase *
   getOrCreateCompilationDatabaseForFile(StringRef Uri);
   // Creates a new ASTUnit for the document at Uri.
   // FIXME: This calls chdir internally, which is thread unsafe.
   std::unique_ptr<clang::ASTUnit>
   createASTUnitForFile(StringRef Uri, const DocumentStore &Docs);
 
-  void runWorker();
-  void parseFileAndPublishDiagnostics(StringRef File);
+  /// 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 Uri);
 
-  /// Clang objects.
+  void runWorker();
+  void handleRequest(ASTManagerRequestType RequestType, StringRef Uri);
 
-  /// A map from Uri-s to ASTUnit-s. Guarded by \c ASTLock. ASTUnit-s are used
-  /// for generating diagnostics and fix-it-s asynchronously by the worker
-  /// thread and synchronously for code completion.
-  ///
-  /// TODO(krasimir): code completion should always have priority over parsing
-  /// for diagnostics.
-  llvm::StringMap<std::unique_ptr<clang::ASTUnit>> ASTs;
-  /// A lock for access to the map \c ASTs.
-  std::mutex ASTLock;
+  /// 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 Uri);
 
+  /// Caches compilation databases loaded from directories(keys are directories).
   llvm::StringMap<std::unique_ptr<clang::tooling::CompilationDatabase>>
       CompilationDatabases;
+
+  /// Clang objects.
+  /// A map from Uri-s 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;
 
-  typedef std::map<clangd::Diagnostic, std::vector<clang::tooling::Replacement>>
-      DiagnosticToReplacementMap;
-  DiagnosticToReplacementMap FixIts;
-  std::mutex FixItLock;
+  /// 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;
 
-  /// Queue of requests.
-  std::deque<std::string> RequestQueue;
+  /// 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 and Done.
+  /// Lock for accesses to RequestQueue, DocVersions and Done.
   std::mutex RequestLock;
 
-  /// We run parsing on a separate thread. This thread looks into PendingRequest
-  /// as a 'one element work queue' as the queue is non-empty.
+  /// 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/ASTManager.cpp

 #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;
 
+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 {};
+}
+
+ASTManagerRequest::ASTManagerRequest(ASTManagerRequestType Type,
+                                     std::string Uri,
+                                     DocVersion Version)
+    : Type(Type), Uri(Uri), 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;

 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() {
   while (true) {
-    std::string File;
+    ASTManagerRequest Request;
 
+    // Pick request from the queue
     {
       std::unique_lock<std::mutex> Lock(RequestLock);
-      // Check if there's another request pending. We keep parsing until
-      // our one-element queue is empty.
+      // Wait for more requests.
       ClangRequestCV.wait(Lock,
                           [this] { return !RequestQueue.empty() || Done; });
-
-      if (RequestQueue.empty() && Done)
+      if (Done)
         return;
+      assert(!RequestQueue.empty() && "RequestQueue was empty");
 
-      File = std::move(RequestQueue.back());
+      Request = std::move(RequestQueue.back());
       RequestQueue.pop_back();
-    } // unlock.
 
-    parseFileAndPublishDiagnostics(File);
+      // Skip outdated requests
+      if (Request.Version != DocVersions.find(Request.Uri)->second) {
+        Output.log("Version for " + Twine(Request.Uri) +
+                   " in request is outdated, skipping request\n");
+        continue;
+      }
+    } // unlock RequestLock
+
+    handleRequest(Request.Type, Request.Uri);
   }
 }
 
-void ASTManager::parseFileAndPublishDiagnostics(StringRef File) {
-  DiagnosticToReplacementMap LocalFixIts; // Temporary storage
-  std::string Diagnostics;
-  {
-    std::lock_guard<std::mutex> ASTGuard(ASTLock);
-    auto &Unit = ASTs[File]; // Only one thread can access this at a time.
+void ASTManager::queueOrRun(ASTManagerRequestType RequestType, StringRef Uri) {
+  if (RunSynchronously) {
+    handleRequest(RequestType, Uri);
+    return;
+  }
+
+  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[Uri];
+  RequestQueue.push_back(ASTManagerRequest(RequestType, Uri, version));
+  ClangRequestCV.notify_one();
+}
+
+void ASTManager::handleRequest(ASTManagerRequestType RequestType,
+                               StringRef Uri) {
+  switch (RequestType) {
+  case ASTManagerRequestType::ParseAndPublishDiagnostics:
+    parseFileAndPublishDiagnostics(Uri);
+    break;
+  case ASTManagerRequestType::RemoveDocData: {
+    std::lock_guard<std::mutex> Lock(ClangObjectLock);
+    auto DocDataIt = DocDatas.find(Uri);
+    // 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
+  }
+}
 
-    if (!Unit) {
-      Unit = createASTUnitForFile(File, this->Store);
-    } 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));
-    }
+void ASTManager::parseFileAndPublishDiagnostics(StringRef Uri) {
+  std::unique_lock<std::mutex> ClangObjectLockGuard(ClangObjectLock);
+
+  auto &DocData = DocDatas[Uri];
+  ASTUnit *Unit = DocData.getAST();
+    auto newAST = createASTUnitForFile(Uri, this->Store);
+    Unit = newAST.get();
+
+    DocData.setAST(std::move(newAST));
 
-    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.
-    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"("},)";
+  std::string Diagnostics;
+  DocData::DiagnosticToReplacementMap LocalFixIts; // Temporary storage
 
-      // 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));
-      }
-    }
-  } // unlock ASTLock
+    clangd::Diagnostic Diag = {R, getSeverity(D->getLevel()), D->getMessage()};
 
   // Put FixIts into place.
-  {
-    std::lock_guard<std::mutex> Guard(FixItLock);
-    FixIts = std::move(LocalFixIts);
-  }
+  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":")" +
-      File + R"(","diagnostics":[)" + Diagnostics + R"(]}})");
+      Uri + R"(","diagnostics":[)" + Diagnostics + R"(]}})");
 }
 
 ASTManager::~ASTManager() {
   {
     std::lock_guard<std::mutex> Guard(RequestLock);
     // Wake up the clang worker thread, then exit.
     Done = true;
     ClangRequestCV.notify_one();
-  }
+  } // unlock DocDataLock
   ClangWorker.join();
 }
 
 void ASTManager::onDocumentAdd(StringRef Uri) {
-  if (RunSynchronously) {
-    parseFileAndPublishDiagnostics(Uri);
-    return;
-  }
-  std::lock_guard<std::mutex> Guard(RequestLock);
-  // Currently we discard all pending requests and just enqueue the latest one.
-  RequestQueue.clear();
-  RequestQueue.push_back(Uri);
-  ClangRequestCV.notify_one();
+  queueOrRun(ASTManagerRequestType::ParseAndPublishDiagnostics, Uri);
+
+void ASTManager::onDocumentRemove(StringRef Uri) {
+  queueOrRun(ASTManagerRequestType::RemoveDocData, Uri);
 }
 
 tooling::CompilationDatabase *
 ASTManager::getOrCreateCompilationDatabaseForFile(StringRef Uri) {
-  auto &I = CompilationDatabases[Uri];
-  if (I)
-    return I.get();
+  namespace path = llvm::sys::path;
 
   Uri.consume_front("file://");
+  assert(path::is_absolute(Uri) && "path must be absolute");
+
+  for (auto Path = path::parent_path(Uri); !Path.empty();
+       Path = path::parent_path(Path)) {
+
+    auto CachedIt = CompilationDatabases.find(Path);
+    if (CachedIt != CompilationDatabases.end())
+      return CachedIt->second.get();
 
     std::string Error;
-  I = tooling::CompilationDatabase::autoDetectFromSource(Uri, Error);
-  Output.log("Failed to load compilation database: " + Twine(Error) + "\n");
-  return I.get();
+    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;
+    }
+
+    // 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(Uri) + "\n");
+  return nullptr;
 }
 
 std::unique_ptr<clang::ASTUnit>
 ASTManager::createASTUnitForFile(StringRef Uri, const DocumentStore &Docs) {
   tooling::CompilationDatabase *CDB =
       getOrCreateCompilationDatabaseForFile(Uri);
 
   Uri.consume_front("file://");
return std::unique_ptr<clang::ASTUnit>(ASTUnit::LoadFromCommandLine(
       /*RemappedFilesKeepOriginalName=*/true,
       /*PrecompilePreambleAfterNParses=*/1, /*TUKind=*/TU_Complete,
       /*CacheCodeCompletionResults=*/true,
       /*IncludeBriefCommentsInCodeCompletion=*/true,
       /*AllowPCHWithCompilerErrors=*/true));
 }
 
 std::vector<clang::tooling::Replacement>
-ASTManager::getFixIts(const clangd::Diagnostic &D) {
-  std::lock_guard<std::mutex> Guard(FixItLock);
-  auto I = FixIts.find(D);
-  if (I != FixIts.end())
-    return I->second;
-  return {};
+ASTManager::getFixIts(StringRef Uri, 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[Uri].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)
ASTManager::codeComplete(StringRef Uri, unsigned Line, unsigned Column) {
   // 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(ASTLock);
-  auto &Unit = ASTs[Uri];
-  if (!Unit)
-    Unit = createASTUnitForFile(Uri, this->Store);
+
+  std::lock_guard<std::mutex> Guard(ClangObjectLock);
+  auto &DocData = DocDatas[Uri];
+  auto Unit = DocData.getAST();
+  if (!Unit) {
+    auto newAST = createASTUnitForFile(Uri, this->Store);
+    Unit = newAST.get();
+    DocData.setAST(std::move(newAST));
+  }
   if (!Unit)
     return {};
   IntrusiveRefCntPtr<SourceManager> SourceMgr(
       new SourceManager(*DiagEngine, Unit->getFileManager()));
   StringRef File(Uri);
   File.consume_front("file://");
   // CodeComplete seems to require fresh LangOptions.
   LangOptions LangOpts = Unit->getLangOpts();

clangd/ClangDMain.cpp

int main(int argc, char *argv[]) {
   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));
-  // FIXME: Implement textDocument/didClose.
+  Dispatcher.registerHandler(
+      "textDocument/didClose",
+      llvm::make_unique<TextDocumentDidCloseHandler>(Out, Store));
   Dispatcher.registerHandler(
       "textDocument/didChange",
       llvm::make_unique<TextDocumentDidChangeHandler>(Out, Store));
   Dispatcher.registerHandler(
       "textDocument/rangeFormatting",
       llvm::make_unique<TextDocumentRangeFormattingHandler>(Out, Store));
   Dispatcher.registerHandler(
       "textDocument/onTypeFormatting",

clangd/Protocol.h

 struct DidOpenTextDocumentParams {
   /// The document that was opened.
   TextDocumentItem textDocument;
 
   static llvm::Optional<DidOpenTextDocumentParams>
   parse(llvm::yaml::MappingNode *Params);
 };
 
+struct DidCloseTextDocumentParams {
+  /// The document that was closed.
+  TextDocumentIdentifier textDocument;
+
+  static llvm::Optional<DidCloseTextDocumentParams>
+  parse(llvm::yaml::MappingNode *Params);
+};
+
 struct TextDocumentContentChangeEvent {
   /// The new text of the document.
   std::string text;
 
   static llvm::Optional<TextDocumentContentChangeEvent>
   parse(llvm::yaml::MappingNode *Params);
 };
 

clangd/Protocol.cpp

if (KeyValue == "textDocument") {
       Result.textDocument = std::move(*Parsed);
     } else {
       return llvm::None;
     }
   }
   return Result;
 }
 
+llvm::Optional<DidCloseTextDocumentParams>
+DidCloseTextDocumentParams::parse(llvm::yaml::MappingNode *Params) {
+  DidCloseTextDocumentParams Result;
+  for (auto &NextKeyValue : *Params) {
+    auto *KeyString = dyn_cast<llvm::yaml::ScalarNode>(NextKeyValue.getKey());
+    if (!KeyString)
+      return llvm::None;
+
+    llvm::SmallString<10> KeyStorage;
+    StringRef KeyValue = KeyString->getValue(KeyStorage);
+    auto *Value = NextKeyValue.getValue();
+
+    if (KeyValue == "textDocument") {
+      auto *Map = dyn_cast<llvm::yaml::MappingNode>(Value);
+      if (!Map)
+        return llvm::None;
+      auto Parsed = TextDocumentIdentifier::parse(Map);
+      if (!Parsed)
+        return llvm::None;
+      Result.textDocument = std::move(*Parsed);
+    } else {
+      return llvm::None;
+    }
+  }
+  return Result;
+}
+
 llvm::Optional<DidChangeTextDocumentParams>
 DidChangeTextDocumentParams::parse(llvm::yaml::MappingNode *Params) {
   DidChangeTextDocumentParams Result;
   for (auto &NextKeyValue : *Params) {
     auto *KeyString = dyn_cast<llvm::yaml::ScalarNode>(NextKeyValue.getKey());
     if (!KeyString)
       return llvm::None;
 

clangd/ProtocolHandlers.h

TextDocumentDidChangeHandler(JSONOutput &Output, DocumentStore &Store)
       : Handler(Output), Store(Store) {}
 
   void handleNotification(llvm::yaml::MappingNode *Params) override;
 
 private:
   DocumentStore &Store;
 };
 
+struct TextDocumentDidCloseHandler : Handler {
+  TextDocumentDidCloseHandler(JSONOutput &Output, DocumentStore &Store)
+      : Handler(Output), Store(Store) {}
+
+  void handleNotification(llvm::yaml::MappingNode *Params) override;
+
+private:
+  DocumentStore &Store;
+};
+
 struct TextDocumentOnTypeFormattingHandler : Handler {
   TextDocumentOnTypeFormattingHandler(JSONOutput &Output, DocumentStore &Store)
       : Handler(Output), Store(Store) {}
 
   void handleMethod(llvm::yaml::MappingNode *Params, StringRef ID) override;
 
 private:
   DocumentStore &Store;

clangd/ProtocolHandlers.cpp

void TextDocumentDidOpenHandler::handleNotification(
   auto DOTDP = DidOpenTextDocumentParams::parse(Params);
   if (!DOTDP) {
     Output.log("Failed to decode DidOpenTextDocumentParams!\n");
     return;
   }
   Store.addDocument(DOTDP->textDocument.uri, 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);
+}
+
 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.
if (!CAP) {
     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);
   std::string Commands;
   for (Diagnostic &D : CAP->context.diagnostics) {
-    std::vector<clang::tooling::Replacement> Fixes = AST.getFixIts(D);
+    std::vector<clang::tooling::Replacement> Fixes = AST.getFixIts(CAP->textDocument.uri, D);
     std::string Edits = replacementsToEdits(Code, Fixes);
 
     if (!Edits.empty())
       Commands +=
           R"({"title":"Apply FixIt ')" + llvm::yaml::escape(D.message) +
           R"('", "command": "clangd.applyFix", "arguments": [")" +
           llvm::yaml::escape(CAP->textDocument.uri) +
           R"(", [)" + Edits +