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

Contributed by ilya-biryukov!

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

llvm-svn: 299843

clang-tools-extra/clangd/ASTManager.cpp

@@ -20,6 +20,29 @@
 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 File,
+                                     DocVersion Version)
+    : Type(Type), 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>
@@ -61,82 +84,125 @@ ASTManager::ASTManager(JSONOutput &Output, DocumentStore &Store,
 
 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.
 
+      // 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");
+        continue;
+      }
+    } // unlock RequestLock
+
+    handleRequest(Request.Type, Request.File);
+  }
+}
+
+void ASTManager::queueOrRun(ASTManagerRequestType RequestType, StringRef File) {
+  if (RunSynchronously) {
+    handleRequest(RequestType, File);
+    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[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
   }
 }
 
 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.
+  std::unique_lock<std::mutex> ClangObjectLockGuard(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));
-    }
+  auto &DocData = DocDatas[File];
+  ASTUnit *Unit = DocData.getAST();
+  if (!Unit) {
+    auto newAST = createASTUnitForFile(File, this->Store);
+    Unit = newAST.get();
 
-    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"("},)";
-
-      // 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
-
-  // Put FixIts into place.
-  {
-    std::lock_guard<std::mutex> Guard(FixItLock);
-    FixIts = std::move(LocalFixIts);
+    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(
@@ -150,32 +216,48 @@ ASTManager::~ASTManager() {
     // Wake up the clang worker thread, then exit.
     Done = true;
     ClangRequestCV.notify_one();
-  }
+  } // unlock DocDataLock
   ClangWorker.join();
 }
 
 void ASTManager::onDocumentAdd(StringRef File) {
-  if (RunSynchronously) {
-    parseFileAndPublishDiagnostics(File);
-    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(File);
-  ClangRequestCV.notify_one();
+  queueOrRun(ASTManagerRequestType::ParseAndPublishDiagnostics, File);
+}
+
+void ASTManager::onDocumentRemove(StringRef File) {
+  queueOrRun(ASTManagerRequestType::RemoveDocData, File);
 }
 
 tooling::CompilationDatabase *
 ASTManager::getOrCreateCompilationDatabaseForFile(StringRef File) {
-  auto &I = CompilationDatabases[File];
-  if (I)
-    return I.get();
+  namespace path = llvm::sys::path;
 
-  std::string Error;
-  I = tooling::CompilationDatabase::autoDetectFromSource(File, Error);
-  Output.log("Failed to load compilation database: " + Twine(Error) + "\n");
-  return I.get();
+  assert(path::is_absolute(File) && "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()) {
+        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(File) + "\n");
+  return nullptr;
 }
 
 std::unique_ptr<clang::ASTUnit>
@@ -225,16 +307,14 @@ ASTManager::createASTUnitForFile(StringRef File, const DocumentStore &Docs) {
 }
 
 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 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;
@@ -285,10 +365,15 @@ ASTManager::codeComplete(StringRef File, unsigned Line, unsigned Column) {
       new DiagnosticsEngine(new DiagnosticIDs, new DiagnosticOptions));
   std::vector<CompletionItem> Items;
   CompletionItemsCollector Collector(&Items, CCO);
-  std::lock_guard<std::mutex> Guard(ASTLock);
-  auto &Unit = ASTs[File];
-  if (!Unit)
-    Unit = createASTUnitForFile(File, this->Store);
+
+  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(

clang-tools-extra/clangd/ASTManager.h

@@ -29,13 +29,49 @@ class CompilationDatabase;
 
 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;
-  // FIXME: Implement onDocumentRemove
+  void onDocumentRemove(StringRef File) override;
 
   /// Get code completions at a specified \p Line and \p Column in \p File.
   ///
@@ -44,12 +80,13 @@ public:
   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 File, const clangd::Diagnostic &D);
 
   DocumentStore &getStore() const { return Store; }
 
@@ -70,41 +107,52 @@ private:
   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);
 
-  /// Clang objects.
-
-  /// A map from File-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;
-
+  /// 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;
 
-  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;
 };
 

clang-tools-extra/clangd/ClangDMain.cpp

@@ -46,7 +46,9 @@ int main(int argc, char *argv[]) {
   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));

clang-tools-extra/clangd/Protocol.cpp

@@ -262,6 +262,33 @@ DidOpenTextDocumentParams::parse(llvm::yaml::MappingNode *Params) {
   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;

clang-tools-extra/clangd/Protocol.h

@@ -124,6 +124,14 @@ struct 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;

clang-tools-extra/clangd/ProtocolHandlers.cpp

@@ -24,6 +24,17 @@ void TextDocumentDidOpenHandler::handleNotification(
   Store.addDocument(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);
+}
+
 void TextDocumentDidChangeHandler::handleNotification(
     llvm::yaml::MappingNode *Params) {
   auto DCTDP = DidChangeTextDocumentParams::parse(Params);
@@ -156,7 +167,7 @@ void CodeActionHandler::handleMethod(llvm::yaml::MappingNode *Params,
   std::string Code = AST.getStore().getDocument(CAP->textDocument.uri.file);
   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.file, D);
     std::string Edits = replacementsToEdits(Code, Fixes);
 
     if (!Edits.empty())

clang-tools-extra/clangd/ProtocolHandlers.h

@@ -75,6 +75,16 @@ 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) {}