[clangd] Allow to override contents of the file during completion.

Reviewers: krasimir

Reviewed By: krasimir

Subscribers: klimek, cfe-commits

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

llvm-svn: 305280

3 files changed, 91 insertions, 10 deletions

diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index 459b6a64217..ee845e6e1df 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -184,17 +184,27 @@ void ClangdServer::forceReparse(PathRef File) {
   addDocument(File, getDocument(File));
 }
 
-Tagged<std::vector<CompletionItem>> ClangdServer::codeComplete(PathRef File,
-                                                               Position Pos) {
-  auto FileContents = DraftMgr.getDraft(File);
-  assert(FileContents.Draft && "codeComplete is called for non-added document");
+Tagged<std::vector<CompletionItem>>
+ClangdServer::codeComplete(PathRef File, Position Pos,
+                           llvm::Optional<StringRef> OverridenContents) {
+  if (!OverridenContents) {
+    auto FileContents = DraftMgr.getDraft(File);
+    assert(FileContents.Draft &&
+           "codeComplete is called for non-added document");
+
+    OverridenContents = *FileContents.Draft;
+  }
 
   std::vector<CompletionItem> Result;
   auto TaggedFS = FSProvider->getTaggedFileSystem();
-  Units.runOnUnitWithoutReparse(
-      File, *FileContents.Draft, *CDB, PCHs, TaggedFS.Value, [&](ClangdUnit &Unit) {
-        Result = Unit.codeComplete(*FileContents.Draft, Pos, TaggedFS.Value);
-      });
+  // It would be nice to use runOnUnitWithoutReparse here, but we can't
+  // guarantee the correctness of code completion cache here if we don't do the
+  // reparse.
+  Units.runOnUnit(File, *OverridenContents, *CDB, PCHs, TaggedFS.Value,
+                  [&](ClangdUnit &Unit) {
+                    Result = Unit.codeComplete(*OverridenContents, Pos,
+                                               TaggedFS.Value);
+                  });
   return make_tagged(std::move(Result), TaggedFS.Tag);
 }
 
diff --git a/clang-tools-extra/clangd/ClangdServer.h b/clang-tools-extra/clangd/ClangdServer.h
index 2b6a3ff92ff..73e45f8800f 100644
--- a/clang-tools-extra/clangd/ClangdServer.h
+++ b/clang-tools-extra/clangd/ClangdServer.h
@@ -152,8 +152,15 @@ public:
   /// Force \p File to be reparsed using the latest contents.
   void forceReparse(PathRef File);
 
-  /// Run code completion for \p File at \p Pos.
-  Tagged<std::vector<CompletionItem>> codeComplete(PathRef File, Position Pos);
+  /// Run code completion for \p File at \p Pos. If \p OverridenContents is not
+  /// None, they will used only for code completion, i.e. no diagnostics update
+  /// will be scheduled and a draft for \p File will not be updated.
+  /// If \p OverridenContents is None, contents of the current draft for \p File
+  /// will be used.
+  /// This method should only be called for currently tracked files.
+  Tagged<std::vector<CompletionItem>>
+  codeComplete(PathRef File, Position Pos,
+               llvm::Optional<StringRef> OverridenContents = llvm::None);
 
   /// Run formatting for \p Rng inside \p File.
   std::vector<tooling::Replacement> formatRange(PathRef File, Range Rng);
diff --git a/clang-tools-extra/unittests/clangd/ClangdTests.cpp b/clang-tools-extra/unittests/clangd/ClangdTests.cpp
index ef24e64294d..3a60fe518b6 100644
--- a/clang-tools-extra/unittests/clangd/ClangdTests.cpp
+++ b/clang-tools-extra/unittests/clangd/ClangdTests.cpp
@@ -398,5 +398,69 @@ TEST_F(ClangdVFSTest, CheckVersions) {
   EXPECT_EQ(Server.codeComplete(FooCpp, Position{0, 0}).Tag, FS->Tag);
 }
 
+class ClangdCompletionTest : public ClangdVFSTest {
+protected:
+  bool ContainsItem(std::vector<CompletionItem> const &Items, StringRef Name) {
+    for (const auto &Item : Items) {
+      if (Item.insertText == Name)
+        return true;
+    }
+    return false;
+  }
+};
+
+TEST_F(ClangdCompletionTest, CheckContentsOverride) {
+  MockFSProvider *FS;
+
+  ClangdServer Server(llvm::make_unique<MockCompilationDatabase>(),
+                      llvm::make_unique<ErrorCheckingDiagConsumer>(),
+                      getAndMove(llvm::make_unique<MockFSProvider>(), FS),
+                      /*RunSynchronously=*/false);
+
+  auto FooCpp = getVirtualTestFilePath("foo.cpp");
+  const auto SourceContents = R"cpp(
+int aba;
+int b =   ;
+)cpp";
+
+  const auto OverridenSourceContents = R"cpp(
+int cbc;
+int b =   ;
+)cpp";
+  // Complete after '=' sign. We need to be careful to keep the SourceContents'
+  // size the same.
+  // We complete on the 3rd line (2nd in zero-based numbering), because raw
+  // string literal of the SourceContents starts with a newline(it's easy to
+  // miss).
+  Position CompletePos = {2, 8};
+  FS->Files[FooCpp] = SourceContents;
+
+  Server.addDocument(FooCpp, SourceContents);
+
+  {
+    auto CodeCompletionResults1 =
+        Server.codeComplete(FooCpp, CompletePos, None).Value;
+    EXPECT_TRUE(ContainsItem(CodeCompletionResults1, "aba"));
+    EXPECT_FALSE(ContainsItem(CodeCompletionResults1, "cbc"));
+  }
+
+  {
+    auto CodeCompletionResultsOverriden =
+        Server
+            .codeComplete(FooCpp, CompletePos,
+                          StringRef(OverridenSourceContents))
+            .Value;
+    EXPECT_TRUE(ContainsItem(CodeCompletionResultsOverriden, "cbc"));
+    EXPECT_FALSE(ContainsItem(CodeCompletionResultsOverriden, "aba"));
+  }
+
+  {
+    auto CodeCompletionResults2 =
+        Server.codeComplete(FooCpp, CompletePos, None).Value;
+    EXPECT_TRUE(ContainsItem(CodeCompletionResults2, "aba"));
+    EXPECT_FALSE(ContainsItem(CodeCompletionResults2, "cbc"));
+  }
+}
+
 } // namespace clangd
 } // namespace clang