[clangd] Use Dirty Filesystem for cross file rename.

[njames93]

Summary

Refactor cross file rename to use a Filesystem instead of a function for getting buffer contents of open files.

Depends on D94554

[njames93]

Rebase and fix up tests after changes to rename.

[njames93]

Rebase and remove an addressed FIXME.

[sammccall]

This looks pretty good, great cleanup!

clang-tools-extra/clangd/ClangdServer.cpp

495

we shouldn't need to pass a nontrivial FS in here, right?

clang-tools-extra/clangd/refactor/Rename.h

40

Index + FS are both only used for cross-file renames at this point.

(We used to support a single-file rename mode where the index was used to validate that the symbol wasn't used elsewhere, but that's gone now)

I think we should document this, and assert that they're either both set (cross-file) or neither is.

[njames93]

clang-tools-extra/clangd/refactor/Rename.h

40

I wasn't sure the best way to approach this, but i agree, this is probably the safest.

[kadircet]

clang-tools-extra/clangd/ClangdServer.cpp

495

i think this makes sense, as the idea is try to keep it fast by shying away from copying over all the dirty buffers.

[njames93]

Update to assert that either Index and FS are both set, or neither set.
Document this behaviour.

clang-tools-extra/clangd/ClangdServer.h

llvm::StringMap<llvm::Optional<FuzzyFindRequest>>
       CachedCompletionFuzzyFindRequestByFile;
   mutable std::mutex CachedCompletionFuzzyFindRequestMutex;
 
   llvm::Optional<std::string> WorkspaceRoot;
   llvm::Optional<TUScheduler> WorkScheduler;
 
   // Store of the current versions of the open documents.
   // Only written from the main thread (despite being threadsafe).
-  // FIXME: TUScheduler also keeps these, unify?
   DraftStore DraftMgr;
 
   std::unique_ptr<ThreadsafeFS> DirtyFS;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/clangd/ClangdServer.cpp

