Move DraftMgr from ClangdServer to ClangdLSPServer

Summary:
This patch moves the draft manager closer to the edge of Clangd, from
ClangdServer to ClangdLSPServer.  This will make it easier to implement
incremental document sync, by making ClangdServer only deal with
complete documents.

As a result, DraftStore doesn't have to deal with versioning, and thus
its API can be simplified.  It is replaced by a StringMap in
ClangdServer holding a current version number for each file.

Signed-off-by: Simon Marchi <simon.marchi@ericsson.com>

Subscribers: klimek, mgorny, ilya-biryukov, jkorous-apple, ioeric, cfe-commits

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

llvm-svn: 327711

7 files changed, 75 insertions, 126 deletions

diff --git a/clang-tools-extra/clangd/ClangdLSPServer.cpp b/clang-tools-extra/clangd/ClangdLSPServer.cpp
index 7eb02a32bbb..8b7e7148584 100644
--- a/clang-tools-extra/clangd/ClangdLSPServer.cpp
+++ b/clang-tools-extra/clangd/ClangdLSPServer.cpp
@@ -12,6 +12,7 @@
 #include "JSONRPCDispatcher.h"
 #include "SourceCode.h"
 #include "URI.h"
+#include "llvm/Support/Errc.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 
@@ -142,8 +143,12 @@ void ClangdLSPServer::onDocumentDidOpen(DidOpenTextDocumentParams &Params) {
   if (Params.metadata && !Params.metadata->extraFlags.empty())
     CDB.setExtraFlagsForFile(Params.textDocument.uri.file(),
                              std::move(Params.metadata->extraFlags));
-  Server.addDocument(Params.textDocument.uri.file(), Params.textDocument.text,
-                     WantDiagnostics::Yes);
+
+  PathRef File = Params.textDocument.uri.file();
+  std::string &Contents = Params.textDocument.text;
+
+  DraftMgr.updateDraft(File, Contents);
+  Server.addDocument(File, Contents, WantDiagnostics::Yes);
 }
 
 void ClangdLSPServer::onDocumentDidChange(DidChangeTextDocumentParams &Params) {
@@ -154,9 +159,13 @@ void ClangdLSPServer::onDocumentDidChange(DidChangeTextDocumentParams &Params) {
   if (Params.wantDiagnostics.hasValue())
     WantDiags = Params.wantDiagnostics.getValue() ? WantDiagnostics::Yes
                                                   : WantDiagnostics::No;
+
+  PathRef File = Params.textDocument.uri.file();
+  std::string &Contents = Params.contentChanges[0].text;
+
   // We only support full syncing right now.
-  Server.addDocument(Params.textDocument.uri.file(),
-                     Params.contentChanges[0].text, WantDiags);
+  DraftMgr.updateDraft(File, Contents);
+  Server.addDocument(File, Contents, WantDiags);
 }
 
 void ClangdLSPServer::onFileEvent(DidChangeWatchedFilesParams &Params) {
@@ -188,7 +197,7 @@ void ClangdLSPServer::onCommand(ExecuteCommandParams &Params) {
   } else if (Params.command ==
              ExecuteCommandParams::CLANGD_INSERT_HEADER_INCLUDE) {
     auto &FileURI = Params.includeInsertion->textDocument.uri;
-    auto Code = Server.getDocument(FileURI.file());
+    auto Code = DraftMgr.getDraft(FileURI.file());
     if (!Code)
       return replyError(ErrorCode::InvalidParams,
                         ("command " +
@@ -233,7 +242,7 @@ void ClangdLSPServer::onCommand(ExecuteCommandParams &Params) {
 
 void ClangdLSPServer::onRename(RenameParams &Params) {
   Path File = Params.textDocument.uri.file();
-  llvm::Optional<std::string> Code = Server.getDocument(File);
+  llvm::Optional<std::string> Code = DraftMgr.getDraft(File);
   if (!Code)
     return replyError(ErrorCode::InvalidParams,
                       "onRename called for non-added file");
@@ -254,13 +263,15 @@ void ClangdLSPServer::onRename(RenameParams &Params) {
 }
 
 void ClangdLSPServer::onDocumentDidClose(DidCloseTextDocumentParams &Params) {
-  Server.removeDocument(Params.textDocument.uri.file());
+  PathRef File = Params.textDocument.uri.file();
+  DraftMgr.removeDraft(File);
+  Server.removeDocument(File);
 }
 
 void ClangdLSPServer::onDocumentOnTypeFormatting(
     DocumentOnTypeFormattingParams &Params) {
   auto File = Params.textDocument.uri.file();
-  auto Code = Server.getDocument(File);
+  auto Code = DraftMgr.getDraft(File);
   if (!Code)
     return replyError(ErrorCode::InvalidParams,
                       "onDocumentOnTypeFormatting called for non-added file");
@@ -276,7 +287,7 @@ void ClangdLSPServer::onDocumentOnTypeFormatting(
 void ClangdLSPServer::onDocumentRangeFormatting(
     DocumentRangeFormattingParams &Params) {
   auto File = Params.textDocument.uri.file();
-  auto Code = Server.getDocument(File);
+  auto Code = DraftMgr.getDraft(File);
   if (!Code)
     return replyError(ErrorCode::InvalidParams,
                       "onDocumentRangeFormatting called for non-added file");
@@ -291,7 +302,7 @@ void ClangdLSPServer::onDocumentRangeFormatting(
 
 void ClangdLSPServer::onDocumentFormatting(DocumentFormattingParams &Params) {
   auto File = Params.textDocument.uri.file();
-  auto Code = Server.getDocument(File);
+  auto Code = DraftMgr.getDraft(File);
   if (!Code)
     return replyError(ErrorCode::InvalidParams,
                       "onDocumentFormatting called for non-added file");
@@ -307,7 +318,7 @@ void ClangdLSPServer::onDocumentFormatting(DocumentFormattingParams &Params) {
 void ClangdLSPServer::onCodeAction(CodeActionParams &Params) {
   // We provide a code action for each diagnostic at the requested location
   // which has FixIts available.
-  auto Code = Server.getDocument(Params.textDocument.uri.file());
+  auto Code = DraftMgr.getDraft(Params.textDocument.uri.file());
   if (!Code)
     return replyError(ErrorCode::InvalidParams,
                       "onCodeAction called for non-added file");
@@ -397,7 +408,7 @@ void ClangdLSPServer::onChangeConfiguration(
   // Compilation database change.
   if (Settings.compilationDatabasePath.hasValue()) {
     CDB.setCompileCommandsDir(Settings.compilationDatabasePath.getValue());
-    Server.reparseOpenedFiles();
+    reparseOpenedFiles();
   }
 }
 
@@ -479,3 +490,10 @@ void ClangdLSPServer::onDiagnosticsReady(PathRef File,
        }},
   });
 }
+
+void ClangdLSPServer::reparseOpenedFiles() {
+  for (const Path &FilePath : DraftMgr.getActiveFiles())
+    Server.addDocument(FilePath, *DraftMgr.getDraft(FilePath),
+                       WantDiagnostics::Auto,
+                       /*SkipCache=*/true);
+}
diff --git a/clang-tools-extra/clangd/ClangdLSPServer.h b/clang-tools-extra/clangd/ClangdLSPServer.h
index 706eb4df457..82cf16f2637 100644
--- a/clang-tools-extra/clangd/ClangdLSPServer.h
+++ b/clang-tools-extra/clangd/ClangdLSPServer.h
@@ -11,6 +11,7 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDLSPSERVER_H
 
 #include "ClangdServer.h"
+#include "DraftStore.h"
 #include "GlobalCompilationDatabase.h"
 #include "Path.h"
 #include "Protocol.h"
@@ -74,6 +75,11 @@ private:
 
   std::vector<Fix> getFixes(StringRef File, const clangd::Diagnostic &D);
 
+  /// Forces a reparse of all currently opened files.  As a result, this method
+  /// may be very expensive.  This method is normally called when the
+  /// compilation database is changed.
+  void reparseOpenedFiles();
+
   JSONOutput &Out;
   /// Used to indicate that the 'shutdown' request was received from the
   /// Language Server client.
@@ -96,6 +102,10 @@ private:
   RealFileSystemProvider FSProvider;
   /// Options used for code completion
   clangd::CodeCompleteOptions CCOpts;
+
+  // Store of the current versions of the open documents.
+  DraftStore DraftMgr;
+
   // Server must be the last member of the class to allow its destructor to exit
   // the worker thread that may otherwise run an async callback on partially
   // destructed instance of ClangdLSPServer.
diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index 4e0e9553b39..4d49072986a 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -120,7 +120,7 @@ void ClangdServer::addDocument(PathRef File, StringRef Contents,
   if (SkipCache)
     CompileArgs.invalidate(File);
 
-  DocVersion Version = DraftMgr.updateDraft(File, Contents);
+  DocVersion Version = ++InternalVersion[File];
   ParseInputs Inputs = {CompileArgs.getCompileCommand(File),
                         FSProvider.getFileSystem(), Contents.str()};
 
@@ -132,7 +132,7 @@ void ClangdServer::addDocument(PathRef File, StringRef Contents,
 }
 
 void ClangdServer::removeDocument(PathRef File) {
-  DraftMgr.removeDraft(File);
+  ++InternalVersion[File];
   CompileArgs.invalidate(File);
   WorkScheduler.remove(File);
 }
@@ -145,19 +145,15 @@ void ClangdServer::codeComplete(PathRef File, Position Pos,
   if (!CodeCompleteOpts.Index) // Respect overridden index.
     CodeCompleteOpts.Index = Index;
 
-  VersionedDraft Latest = DraftMgr.getDraft(File);
-  if (!Latest.Draft)
-    return CB(llvm::make_error<llvm::StringError>(
-        "codeComplete called for non-added document",
-        llvm::errc::invalid_argument));
-
   // Copy PCHs to avoid accessing this->PCHs concurrently
   std::shared_ptr<PCHContainerOperations> PCHs = this->PCHs;
   auto FS = FSProvider.getFileSystem();
   auto Task = [PCHs, Pos, FS,
                CodeCompleteOpts](Path File, Callback<CompletionList> CB,
                                  llvm::Expected<InputsAndPreamble> IP) {
-    assert(IP && "error when trying to read preamble for codeComplete");
+    if (!IP)
+      return CB(IP.takeError());
+
     auto PreambleData = IP->Preamble;
 
     // FIXME(ibiryukov): even if Preamble is non-null, we may want to check
@@ -174,11 +170,6 @@ void ClangdServer::codeComplete(PathRef File, Position Pos,
 
 void ClangdServer::signatureHelp(PathRef File, Position Pos,
                                  Callback<SignatureHelp> CB) {
-  VersionedDraft Latest = DraftMgr.getDraft(File);
-  if (!Latest.Draft)
-    return CB(llvm::make_error<llvm::StringError>(
-        "signatureHelp is called for non-added document",
-        llvm::errc::invalid_argument));
 
   auto PCHs = this->PCHs;
   auto FS = FSProvider.getFileSystem();
@@ -333,13 +324,6 @@ ClangdServer::insertInclude(PathRef File, StringRef Code,
   return formatReplacements(Code, *Replaces, *Style);
 }
 
-llvm::Optional<std::string> ClangdServer::getDocument(PathRef File) {
-  auto Latest = DraftMgr.getDraft(File);
-  if (!Latest.Draft)
-    return llvm::None;
-  return std::move(*Latest.Draft);
-}
-
 void ClangdServer::dumpAST(PathRef File,
                            UniqueFunction<void(std::string)> Callback) {
   auto Action = [](decltype(Callback) Callback,
@@ -455,11 +439,6 @@ ClangdServer::formatCode(llvm::StringRef Code, PathRef File,
 
 void ClangdServer::findDocumentHighlights(
     PathRef File, Position Pos, Callback<std::vector<DocumentHighlight>> CB) {
-  auto FileContents = DraftMgr.getDraft(File);
-  if (!FileContents.Draft)
-    return CB(llvm::make_error<llvm::StringError>(
-        "findDocumentHighlights called on non-added file",
-        llvm::errc::invalid_argument));
 
   auto FS = FSProvider.getFileSystem();
   auto Action = [FS, Pos](Callback<std::vector<DocumentHighlight>> CB,
@@ -473,12 +452,6 @@ void ClangdServer::findDocumentHighlights(
 }
 
 void ClangdServer::findHover(PathRef File, Position Pos, Callback<Hover> CB) {
-  Hover FinalHover;
-  auto FileContents = DraftMgr.getDraft(File);
-  if (!FileContents.Draft)
-    return CB(llvm::make_error<llvm::StringError>(
-        "findHover called on non-added file", llvm::errc::invalid_argument));
-
   auto FS = FSProvider.getFileSystem();
   auto Action = [Pos, FS](Callback<Hover> CB,
                           llvm::Expected<InputsAndAST> InpAST) {
@@ -508,12 +481,6 @@ void ClangdServer::consumeDiagnostics(PathRef File, DocVersion Version,
   DiagConsumer.onDiagnosticsReady(File, std::move(Diags));
 }
 
-void ClangdServer::reparseOpenedFiles() {
-  for (const Path &FilePath : DraftMgr.getActiveFiles())
-    addDocument(FilePath, *DraftMgr.getDraft(FilePath).Draft,
-                WantDiagnostics::Auto, /*SkipCache=*/true);
-}
-
 void ClangdServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
   // FIXME: Do nothing for now. This will be used for indexing and potentially
   // invalidating other caches.
diff --git a/clang-tools-extra/clangd/ClangdServer.h b/clang-tools-extra/clangd/ClangdServer.h
index 1dd61739f4d..34af4b9e29a 100644
--- a/clang-tools-extra/clangd/ClangdServer.h
+++ b/clang-tools-extra/clangd/ClangdServer.h
@@ -13,7 +13,6 @@
 #include "ClangdUnit.h"
 #include "CodeComplete.h"
 #include "CompileArgsCache.h"
-#include "DraftStore.h"
 #include "Function.h"
 #include "GlobalCompilationDatabase.h"
 #include "Protocol.h"
@@ -127,18 +126,9 @@ public:
   /// resources associated with it.
   void removeDocument(PathRef File);
 
-  /// Calls forceReparse() on all currently opened files.
-  /// As a result, this method may be very expensive.
-  /// This method is normally called when the compilation database is changed.
-  /// FIXME: this method must be moved to ClangdLSPServer along with DraftMgr.
-  void reparseOpenedFiles();
-
   /// Run code completion for \p File at \p Pos.
   /// Request is processed asynchronously.
   ///
-  /// The current draft for \p File will be used. If \p UsedFS is non-null, it
-  /// will be overwritten by vfs::FileSystem used for completion.
-  ///
   /// This method should only be called for currently tracked files. However, it
   /// is safe to call removeDocument for \p File after this method returns, even
   /// while returned future is not yet ready.
@@ -148,11 +138,8 @@ public:
                     const clangd::CodeCompleteOptions &Opts,
                     Callback<CompletionList> CB);
 
-  /// Provide signature help for \p File at \p Pos. If \p OverridenContents is
-  /// not None, they will used only for signature help, i.e. no diagnostics
-  /// update will be scheduled and a draft for \p File will not be updated. If
-  /// If \p UsedFS is non-null, it will be overwritten by vfs::FileSystem used
-  /// for signature help. This method should only be called for tracked files.
+  /// Provide signature help for \p File at \p Pos.  This method should only be
+  /// called for tracked files.
   void signatureHelp(PathRef File, Position Pos, Callback<SignatureHelp> CB);
 
   /// Get definition of symbol at a specified \p Line and \p Column in \p File.
@@ -204,12 +191,6 @@ public:
                                                 StringRef DeclaringHeader,
                                                 StringRef InsertedHeader);
 
-  /// Gets current document contents for \p File. Returns None if \p File is not
-  /// currently tracked.
-  /// FIXME(ibiryukov): This function is here to allow offset-to-Position
-  /// conversions in outside code, maybe there's a way to get rid of it.
-  llvm::Optional<std::string> getDocument(PathRef File);
-
   /// Only for testing purposes.
   /// Waits until all requests to worker thread are finished and dumps AST for
   /// \p File. \p File must be in the list of added documents.
@@ -238,13 +219,18 @@ private:
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
+  typedef uint64_t DocVersion;
+
   void consumeDiagnostics(PathRef File, DocVersion Version,
                           std::vector<Diag> Diags);
 
   CompileArgsCache CompileArgs;
   DiagnosticsConsumer &DiagConsumer;
   FileSystemProvider &FSProvider;
-  DraftStore DraftMgr;
+
+  /// Used to synchronize diagnostic responses for added and removed files.
+  llvm::StringMap<DocVersion> InternalVersion;
+
   // The index used to look up symbols. This could be:
   //   - null (all index functionality is optional)
   //   - the dynamic index owned by ClangdServer (FileIdx)
diff --git a/clang-tools-extra/clangd/DraftStore.cpp b/clang-tools-extra/clangd/DraftStore.cpp
index 4bda859aa73..e3403abc6ad 100644
--- a/clang-tools-extra/clangd/DraftStore.cpp
+++ b/clang-tools-extra/clangd/DraftStore.cpp
@@ -12,12 +12,13 @@
 using namespace clang;
 using namespace clang::clangd;
 
-VersionedDraft DraftStore::getDraft(PathRef File) const {
+llvm::Optional<std::string> 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 llvm::None;
+
   return It->second;
 }
 
@@ -26,35 +27,19 @@ std::vector<Path> DraftStore::getActiveFiles() const {
   std::vector<Path> ResultVector;
 
   for (auto DraftIt = Drafts.begin(); DraftIt != Drafts.end(); DraftIt++)
-    if (DraftIt->second.Draft)
-      ResultVector.push_back(DraftIt->getKey());
+    ResultVector.push_back(DraftIt->getKey());
 
   return ResultVector;
 }
 
-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) {
+void 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;
+  Drafts[File] = Contents;
 }
 
-DocVersion DraftStore::removeDraft(PathRef File) {
+void 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;
+  Drafts.erase(File);
 }
diff --git a/clang-tools-extra/clangd/DraftStore.h b/clang-tools-extra/clangd/DraftStore.h
index 4757768a377..c84b0c1391a 100644
--- a/clang-tools-extra/clangd/DraftStore.h
+++ b/clang-tools-extra/clangd/DraftStore.h
@@ -13,7 +13,6 @@
 #include "Path.h"
 #include "clang/Basic/LLVM.h"
 #include "llvm/ADT/StringMap.h"
-#include <cstdint>
 #include <mutex>
 #include <string>
 #include <vector>
@@ -21,45 +20,26 @@
 namespace clang {
 namespace clangd {
 
-/// Using unsigned int type here to avoid undefined behaviour on overflow.
-typedef uint64_t 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 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.
+/// filenames. The contents are owned by the DraftStore.
 class DraftStore {
 public:
-  /// \return version and contents of the stored document.
-  /// For untracked files, a (0, None) pair is returned.
-  VersionedDraft getDraft(PathRef File) const;
+  /// \return Contents of the stored document.
+  /// For untracked files, a llvm::None is returned.
+  llvm::Optional<std::string> getDraft(PathRef File) const;
 
-  /// \return List of names of active drafts in this store.  Drafts that were
-  /// removed (which still have an entry in the Drafts map) are not returned by
-  /// this function.
+  /// \return List of names of the drafts in this store.
   std::vector<Path> getActiveFiles() 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);
+  void updateDraft(PathRef File, StringRef Contents);
+
+  /// Remove the draft from the store.
+  void removeDraft(PathRef File);
 
 private:
   mutable std::mutex Mutex;
-  llvm::StringMap<VersionedDraft> Drafts;
+  llvm::StringMap<std::string> Drafts;
 };
 
 } // namespace clangd
diff --git a/clang-tools-extra/unittests/clangd/ClangdTests.cpp b/clang-tools-extra/unittests/clangd/ClangdTests.cpp
index 3db20753a29..13671dd32fe 100644
--- a/clang-tools-extra/unittests/clangd/ClangdTests.cpp
+++ b/clang-tools-extra/unittests/clangd/ClangdTests.cpp
@@ -478,7 +478,10 @@ int hello;
   CDB.ExtraClangFlags.clear();
   DiagConsumer.clear();
   Server.removeDocument(BazCpp);
-  Server.reparseOpenedFiles();
+  Server.addDocument(FooCpp, FooSource.code(), WantDiagnostics::Auto,
+                     /*SkipCache=*/true);
+  Server.addDocument(BarCpp, BarSource.code(), WantDiagnostics::Auto,
+                     /*SkipCache=*/true);
   ASSERT_TRUE(Server.blockUntilIdleForTest());
 
   EXPECT_THAT(DiagConsumer.filesWithDiags(),