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

[ilya-biryukov]

Summary

This is a reapplied r305280 with a fix to the crash found by build bots
(StringRef to an out-of-scope local std::string).

[ilya-biryukov]

clangd/ClangdServer.cpp

190

This is the relevant line that changed from the previous review.

clangd/ClangdServer.h

public:
   /// 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);
   /// 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);
   /// Run formatting for the whole \p File.
   std::vector<tooling::Replacement> formatFile(PathRef File);
   /// Run formatting after a character was typed at \p Pos in \p File.
   std::vector<tooling::Replacement> formatOnType(PathRef File, Position Pos);
 

clangd/ClangdServer.cpp

 }
 
 void ClangdServer::forceReparse(PathRef File) {
   // The addDocument schedules the reparse even if the contents of the file
   // never changed, so we just call it here.
   addDocument(File, getDocument(File));
 }
 
-Tagged<std::vector<CompletionItem>> ClangdServer::codeComplete(PathRef File,
-                                                               Position Pos) {
+Tagged<std::vector<CompletionItem>>
+ClangdServer::codeComplete(PathRef File, Position Pos,
+                           llvm::Optional<StringRef> OverridenContents) {
+  std::string DraftStorage;
+  if (!OverridenContents) {
     auto FileContents = DraftMgr.getDraft(File);
-  assert(FileContents.Draft && "codeComplete is called for non-added document");
+    assert(FileContents.Draft &&
+           "codeComplete is called for non-added document");
+
+    DraftStorage = std::move(*FileContents.Draft);
+    OverridenContents = DraftStorage;
+  }
 
   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);
 }
 
 std::vector<tooling::Replacement> ClangdServer::formatRange(PathRef File,
                                                             Range Rng) {
   std::string Code = getDocument(File);
 
   size_t Begin = positionToOffset(Code, Rng.start);

unittests/clangd/ClangdTests.cpp

TEST_F(ClangdVFSTest, CheckVersions) {
   EXPECT_EQ(Server.codeComplete(FooCpp, Position{0, 0}).Tag, FS->Tag);
 
   FS->Tag = "321";
   Server.addDocument(FooCpp, SourceContents);
   EXPECT_EQ(DiagConsumer->lastVFSTag(), FS->Tag);
   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