void ClangdServer::prepareRename(PathRef File, Position Pos,
                                  Callback<RenameResult> CB) {
   auto Action = [Pos, File = File.str(), CB = std::move(CB),
                  NewName = std::move(NewName),
                  RenameOpts](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     // prepareRename is latency-sensitive: we don't query the index, as we
     // only need main-file references
-    auto Results = clangd::rename(
-        {Pos, NewName.getValueOr("__clangd_rename_dummy"), InpAST->AST, File,
+    auto Results =
+        clangd::rename({Pos, NewName.getValueOr("__clangd_rename_dummy"),
+                        InpAST->AST, File, /*FS=*/nullptr,
                         /*Index=*/nullptr, RenameOpts});
     if (!Results) {
       // LSP says to return null on failure, but that will result in a generic
       // failure message. If we send an LSP error response, clients can surface
       // the message to users (VSCode does).
       return CB(Results.takeError());
     }
     return CB(*Results);
   };
   WorkScheduler->runWithAST("PrepareRename", File, std::move(Action));
 }
 
 void ClangdServer::rename(PathRef File, Position Pos, llvm::StringRef NewName,
                           const RenameOptions &Opts,
                           Callback<RenameResult> CB) {
-  // A snapshot of all file dirty buffers.
-  llvm::StringMap<std::string> Snapshot = WorkScheduler->getAllFileContents();
   auto Action = [File = File.str(), NewName = NewName.str(), Pos, Opts,
-                 CB = std::move(CB), Snapshot = std::move(Snapshot),
+                 CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     // Tracks number of files edited per invocation.
     static constexpr trace::Metric RenameFiles("rename_files",
                                                trace::Metric::Distribution);
     if (!InpAST)
       return CB(InpAST.takeError());
-    auto GetDirtyBuffer =
-        [&Snapshot](PathRef AbsPath) -> llvm::Optional<std::string> {
-      auto It = Snapshot.find(AbsPath);
-      if (It == Snapshot.end())
-        return llvm::None;
-      return It->second;
-    };
-    auto R = clangd::rename(
-        {Pos, NewName, InpAST->AST, File, Index, Opts, GetDirtyBuffer});
+    auto R = clangd::rename({Pos, NewName, InpAST->AST, File,
+                             DirtyFS->view(llvm::None), Index, Opts});
     if (!R)
       return CB(R.takeError());
 
     if (Opts.WantFormat) {
       auto Style = getFormatStyleForFile(File, InpAST->Inputs.Contents,
                                          *InpAST->Inputs.TFS);
       llvm::Error Err = llvm::Error::success();
       for (auto &E : R->GlobalChanges)

clang-tools-extra/clangd/TUScheduler.h

public:
   /// Returns true if the file was not previously tracked.
   bool update(PathRef File, ParseInputs Inputs, WantDiagnostics WD);
 
   /// Remove \p File from the list of tracked files and schedule removal of its
   /// resources. Pending diagnostics for closed files may not be delivered, even
   /// if requested with WantDiags::Auto or WantDiags::Yes.
   void remove(PathRef File);
 
-  /// Returns a snapshot of all file buffer contents, per last update().
-  llvm::StringMap<std::string> getAllFileContents() const;
-
   /// Schedule an async task with no dependencies.
   /// Path may be empty (it is used only to set the Context).
   void run(llvm::StringRef Name, llvm::StringRef Path,
            llvm::unique_function<void()> Action);
 
   /// Similar to run, except the task is expected to be quick.
   /// This function will not honor AsyncThreadsCount (except
   /// if threading is disabled with AsyncThreadsCount=0)

clang-tools-extra/clangd/TUScheduler.cpp

if (!Removed)
     elog("Trying to remove file from TUScheduler that is not tracked: {0}",
          File);
   // We don't call HeaderIncluders.remove(File) here.
   // If we did, we'd avoid potentially stale header/mainfile associations.
   // However, it would mean that closing a mainfile could invalidate the
   // preamble of several open headers.
 }
 
-llvm::StringMap<std::string> TUScheduler::getAllFileContents() const {
-  llvm::StringMap<std::string> Results;
-  for (auto &It : Files)
-    Results.try_emplace(It.getKey(), It.getValue()->Contents);
-  return Results;
-}
-
 void TUScheduler::run(llvm::StringRef Name, llvm::StringRef Path,
                       llvm::unique_function<void()> Action) {
   runWithSemaphore(Name, Path, std::move(Action), Barrier);
 }
 
 void TUScheduler::runQuick(llvm::StringRef Name, llvm::StringRef Path,
                            llvm::unique_function<void()> Action) {
   // Use QuickRunBarrier to serialize quick tasks: we are ignoring

clang-tools-extra/clangd/refactor/Rename.h

 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/Support/Error.h"
 
 namespace clang {
 namespace clangd {
 class ParsedAST;
 class SymbolIndex;
 
-/// Gets dirty buffer for a given file \p AbsPath.
-/// Returns None if there is no dirty buffer for the given file.
-using DirtyBufferGetter =
-    llvm::function_ref<llvm::Optional<std::string>(PathRef AbsPath)>;
-
 struct RenameOptions {
   /// The maximum number of affected files (0 means no limit), only meaningful
   /// when AllowCrossFile = true.
   /// If the actual number exceeds the limit, rename is forbidden.
   size_t LimitFiles = 50;
   /// If true, format the rename edits, only meaningful in ClangdServer layer.
   bool WantFormat = false;
 };
 
 struct RenameInputs {
   Position Pos; // the position triggering the rename
   llvm::StringRef NewName;
 
   ParsedAST &AST;
   llvm::StringRef MainFilePath;
 
+  // The filesystem to query when performing cross file renames.
+  // If this is set, Index must also be set, likewise if this is nullptr, Index
+  // must also be nullptr.
+  llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FS = nullptr;
+
   const SymbolIndex *Index = nullptr;
 
   RenameOptions Opts = {};
-  // When set, used by the rename to get file content for all rename-related
-  // files.
-  // If there is no corresponding dirty buffer, we will use the file content
-  // from disk.
-  DirtyBufferGetter GetDirtyBuffer = nullptr;
 };
 
 struct RenameResult {
   // The range of the symbol that the user can attempt to rename.
   Range Target;
   // Rename occurrences for the current main file.
   std::vector<Range> LocalChanges;
   // Complete edits for the rename, including LocalChanges.

clang-tools-extra/clangd/refactor/Rename.cpp

 // We choose to trade off some correctness for performance and scalability.
 //
 // Clangd builds a dynamic index for all opened files on top of the static
 // index of the whole codebase. Dynamic index is up-to-date (respects dirty
 // buffers) as long as clangd finishes processing opened files, while static
 // index (background index) is relatively stale. We choose the dirty buffers
 // as the file content we rename on, and fallback to file content on disk if
 // there is no dirty buffer.
-llvm::Expected<FileEdits> renameOutsideFile(
-    const NamedDecl &RenameDecl, llvm::StringRef MainFilePath,
-    llvm::StringRef NewName, const SymbolIndex &Index, size_t MaxLimitFiles,
-    llvm::function_ref<llvm::Expected<std::string>(PathRef)> GetFileContent) {
+llvm::Expected<FileEdits>
+renameOutsideFile(const NamedDecl &RenameDecl, llvm::StringRef MainFilePath,
+                  llvm::StringRef NewName, const SymbolIndex &Index,
+                  size_t MaxLimitFiles, llvm::vfs::FileSystem &FS) {
   trace::Span Tracer("RenameOutsideFile");
   auto AffectedFiles = findOccurrencesOutsideFile(RenameDecl, MainFilePath,
                                                   Index, MaxLimitFiles);
   if (!AffectedFiles)
     return AffectedFiles.takeError();
   FileEdits Results;
   for (auto &FileAndOccurrences : *AffectedFiles) {
     llvm::StringRef FilePath = FileAndOccurrences.first();
 
-    auto AffectedFileCode = GetFileContent(FilePath);
-    if (!AffectedFileCode) {
-      elog("Fail to read file content: {0}", AffectedFileCode.takeError());
+    auto ExpBuffer = FS.getBufferForFile(FilePath);
+    if (!ExpBuffer) {
+      elog("Fail to read file content: Fail to open file {0}: {1}", FilePath,
+           ExpBuffer.getError().message());
       continue;
     }
+
+    auto AffectedFileCode = (*ExpBuffer)->getBuffer();
     auto RenameRanges =
-        adjustRenameRanges(*AffectedFileCode, RenameDecl.getNameAsString(),
+        adjustRenameRanges(AffectedFileCode, RenameDecl.getNameAsString(),
                            std::move(FileAndOccurrences.second),
                            RenameDecl.getASTContext().getLangOpts());
     if (!RenameRanges) {
       // Our heuristics fails to adjust rename ranges to the current state of
       // the file, it is most likely the index is stale, so we give up the
       // entire rename.
       return error("Index results don't match the content of file {0} "
                    "(the index may be stale)",
                    FilePath);
     }
     auto RenameEdit =
-        buildRenameEdit(FilePath, *AffectedFileCode, *RenameRanges, NewName);
+        buildRenameEdit(FilePath, AffectedFileCode, *RenameRanges, NewName);
     if (!RenameEdit)
       return error("failed to rename in file {0}: {1}", FilePath,
                    RenameEdit.takeError());
     if (!RenameEdit->Replacements.empty())
       Results.insert({FilePath, std::move(*RenameEdit)});
   }
   return Results;
 }
void findNearMiss(
   }
   findNearMiss(PartialMatch, IndexedRest, LexedRest.drop_front(),
                LexedIndex + 1, Fuel, MatchedCB);
 }
 
 } // namespace
 
 llvm::Expected<RenameResult> rename(const RenameInputs &RInputs) {
+  assert(!RInputs.Index == !RInputs.FS &&
+         "Index and FS must either both be specified or both null.");
   trace::Span Tracer("Rename flow");
   const auto &Opts = RInputs.Opts;
   ParsedAST &AST = RInputs.AST;
   const SourceManager &SM = AST.getSourceManager();
   llvm::StringRef MainFileCode = SM.getBufferData(SM.getMainFileID());
-  auto GetFileContent = [&RInputs,
-                         &SM](PathRef AbsPath) -> llvm::Expected<std::string> {
-    llvm::Optional<std::string> DirtyBuffer;
-    if (RInputs.GetDirtyBuffer &&
-        (DirtyBuffer = RInputs.GetDirtyBuffer(AbsPath)))
-      return std::move(*DirtyBuffer);
-
-    auto Content =
-        SM.getFileManager().getVirtualFileSystem().getBufferForFile(AbsPath);
-    if (!Content)
-      return error("Fail to open file {0}: {1}", AbsPath,
-                   Content.getError().message());
-    if (!*Content)
-      return error("Got no buffer for file {0}", AbsPath);
-
-    return (*Content)->getBuffer().str();
-  };
   // Try to find the tokens adjacent to the cursor position.
   auto Loc = sourceLocationInMainFile(SM, RInputs.Pos);
   if (!Loc)
     return Loc.takeError();
   const syntax::Token *IdentifierToken =
       spelledIdentifierTouching(*Loc, AST.getTokens());
 
   // Renames should only triggered on identifiers.
if (!RInputs.Index) {
     assert(Result.GlobalChanges.empty());
     return Result;
   }
 
   auto OtherFilesEdits = renameOutsideFile(
       RenameDecl, RInputs.MainFilePath, RInputs.NewName, *RInputs.Index,
       Opts.LimitFiles == 0 ? std::numeric_limits<size_t>::max()
                            : Opts.LimitFiles,
-      GetFileContent);
+      *RInputs.FS);
   if (!OtherFilesEdits)
     return OtherFilesEdits.takeError();
   Result.GlobalChanges = *OtherFilesEdits;
   // Attach the rename edits for the main file.
   Result.GlobalChanges.try_emplace(RInputs.MainFilePath,
                                    std::move(MainFileEdits));
   return Result;
 }

clang-tools-extra/clangd/unittests/RenameTests.cpp

 using testing::ElementsAre;
 using testing::Eq;
 using testing::IsEmpty;
 using testing::Pair;
 using testing::SizeIs;
 using testing::UnorderedElementsAre;
 using testing::UnorderedElementsAreArray;
 
+llvm::IntrusiveRefCntPtr<llvm::vfs::OverlayFileSystem>
+createOverlay(llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> Base,
+              llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> Overlay) {
+  auto OFS =
+      llvm::makeIntrusiveRefCnt<llvm::vfs::OverlayFileSystem>(std::move(Base));
+  OFS->pushOverlay(std::move(Overlay));
+  return OFS;
+}
+
+llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> getVFSFromAST(ParsedAST &AST) {
+  return &AST.getSourceManager().getFileManager().getVirtualFileSystem();
+}
+
 // Convert a Range to a Ref.
 Ref refWithRange(const clangd::Range &Range, const std::string &URI) {
   Ref Result;
   Result.Kind = RefKind::Reference | RefKind::Spelled;
   Result.Location.Start.setLine(Range.start.line);
   Result.Location.Start.setColumn(Range.start.character);
   Result.Location.End.setLine(Range.end.line);
   Result.Location.End.setColumn(Range.end.character);
for (llvm::StringRef T : Tests) {
     SCOPED_TRACE(T);
     Annotations Code(T);
     auto TU = TestTU::withCode(Code.code());
     TU.ExtraArgs.push_back("-xobjective-c++");
     auto AST = TU.build();
     auto Index = TU.index();
     for (const auto &RenamePos : Code.points()) {
       auto RenameResult =
-          rename({RenamePos, NewName, AST, testPath(TU.Filename), Index.get()});
+          rename({RenamePos, NewName, AST, testPath(TU.Filename),
+                  getVFSFromAST(AST), Index.get()});
       ASSERT_TRUE(bool(RenameResult)) << RenameResult.takeError();
       ASSERT_EQ(1u, RenameResult->GlobalChanges.size());
       EXPECT_EQ(
           applyEdits(std::move(RenameResult->GlobalChanges)).front().second,
           expectedResult(Code, NewName));
     }
   }
 }

 
 TEST(RenameTest, IndexMergeMainFile) {
   Annotations Code("int ^x();");
   TestTU TU = TestTU::withCode(Code.code());
   TU.Filename = "main.cc";
   auto AST = TU.build();
 
   auto Main = testPath("main.cc");
+  auto InMemFS = llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
+  InMemFS->addFile(testPath("main.cc"), 0,
+                   llvm::MemoryBuffer::getMemBuffer(Code.code()));
+  InMemFS->addFile(testPath("other.cc"), 0,
+                   llvm::MemoryBuffer::getMemBuffer(Code.code()));
 
   auto Rename = [&](const SymbolIndex *Idx) {
-    auto GetDirtyBuffer = [&](PathRef Path) -> llvm::Optional<std::string> {
-      return Code.code().str(); // Every file has the same content.
-    };
-    RenameInputs Inputs{Code.point(), "xPrime",        AST,           Main,
-                        Idx,          RenameOptions(), GetDirtyBuffer};
+    RenameInputs Inputs{Code.point(),
+                        "xPrime",
+                        AST,
+                        Main,
+                        Idx ? createOverlay(getVFSFromAST(AST), InMemFS)
+                            : nullptr,
+                        Idx,
+                        RenameOptions()};
     auto Results = rename(Inputs);
     EXPECT_TRUE(bool(Results)) << llvm::toString(Results.takeError());
     return std::move(*Results);
   };
 
   // We do not expect to see duplicated edits from AST vs index.
   auto Results = Rename(TU.index().get());
   EXPECT_THAT(Results.GlobalChanges.keys(), ElementsAre(Main));
TEST(CrossFileRenameTests, DirtyBuffer) {
   FSymbols.update(FooPath, nullptr, buildRefSlab(FooCode, "Foo", FooPath),
                   nullptr, false);
   FSymbols.update(BarPath, nullptr, buildRefSlab(BarCode, "Bar", BarPath),
                   nullptr, false);
   auto Index = FSymbols.buildIndex(IndexType::Light);
 
   Annotations MainCode("class  [[Fo^o]] {};");
   auto MainFilePath = testPath("main.cc");
-  // Dirty buffer for foo.cc.
-  auto GetDirtyBuffer = [&](PathRef Path) -> llvm::Optional<std::string> {
-    if (Path == FooPath)
-      return FooDirtyBuffer.code().str();
-    return llvm::None;
-  };
+  llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> InMemFS =
+      new llvm::vfs::InMemoryFileSystem;
+  InMemFS->addFile(FooPath, 0,
+                   llvm::MemoryBuffer::getMemBuffer(FooDirtyBuffer.code()));
 
   // Run rename on Foo, there is a dirty buffer for foo.cc, rename should
   // respect the dirty buffer.
   TestTU TU = TestTU::withCode(MainCode.code());
   auto AST = TU.build();
   llvm::StringRef NewName = "newName";
-  auto Results = rename({MainCode.point(),
-                         NewName,
-                         AST,
-                         MainFilePath,
-                         Index.get(),
-                         {},
-                         GetDirtyBuffer});
+  auto Results =
+      rename({MainCode.point(), NewName, AST, MainFilePath,
+              createOverlay(getVFSFromAST(AST), InMemFS), Index.get()});
   ASSERT_TRUE(bool(Results)) << Results.takeError();
   EXPECT_THAT(
       applyEdits(std::move(Results->GlobalChanges)),
       UnorderedElementsAre(
           Pair(Eq(FooPath), Eq(expectedResult(FooDirtyBuffer, NewName))),
           Pair(Eq(MainFilePath), Eq(expectedResult(MainCode, NewName)))));
 
   // Run rename on Bar, there is no dirty buffer for the affected file bar.cc,
   // so we should read file content from VFS.
   MainCode = Annotations("void [[Bar]]() { [[B^ar]](); }");
   TU = TestTU::withCode(MainCode.code());
   // Set a file "bar.cc" on disk.
   TU.AdditionalFiles["bar.cc"] = std::string(BarCode.code());
   AST = TU.build();
-  Results = rename({MainCode.point(),
-                    NewName,
-                    AST,
-                    MainFilePath,
-                    Index.get(),
-                    {},
-                    GetDirtyBuffer});
+  Results = rename({MainCode.point(), NewName, AST, MainFilePath,
+                    createOverlay(getVFSFromAST(AST), InMemFS), Index.get()});
   ASSERT_TRUE(bool(Results)) << Results.takeError();
   EXPECT_THAT(
       applyEdits(std::move(Results->GlobalChanges)),
       UnorderedElementsAre(
           Pair(Eq(BarPath), Eq(expectedResult(BarCode, NewName))),
           Pair(Eq(MainFilePath), Eq(expectedResult(MainCode, NewName)))));
 
   // Run rename on a pagination index which couldn't return all refs in one
class PaginationIndex : public SymbolIndex {
 
     llvm::unique_function<IndexContents(llvm::StringRef) const>
     indexedFiles() const override {
       return [](llvm::StringRef) { return IndexContents::None; };
     }
 
     size_t estimateMemoryUsage() const override { return 0; }
   } PIndex;
-  Results = rename({MainCode.point(),
-                    NewName,
-                    AST,
-                    MainFilePath,
-                    &PIndex,
-                    {},
-                    GetDirtyBuffer});
+  Results = rename({MainCode.point(), NewName, AST, MainFilePath,
+                    createOverlay(getVFSFromAST(AST), InMemFS), &PIndex});
   EXPECT_FALSE(Results);
   EXPECT_THAT(llvm::toString(Results.takeError()),
               testing::HasSubstr("too many occurrences"));
 }
 
 TEST(CrossFileRenameTests, DeduplicateRefsFromIndex) {
   auto MainCode = Annotations("int [[^x]] = 2;");
   auto MainFilePath = testPath("main.cc");
public:
     indexedFiles() const override {
       return [](llvm::StringRef) { return IndexContents::None; };
     }
 
     size_t estimateMemoryUsage() const override { return 0; }
     Ref ReturnedRef;
   } DIndex(XRefInBarCC);
   llvm::StringRef NewName = "newName";
-  auto Results =
-      rename({MainCode.point(), NewName, AST, MainFilePath, &DIndex});
+  auto Results = rename({MainCode.point(), NewName, AST, MainFilePath,
+                         getVFSFromAST(AST), &DIndex});
   ASSERT_TRUE(bool(Results)) << Results.takeError();
   EXPECT_THAT(
       applyEdits(std::move(Results->GlobalChanges)),
       UnorderedElementsAre(
           Pair(Eq(BarPath), Eq(expectedResult(BarCode, NewName))),
           Pair(Eq(MainFilePath), Eq(expectedResult(MainCode, NewName)))));
 }