njames93 created this revision.
Herald added subscribers: arphaman, javed.absar.
njames93 requested review of this revision.
Herald added subscribers: cfe-commits, MaskRay.
Herald added a project: clang.

Before I start, this should probably be on clangd-dev, but that list seems 
kinda dead so I've put my working here instead.

This definitely needs splitting up and tweaking most of the specifics but here 
is a proof of concept so to say.
The majority of changes here are:

- Change Drafts to hold ref counted strings.
- Add a ThreadsafeFS to Draftstore that overlays its contents over another 
ThreadsafeFS.
- Use said Filesystem in the ClangdServer to provide a way to access its 
contents in a much friendlier way.
- Add a command line option to control using that Filesystem when building 
pre-ambles.
- Always use this Filesystem when performing x-file renames and tweaks.

With all these changes I personally find the UX so much nicer.

I've put the Option for using DirtyBuffers when building pre-ambles as a hidden 
option for now, defaulted to off.
However as per LSP spec, we should think about in the future, once this is 
fully ready and tested to default it to on.
Also I know we shouldn't be using the command line really, when we have 
`.clangd` config instead.
This will be migrated to there in due time.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D94554

Files:
  clang-tools-extra/clangd/ClangdLSPServer.cpp
  clang-tools-extra/clangd/ClangdServer.cpp
  clang-tools-extra/clangd/ClangdServer.h
  clang-tools-extra/clangd/DraftStore.cpp
  clang-tools-extra/clangd/DraftStore.h
  clang-tools-extra/clangd/TUScheduler.cpp
  clang-tools-extra/clangd/TUScheduler.h
  clang-tools-extra/clangd/refactor/Rename.cpp
  clang-tools-extra/clangd/refactor/Rename.h
  clang-tools-extra/clangd/refactor/Tweak.cpp
  clang-tools-extra/clangd/refactor/Tweak.h
  clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp
  clang-tools-extra/clangd/tool/Check.cpp
  clang-tools-extra/clangd/tool/ClangdMain.cpp
  clang-tools-extra/clangd/unittests/DraftStoreTests.cpp
  clang-tools-extra/clangd/unittests/RenameTests.cpp
  clang-tools-extra/clangd/unittests/tweaks/TweakTesting.cpp

Index: clang-tools-extra/clangd/unittests/tweaks/TweakTesting.cpp
===================================================================
--- clang-tools-extra/clangd/unittests/tweaks/TweakTesting.cpp
+++ clang-tools-extra/clangd/unittests/tweaks/TweakTesting.cpp
@@ -74,7 +74,8 @@
   SelectionTree::createEach(AST.getASTContext(), AST.getTokens(), Range.first,
                             Range.second, [&](SelectionTree ST) {
                               Tweak::Selection S(Index, AST, Range.first,
-                                                 Range.second, std::move(ST));
+                                                 Range.second, std::move(ST),
+                                                 nullptr);
                               if (auto T = prepareTweak(TweakID, S)) {
                                 Result = (*T)->apply(S);
                                 return true;
Index: clang-tools-extra/clangd/unittests/RenameTests.cpp
===================================================================
--- clang-tools-extra/clangd/unittests/RenameTests.cpp
+++ clang-tools-extra/clangd/unittests/RenameTests.cpp
@@ -33,6 +33,19 @@
 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;
@@ -833,8 +846,8 @@
     TU.ExtraArgs.push_back("-xobjective-c++");
     auto AST = TU.build();
     for (const auto &RenamePos : Code.points()) {
-      auto RenameResult =
-          rename({RenamePos, NewName, AST, testPath(TU.Filename)});
+      auto RenameResult = rename(
+          {RenamePos, NewName, AST, testPath(TU.Filename), getVFSFromAST(AST)});
       ASSERT_TRUE(bool(RenameResult)) << RenameResult.takeError();
       ASSERT_EQ(1u, RenameResult->GlobalChanges.size());
       EXPECT_EQ(
@@ -1040,8 +1053,8 @@
     }
     auto AST = TU.build();
     llvm::StringRef NewName = Case.NewName;
-    auto Results =
-        rename({T.point(), NewName, AST, testPath(TU.Filename), Case.Index});
+    auto Results = rename({T.point(), NewName, AST, testPath(TU.Filename),
+                           getVFSFromAST(AST), Case.Index});
     bool WantRename = true;
     if (T.ranges().empty())
       WantRename = false;
@@ -1081,8 +1094,8 @@
   auto AST = TU.build();
   llvm::StringRef NewName = "abcde";
 
-  auto RenameResult =
-      rename({Code.point(), NewName, AST, testPath(TU.Filename)});
+  auto RenameResult = rename(
+      {Code.point(), NewName, AST, testPath(TU.Filename), getVFSFromAST(AST)});
   ASSERT_TRUE(bool(RenameResult)) << RenameResult.takeError() << Code.point();
   ASSERT_EQ(1u, RenameResult->GlobalChanges.size());
   EXPECT_EQ(applyEdits(std::move(RenameResult->GlobalChanges)).front().second,
@@ -1098,7 +1111,8 @@
       )cpp";
   TU.HeaderFilename = "protobuf.pb.h";
   auto AST = TU.build();
-  auto Results = rename({Code.point(), "newName", AST, testPath(TU.Filename)});
+  auto Results = rename({Code.point(), "newName", AST, testPath(TU.Filename),
+                         getVFSFromAST(AST)});
   EXPECT_FALSE(Results);
   EXPECT_THAT(llvm::toString(Results.takeError()),
               testing::HasSubstr("not a supported kind"));
@@ -1170,12 +1184,10 @@
 
   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.
@@ -1186,9 +1198,9 @@
                          NewName,
                          AST,
                          MainFilePath,
+                         createOverlay(getVFSFromAST(AST), InMemFS),
                          Index.get(),
-                         {/*CrossFile=*/true},
-                         GetDirtyBuffer});
+                         {/*CrossFile=*/true}});
   ASSERT_TRUE(bool(Results)) << Results.takeError();
   EXPECT_THAT(
       applyEdits(std::move(Results->GlobalChanges)),
@@ -1207,9 +1219,9 @@
                     NewName,
                     AST,
                     MainFilePath,
+                    createOverlay(getVFSFromAST(AST), InMemFS),
                     Index.get(),
-                    {/*CrossFile=*/true},
-                    GetDirtyBuffer});
+                    {/*CrossFile=*/true}});
   ASSERT_TRUE(bool(Results)) << Results.takeError();
   EXPECT_THAT(
       applyEdits(std::move(Results->GlobalChanges)),
@@ -1249,9 +1261,9 @@
                     NewName,
                     AST,
                     MainFilePath,
+                    createOverlay(getVFSFromAST(AST), InMemFS),
                     &PIndex,
-                    {/*CrossFile=*/true},
-                    GetDirtyBuffer});
+                    {/*CrossFile=*/true}});
   EXPECT_FALSE(Results);
   EXPECT_THAT(llvm::toString(Results.takeError()),
               testing::HasSubstr("too many occurrences"));
@@ -1305,6 +1317,7 @@
                          NewName,
                          AST,
                          MainFilePath,
+                         getVFSFromAST(AST),
                          &DIndex,
                          {/*CrossFile=*/true}});
   ASSERT_TRUE(bool(Results)) << Results.takeError();
@@ -1525,7 +1538,7 @@
   auto Path = testPath(TU.Filename);
   auto AST = TU.build();
   llvm::StringRef NewName = "newName";
-  auto Results = rename({Code.point(), NewName, AST, Path});
+  auto Results = rename({Code.point(), NewName, AST, Path, getVFSFromAST(AST)});
   ASSERT_TRUE(bool(Results)) << Results.takeError();
   EXPECT_THAT(
       applyEdits(std::move(Results->GlobalChanges)),
Index: clang-tools-extra/clangd/unittests/DraftStoreTests.cpp
===================================================================
--- clang-tools-extra/clangd/unittests/DraftStoreTests.cpp
+++ clang-tools-extra/clangd/unittests/DraftStoreTests.cpp
@@ -52,8 +52,8 @@
     llvm::Expected<DraftStore::Draft> Result =
         DS.updateDraft(Path, llvm::None, {Event});
     ASSERT_TRUE(!!Result);
-    EXPECT_EQ(Result->Contents, SrcAfter.code());
-    EXPECT_EQ(DS.getDraft(Path)->Contents, SrcAfter.code());
+    EXPECT_EQ(*Result->Contents, SrcAfter.code());
+    EXPECT_EQ(*DS.getDraft(Path)->Contents, SrcAfter.code());
     EXPECT_EQ(Result->Version, static_cast<int64_t>(i));
   }
 }
@@ -84,8 +84,8 @@
       DS.updateDraft(Path, llvm::None, Changes);
 
   ASSERT_TRUE(!!Result) << llvm::toString(Result.takeError());
-  EXPECT_EQ(Result->Contents, FinalSrc.code());
-  EXPECT_EQ(DS.getDraft(Path)->Contents, FinalSrc.code());
+  EXPECT_EQ(*Result->Contents, FinalSrc.code());
+  EXPECT_EQ(*DS.getDraft(Path)->Contents, FinalSrc.code());
   EXPECT_EQ(Result->Version, 1);
 }
 
@@ -345,7 +345,7 @@
 
   Optional<DraftStore::Draft> Contents = DS.getDraft(File);
   EXPECT_TRUE(Contents);
-  EXPECT_EQ(Contents->Contents, OriginalContents);
+  EXPECT_EQ(*Contents->Contents, OriginalContents);
   EXPECT_EQ(Contents->Version, 0);
 }
 
Index: clang-tools-extra/clangd/tool/ClangdMain.cpp
===================================================================
--- clang-tools-extra/clangd/tool/ClangdMain.cpp
+++ clang-tools-extra/clangd/tool/ClangdMain.cpp
@@ -501,6 +501,12 @@
     init(ClangdServer::Options().CollectMainFileRefs),
 };
 
+opt<bool> UseDirtyPreambles{"use-dirty-preambles", cat(Misc),
+                            desc("Use files open in the editor when building "
+                                 "pre-ambles instead of reading from the disk"),
+                            Hidden,
+                            init(ClangdServer::Options().UseDirtyPreambles)};
+
 #if defined(__GLIBC__) && CLANGD_MALLOC_TRIM
 opt<bool> EnableMallocTrim{
     "malloc-trim",
@@ -883,6 +889,7 @@
     Opts.ClangTidyProvider = ClangTidyOptProvider;
   }
   Opts.AsyncPreambleBuilds = AsyncPreamble;
+  Opts.UseDirtyPreambles = UseDirtyPreambles;
   Opts.SuggestMissingIncludes = SuggestMissingIncludes;
   Opts.QueryDriverGlobs = std::move(QueryDriverGlobs);
   Opts.TweakFilter = [&](const Tweak &T) {
Index: clang-tools-extra/clangd/tool/Check.cpp
===================================================================
--- clang-tools-extra/clangd/tool/Check.cpp
+++ clang-tools-extra/clangd/tool/Check.cpp
@@ -203,7 +203,8 @@
       vlog("  {0} {1}", Pos, Tok.text(AST->getSourceManager()));
       auto Tree = SelectionTree::createRight(AST->getASTContext(),
                                              AST->getTokens(), Start, End);
-      Tweak::Selection Selection(&Index, *AST, Start, End, std::move(Tree));
+      Tweak::Selection Selection(&Index, *AST, Start, End, std::move(Tree),
+                                 nullptr);
       for (const auto &T : prepareTweaks(Selection, Opts.TweakFilter)) {
         auto Result = T->apply(Selection);
         if (!Result) {
Index: clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp
===================================================================
--- clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp
+++ clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp
@@ -63,10 +63,9 @@
 }
 
 llvm::Optional<Path> getSourceFile(llvm::StringRef FileName,
-                                   const Tweak::Selection &Sel) {
-  if (auto Source = getCorrespondingHeaderOrSource(
-          FileName,
-          &Sel.AST->getSourceManager().getFileManager().getVirtualFileSystem()))
+                                   const Tweak::Selection &Sel,
+                                   llvm::vfs::FileSystem *FS) {
+  if (auto Source = getCorrespondingHeaderOrSource(FileName, FS))
     return *Source;
   return getCorrespondingHeaderOrSource(FileName, *Sel.AST, Sel.Index);
 }
@@ -403,13 +402,14 @@
     if (!MainFileName)
       return error("Couldn't get absolute path for main file.");
 
-    auto CCFile = getSourceFile(*MainFileName, Sel);
+    auto *FS = Sel.FS;
+    assert(FS && "FS Must be set in apply");
+
+    auto CCFile = getSourceFile(*MainFileName, Sel, FS);
+
     if (!CCFile)
       return error("Couldn't find a suitable implementation file.");
-
-    auto &FS =
-        Sel.AST->getSourceManager().getFileManager().getVirtualFileSystem();
-    auto Buffer = FS.getBufferForFile(*CCFile);
+    auto Buffer = FS->getBufferForFile(*CCFile);
     // FIXME: Maybe we should consider creating the implementation file if it
     // doesn't exist?
     if (!Buffer)
Index: clang-tools-extra/clangd/refactor/Tweak.h
===================================================================
--- clang-tools-extra/clangd/refactor/Tweak.h
+++ clang-tools-extra/clangd/refactor/Tweak.h
@@ -48,7 +48,8 @@
   /// Input to prepare and apply tweaks.
   struct Selection {
     Selection(const SymbolIndex *Index, ParsedAST &AST, unsigned RangeBegin,
-              unsigned RangeEnd, SelectionTree ASTSelection);
+              unsigned RangeEnd, SelectionTree ASTSelection,
+              llvm::vfs::FileSystem *VFS);
     /// The text of the active document.
     llvm::StringRef Code;
     /// The Index for handling codebase related queries.
@@ -64,6 +65,11 @@
     unsigned SelectionEnd;
     /// The AST nodes that were selected.
     SelectionTree ASTSelection;
+    /// File system used to access source code (for cross-file tweaks).
+    /// This can be used to overlay the "dirty" contents of files open in the
+    /// editor, which (in case of headers) may not match the saved contents used
+    /// for building the AST.
+    llvm::vfs::FileSystem *FS = nullptr;
     // FIXME: provide a way to get sources and ASTs for other files.
   };
 
Index: clang-tools-extra/clangd/refactor/Tweak.cpp
===================================================================
--- clang-tools-extra/clangd/refactor/Tweak.cpp
+++ clang-tools-extra/clangd/refactor/Tweak.cpp
@@ -47,9 +47,12 @@
 
 Tweak::Selection::Selection(const SymbolIndex *Index, ParsedAST &AST,
                             unsigned RangeBegin, unsigned RangeEnd,
-                            SelectionTree ASTSelection)
+                            SelectionTree ASTSelection,
+                            llvm::vfs::FileSystem *FS)
     : Index(Index), AST(&AST), SelectionBegin(RangeBegin),
-      SelectionEnd(RangeEnd), ASTSelection(std::move(ASTSelection)) {
+      SelectionEnd(RangeEnd), ASTSelection(std::move(ASTSelection)),
+      FS(FS ? FS
+            : &AST.getSourceManager().getFileManager().getVirtualFileSystem()) {
   auto &SM = AST.getSourceManager();
   Code = SM.getBufferData(SM.getMainFileID());
   Cursor = SM.getComposedLoc(SM.getMainFileID(), RangeBegin);
Index: clang-tools-extra/clangd/refactor/Rename.h
===================================================================
--- clang-tools-extra/clangd/refactor/Rename.h
+++ clang-tools-extra/clangd/refactor/Rename.h
@@ -21,11 +21,6 @@
 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 {
   /// If true, enable cross-file rename; otherwise, only allows to rename a
   /// symbol that's only used in the current file.
@@ -45,14 +40,12 @@
   ParsedAST &AST;
   llvm::StringRef MainFilePath;
 
+  // The filesystem to query when performing cross file renames.
+  llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FS;
+
   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 {
Index: clang-tools-extra/clangd/refactor/Rename.cpp
===================================================================
--- clang-tools-extra/clangd/refactor/Rename.cpp
+++ clang-tools-extra/clangd/refactor/Rename.cpp
@@ -509,10 +509,10 @@
 // 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);
@@ -522,13 +522,16 @@
   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) {
@@ -540,7 +543,7 @@
                    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());
@@ -596,23 +599,6 @@
   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)
@@ -689,7 +675,7 @@
       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;
Index: clang-tools-extra/clangd/TUScheduler.h
===================================================================
--- clang-tools-extra/clangd/TUScheduler.h
+++ clang-tools-extra/clangd/TUScheduler.h
@@ -236,9 +236,6 @@
   /// 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,
Index: clang-tools-extra/clangd/TUScheduler.cpp
===================================================================
--- clang-tools-extra/clangd/TUScheduler.cpp
+++ clang-tools-extra/clangd/TUScheduler.cpp
@@ -1313,13 +1313,6 @@
          File);
 }
 
-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) {
   if (!PreambleTasks) {
Index: clang-tools-extra/clangd/DraftStore.h
===================================================================
--- clang-tools-extra/clangd/DraftStore.h
+++ clang-tools-extra/clangd/DraftStore.h
@@ -11,6 +11,7 @@
 
 #include "Protocol.h"
 #include "support/Path.h"
+#include "support/ThreadsafeFS.h"
 #include "clang/Basic/LLVM.h"
 #include "llvm/ADT/StringMap.h"
 #include <mutex>
@@ -28,10 +29,27 @@
 class DraftStore {
 public:
   struct Draft {
-    std::string Contents;
+    class RefCntString : public llvm::ThreadSafeRefCountedBase<RefCntString> {
+    public:
+      RefCntString(std::string Str) : Data(std::move(Str)) {}
+
+      StringRef str() const { return Data; }
+
+      operator const std::string &() const { return Data; }
+      operator StringRef() const { return str(); }
+
+    private:
+      std::string Data;
+    };
+    llvm::IntrusiveRefCntPtr<RefCntString> Contents;
     int64_t Version = -1;
+    llvm::sys::TimePoint<> MTime = std::chrono::system_clock::now();
   };
 
+  DraftStore(const ThreadsafeFS &BaseFS);
+
+  DraftStore() = default;
+
   /// \return Contents of the stored document.
   /// For untracked files, a llvm::None is returned.
   llvm::Optional<Draft> getDraft(PathRef File) const;
@@ -59,7 +77,14 @@
   /// Remove the draft from the store.
   void removeDraft(PathRef File);
 
+  const ThreadsafeFS &getDraftFS() const {
+    assert(DraftFS && "No draft fs has been set up");
+    return *DraftFS;
+  }
+
 private:
+  class DraftstoreFS;
+  std::unique_ptr<ThreadsafeFS> DraftFS;
   mutable std::mutex Mutex;
   llvm::StringMap<Draft> Drafts;
 };
Index: clang-tools-extra/clangd/DraftStore.cpp
===================================================================
--- clang-tools-extra/clangd/DraftStore.cpp
+++ clang-tools-extra/clangd/DraftStore.cpp
@@ -53,7 +53,7 @@
 
   Draft &D = Drafts[File];
   updateVersion(D, Version);
-  D.Contents = Contents.str();
+  D.Contents = llvm::makeIntrusiveRefCnt<Draft::RefCntString>(Contents.str());
   return D.Version;
 }
 
@@ -69,7 +69,7 @@
                  File);
   }
   Draft &D = EntryIt->second;
-  std::string Contents = EntryIt->second.Contents;
+  std::string Contents = *EntryIt->second.Contents;
 
   for (const TextDocumentContentChangeEvent &Change : Changes) {
     if (!Change.range) {
@@ -109,19 +109,13 @@
                    "computed range length ({1}).",
                    *Change.rangeLength, ComputedRangeLength);
 
-    std::string NewContents;
-    NewContents.reserve(*StartIndex + Change.text.length() +
-                        (Contents.length() - *EndIndex));
-
-    NewContents = Contents.substr(0, *StartIndex);
-    NewContents += Change.text;
-    NewContents += Contents.substr(*EndIndex);
-
-    Contents = std::move(NewContents);
+    Contents.replace(*StartIndex, *EndIndex - *StartIndex, Change.text);
   }
 
   updateVersion(D, Version);
-  D.Contents = std::move(Contents);
+  D.Contents =
+      llvm::makeIntrusiveRefCnt<Draft::RefCntString>(std::move(Contents));
+  D.MTime = std::chrono::system_clock::now();
   return D;
 }
 
@@ -131,5 +125,175 @@
   Drafts.erase(File);
 }
 
+namespace {
+class RefCntStringMemBuffer final : public llvm::MemoryBuffer {
+public:
+  static std::unique_ptr<RefCntStringMemBuffer>
+  create(IntrusiveRefCntPtr<DraftStore::Draft::RefCntString> Data,
+         StringRef Name) {
+    // Allocate space for the FileContentMemBuffer and its name with null
+    // terminator.
+    static_assert(alignof(RefCntStringMemBuffer) <= sizeof(void *),
+                  "Alignment requirements can't be greater than pointer");
+    auto MemSize = sizeof(RefCntStringMemBuffer) + Name.size() + 1;
+    return std::unique_ptr<RefCntStringMemBuffer>(new (operator new(
+        MemSize)) RefCntStringMemBuffer(std::move(Data), Name));
+  }
+
+  BufferKind getBufferKind() const override {
+    return MemoryBuffer::MemoryBuffer_Malloc;
+  }
+
+  StringRef getBufferIdentifier() const override {
+    return StringRef(nameBegin(), NameSize);
+  }
+
+private:
+  RefCntStringMemBuffer(
+      IntrusiveRefCntPtr<DraftStore::Draft::RefCntString> Data, StringRef Name)
+      : BufferContents(std::move(Data)), NameSize(Name.size()) {
+    MemoryBuffer::init(BufferContents->str().begin(),
+                       BufferContents->str().end(), true);
+    memcpy(nameBegin(), Name.begin(), Name.size());
+    nameBegin()[Name.size()] = '\0';
+  }
+
+  char *nameBegin() { return (char *)&this[1]; }
+  const char *nameBegin() const { return (const char *)&this[1]; }
+
+  IntrusiveRefCntPtr<DraftStore::Draft::RefCntString> BufferContents;
+  size_t NameSize;
+};
+class DraftStoreFile : public llvm::vfs::File {
+public:
+  struct FileData {
+    IntrusiveRefCntPtr<DraftStore::Draft::RefCntString> Contents;
+    llvm::sys::fs::UniqueID UID;
+    llvm::sys::TimePoint<> MTime;
+  };
+  DraftStoreFile(std::string RequestedName, FileData Data)
+      : RequestedName(std::move(RequestedName)), Data(std::move(Data)) {}
+
+  llvm::ErrorOr<llvm::vfs::Status> status() override {
+    return llvm::vfs::Status(
+        RequestedName, Data.UID, Data.MTime, 0, 0, Data.Contents->str().size(),
+        llvm::sys::fs::file_type::regular_file, llvm::sys::fs::all_all);
+  }
+
+  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>>
+  getBuffer(const Twine &Name, int64_t FileSize, bool RequiresNullTerminator,
+            bool IsVolatile) override {
+    return RefCntStringMemBuffer::create(Data.Contents, RequestedName);
+  }
+
+  std::error_code close() override { return {}; }
+
+  ~DraftStoreFile() override { close(); }
+
+private:
+  std::string RequestedName;
+  FileData Data;
+};
+
+class DraftStoreFileSystem : public llvm::vfs::FileSystem {
+public:
+  DraftStoreFileSystem(llvm::StringMap<DraftStoreFile::FileData> Contents)
+      : Contents(std::move(Contents)) {}
+
+  llvm::ErrorOr<llvm::vfs::Status> status(const Twine &Path) override {
+    SmallString<256> PathStorage;
+    Path.toVector(PathStorage);
+    auto EC = makeAbsolute(PathStorage);
+    assert(!EC);
+    (void)EC;
+    auto Iter = Contents.find(PathStorage);
+    if (Iter == Contents.end())
+      return llvm::errc::no_such_file_or_directory;
+    return DraftStoreFile(std::string(PathStorage), Iter->getValue()).status();
+  }
+
+  llvm::ErrorOr<std::unique_ptr<llvm::vfs::File>>
+  openFileForRead(const Twine &Path) override {
+    SmallString<256> PathStorage;
+    Path.toVector(PathStorage);
+    auto EC = makeAbsolute(PathStorage);
+    assert(!EC);
+    (void)EC;
+    auto Iter = Contents.find(PathStorage);
+    if (Iter == Contents.end())
+      return llvm::errc::no_such_file_or_directory;
+    return std::make_unique<DraftStoreFile>(std::string(PathStorage),
+                                            Iter->getValue());
+  }
+
+  llvm::ErrorOr<std::string> getCurrentWorkingDirectory() const override {
+    return WorkingDirectory;
+  }
+
+  std::error_code setCurrentWorkingDirectory(const Twine &P) override {
+    SmallString<128> Path;
+    P.toVector(Path);
+
+    // Fix up relative paths. This just prepends the current working
+    // directory.
+    std::error_code EC = makeAbsolute(Path);
+    assert(!EC);
+    (void)EC;
+
+    llvm::sys::path::remove_dots(Path, /*remove_dot_dot=*/true);
+
+    if (!Path.empty())
+      WorkingDirectory = std::string(Path);
+    return {};
+  }
+  llvm::vfs::directory_iterator dir_begin(const Twine &Dir,
+                                          std::error_code &EC) override {
+    EC = llvm::errc::no_such_file_or_directory;
+    return llvm::vfs::directory_iterator();
+  }
+
+private:
+  std::string WorkingDirectory;
+  llvm::StringMap<DraftStoreFile::FileData> Contents;
+};
+
+} // namespace
+
+class DraftStore::DraftstoreFS : public ThreadsafeFS {
+public:
+  DraftstoreFS(const ThreadsafeFS &Base, const DraftStore &DS)
+      : Base(Base), DS(DS) {}
+
+  llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> viewImpl() const override {
+    auto BaseView = Base.view(llvm::None);
+    llvm::StringMap<DraftStoreFile::FileData> Contents;
+    std::lock_guard<std::mutex> Guard(DS.Mutex);
+    for (const auto &KV : DS.Drafts) {
+      // Query the base filesystem for file uniqueids.
+      auto BaseStatus = BaseView->status(KV.getKey());
+      if (!BaseStatus) {
+        elog("Couldn't find file {0} when building DirtyFS", KV.getKey());
+        continue;
+      }
+      // log("Adding dirty file {0} to dirty buffer", KV.getKey());
+      Contents.insert(std::make_pair(
+          KV.getKey(), DraftStoreFile::FileData{KV.getValue().Contents,
+                                                BaseStatus->getUniqueID(),
+                                                KV.getValue().MTime}));
+    }
+    auto OFS = llvm::makeIntrusiveRefCnt<llvm::vfs::OverlayFileSystem>(
+        std::move(BaseView));
+    OFS->pushOverlay(
+        llvm::makeIntrusiveRefCnt<DraftStoreFileSystem>(std::move(Contents)));
+    return OFS;
+  }
+
+private:
+  const ThreadsafeFS &Base;
+  const DraftStore &DS;
+};
+
+DraftStore::DraftStore(const ThreadsafeFS &BaseFS)
+    : DraftFS(std::make_unique<DraftstoreFS>(BaseFS, *this)) {}
 } // namespace clangd
 } // namespace clang
Index: clang-tools-extra/clangd/ClangdServer.h
===================================================================
--- clang-tools-extra/clangd/ClangdServer.h
+++ clang-tools-extra/clangd/ClangdServer.h
@@ -156,6 +156,9 @@
     /// Enable preview of FoldingRanges feature.
     bool FoldingRanges = false;
 
+    /// If true, use the dirty buffer contents when building Preambles.
+    bool UseDirtyPreambles = false;
+
     explicit operator TUScheduler::Options() const;
   };
   // Sensible default options for use in tests.
@@ -170,7 +173,12 @@
   /// those arguments for subsequent reparses. However, ClangdServer will check
   /// if compilation arguments changed on calls to forceReparse().
   ClangdServer(const GlobalCompilationDatabase &CDB, const ThreadsafeFS &TFS,
-               const Options &Opts, Callbacks *Callbacks = nullptr);
+               const ThreadsafeFS &DirtyFS, const Options &Opts,
+               Callbacks *Callbacks = nullptr);
+
+  ClangdServer(const GlobalCompilationDatabase &CDB, const ThreadsafeFS &TFS,
+               const Options &Opts, Callbacks *Callbacks = nullptr)
+      : ClangdServer(CDB, TFS, TFS, Opts, Callbacks) {}
 
   /// Add a \p File to the list of tracked C++ files or update the contents if
   /// \p File is already tracked. Also schedules parsing of the AST for it on a
@@ -363,6 +371,8 @@
   config::Provider *ConfigProvider = nullptr;
 
   const ThreadsafeFS &TFS;
+  /// A File system that overlays open documents over the underlying filesystem.
+  const ThreadsafeFS &DirtyFS;
   Callbacks *ServerCallbacks = nullptr;
   mutable std::mutex ConfigDiagnosticsMu;
 
@@ -392,6 +402,8 @@
   // If true, preserve the type for recovery AST.
   bool PreserveRecoveryASTType = false;
 
+  bool UseDirtyPreambles = false;
+
   // GUARDED_BY(CachedCompletionFuzzyFindRequestMutex)
   llvm::StringMap<llvm::Optional<FuzzyFindRequest>>
       CachedCompletionFuzzyFindRequestByFile;
Index: clang-tools-extra/clangd/ClangdServer.cpp
===================================================================
--- clang-tools-extra/clangd/ClangdServer.cpp
+++ clang-tools-extra/clangd/ClangdServer.cpp
@@ -137,9 +137,10 @@
 }
 
 ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
-                           const ThreadsafeFS &TFS, const Options &Opts,
-                           Callbacks *Callbacks)
-    : ConfigProvider(Opts.ConfigProvider), TFS(TFS), ServerCallbacks(Callbacks),
+                           const ThreadsafeFS &TFS, const ThreadsafeFS &DirtyFS,
+                           const Options &Opts, Callbacks *Callbacks)
+    : ConfigProvider(Opts.ConfigProvider), TFS(TFS), DirtyFS(DirtyFS),
+      ServerCallbacks(Callbacks),
       DynamicIdx(Opts.BuildDynamicSymbolIndex
                      ? new FileIndex(Opts.HeavyweightDynamicSymbolIndex,
                                      Opts.CollectMainFileRefs)
@@ -148,6 +149,7 @@
       SuggestMissingIncludes(Opts.SuggestMissingIncludes),
       BuildRecoveryAST(Opts.BuildRecoveryAST),
       PreserveRecoveryASTType(Opts.PreserveRecoveryASTType),
+      UseDirtyPreambles(Opts.UseDirtyPreambles),
       WorkspaceRoot(Opts.WorkspaceRoot),
       // Pass a callback into `WorkScheduler` to extract symbols from a newly
       // parsed file and rebuild the file index synchronously each time an AST
@@ -206,7 +208,7 @@
 
   // Compile command is set asynchronously during update, as it can be slow.
   ParseInputs Inputs;
-  Inputs.TFS = &TFS;
+  Inputs.TFS = UseDirtyPreambles ? &DirtyFS : &TFS;
   Inputs.Contents = std::string(Contents);
   Inputs.Version = Version.str();
   Inputs.ForceRebuild = ForceRebuild;
@@ -254,7 +256,8 @@
         }
       }
     }
-    ParseInputs ParseInput{IP->Command, &TFS, IP->Contents.str()};
+    ParseInputs ParseInput{IP->Command, UseDirtyPreambles ? &DirtyFS : &TFS,
+                           IP->Contents.str()};
     ParseInput.Index = Index;
     ParseInput.Opts.BuildRecoveryAST = BuildRecoveryAST;
     ParseInput.Opts.PreserveRecoveryASTType = PreserveRecoveryASTType;
@@ -300,7 +303,8 @@
     if (!PreambleData)
       return CB(error("Failed to parse includes"));
 
-    ParseInputs ParseInput{IP->Command, &TFS, IP->Contents.str()};
+    ParseInputs ParseInput{IP->Command, UseDirtyPreambles ? &DirtyFS : &TFS,
+                           IP->Contents.str()};
     ParseInput.Index = Index;
     ParseInput.Opts.BuildRecoveryAST = BuildRecoveryAST;
     ParseInput.Opts.PreserveRecoveryASTType = PreserveRecoveryASTType;
@@ -372,6 +376,7 @@
     //    main-file occurrences;
     auto Results = clangd::rename(
         {Pos, NewName.getValueOr("__clangd_rename_dummy"), InpAST->AST, File,
+         (RenameOpts.AllowCrossFile ? DirtyFS : TFS).view(llvm::None),
          RenameOpts.AllowCrossFile ? nullptr : Index, RenameOpts});
     if (!Results) {
       // LSP says to return null on failure, but that will result in a generic
@@ -387,25 +392,16 @@
 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());
 
@@ -429,7 +425,8 @@
 // May generate several candidate selections, due to SelectionTree ambiguity.
 // vector of pointers because GCC doesn't like non-copyable Selection.
 static llvm::Expected<std::vector<std::unique_ptr<Tweak::Selection>>>
-tweakSelection(const Range &Sel, const InputsAndAST &AST) {
+tweakSelection(const Range &Sel, const InputsAndAST &AST,
+               llvm::vfs::FileSystem *FS) {
   auto Begin = positionToOffset(AST.Inputs.Contents, Sel.start);
   if (!Begin)
     return Begin.takeError();
@@ -441,7 +438,7 @@
       AST.AST.getASTContext(), AST.AST.getTokens(), *Begin, *End,
       [&](SelectionTree T) {
         Result.push_back(std::make_unique<Tweak::Selection>(
-            AST.Inputs.Index, AST.AST, *Begin, *End, std::move(T)));
+            AST.Inputs.Index, AST.AST, *Begin, *End, std::move(T), FS));
         return false;
       });
   assert(!Result.empty() && "Expected at least one SelectionTree");
@@ -454,12 +451,14 @@
   // Tracks number of times a tweak has been offered.
   static constexpr trace::Metric TweakAvailable(
       "tweak_available", trace::Metric::Counter, "tweak_id");
-  auto Action = [File = File.str(), Sel, CB = std::move(CB),
+  auto Action = [File = File.str(), Sel, CB = std::move(CB), &TFS = this->TFS,
                  Filter =
                      std::move(Filter)](Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
-    auto Selections = tweakSelection(Sel, *InpAST);
+    // FIXME: Should we use the dirty fs here?
+    auto FS = TFS.view(llvm::None);
+    auto Selections = tweakSelection(Sel, *InpAST, FS.get());
     if (!Selections)
       return CB(Selections.takeError());
     std::vector<TweakRef> Res;
@@ -497,13 +496,14 @@
                  this](Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
-    auto Selections = tweakSelection(Sel, *InpAST);
+    auto FS = DirtyFS.view(llvm::None);
+    auto Selections = tweakSelection(Sel, *InpAST, FS.get());
     if (!Selections)
       return CB(Selections.takeError());
     llvm::Optional<llvm::Expected<Tweak::Effect>> Effect;
     // Try each selection, take the first one that prepare()s.
     // If they all fail, Effect will hold get the last error.
-    for (const auto &Selection : *Selections) {
+    for (auto &Selection : *Selections) {
       auto T = prepareTweak(TweakID, *Selection);
       if (T) {
         Effect = (*T)->apply(*Selection);
Index: clang-tools-extra/clangd/ClangdLSPServer.cpp
===================================================================
--- clang-tools-extra/clangd/ClangdLSPServer.cpp
+++ clang-tools-extra/clangd/ClangdLSPServer.cpp
@@ -140,7 +140,7 @@
       // If the file is open in user's editor, make sure the version we
       // saw and current version are compatible as this is the text that
       // will be replaced by editors.
-      if (!It.second.canApplyTo(Draft->Contents)) {
+      if (!It.second.canApplyTo(*Draft->Contents)) {
         ++InvalidFileCount;
         LastInvalidFile = It.first();
       }
@@ -518,7 +518,7 @@
     llvm::Optional<WithContextValue> WithOffsetEncoding;
     if (Opts.Encoding)
       WithOffsetEncoding.emplace(kCurrentOffsetEncoding, *Opts.Encoding);
-    Server.emplace(*CDB, TFS, Opts,
+    Server.emplace(*CDB, TFS, DraftMgr.getDraftFS(), Opts,
                    static_cast<ClangdServer::Callbacks *>(this));
   }
   applyConfiguration(Params.initializationOptions.ConfigSettings);
@@ -695,7 +695,7 @@
     return;
   }
 
-  Server->addDocument(File, Draft->Contents, encodeVersion(Draft->Version),
+  Server->addDocument(File, *Draft->Contents, encodeVersion(Draft->Version),
                       WantDiags, Params.forceRebuild);
 }
 
@@ -889,7 +889,7 @@
         "onDocumentOnTypeFormatting called for non-added file",
         ErrorCode::InvalidParams));
 
-  Server->formatOnType(File, Code->Contents, Params.position, Params.ch,
+  Server->formatOnType(File, *Code->Contents, Params.position, Params.ch,
                        std::move(Reply));
 }
 
@@ -904,11 +904,11 @@
         ErrorCode::InvalidParams));
 
   Server->formatRange(
-      File, Code->Contents, Params.range,
+      File, *Code->Contents, Params.range,
       [Code = Code->Contents, Reply = std::move(Reply)](
           llvm::Expected<tooling::Replacements> Result) mutable {
         if (Result)
-          Reply(replacementsToEdits(Code, Result.get()));
+          Reply(replacementsToEdits(*Code, Result.get()));
         else
           Reply(Result.takeError());
       });
@@ -924,11 +924,11 @@
         "onDocumentFormatting called for non-added file",
         ErrorCode::InvalidParams));
 
-  Server->formatFile(File, Code->Contents,
+  Server->formatFile(File, *Code->Contents,
                      [Code = Code->Contents, Reply = std::move(Reply)](
                          llvm::Expected<tooling::Replacements> Result) mutable {
                        if (Result)
-                         Reply(replacementsToEdits(Code, Result.get()));
+                         Reply(replacementsToEdits(*Code, Result.get()));
                        else
                          Reply(Result.takeError());
                      });
@@ -1468,7 +1468,8 @@
       BackgroundContext(Context::current().clone()), Transp(Transp),
       MsgHandler(new MessageHandler(*this)), TFS(TFS),
       SupportedSymbolKinds(defaultSymbolKinds()),
-      SupportedCompletionItemKinds(defaultCompletionItemKinds()), Opts(Opts) {
+      SupportedCompletionItemKinds(defaultCompletionItemKinds()), DraftMgr(TFS),
+      Opts(Opts) {
 
   // clang-format off
   MsgHandler->bind("initialize", &ClangdLSPServer::onInitialize);
@@ -1567,14 +1568,14 @@
   auto Code = DraftMgr.getDraft(Params.textDocument.uri.file());
   if (!Code)
     return true; // completion code will log the error for untracked doc.
-  auto Offset = positionToOffset(Code->Contents, Params.position,
+  auto Offset = positionToOffset(*Code->Contents, Params.position,
                                  /*AllowColumnsBeyondLineLength=*/false);
   if (!Offset) {
     vlog("could not convert position '{0}' to offset for file '{1}'",
          Params.position, Params.textDocument.uri.file());
     return true;
   }
-  return allowImplicitCompletion(Code->Contents, *Offset);
+  return allowImplicitCompletion(*Code->Contents, *Offset);
 }
 
 void ClangdLSPServer::onHighlightingsReady(
@@ -1715,7 +1716,7 @@
   for (const Path &FilePath : DraftMgr.getActiveFiles())
     if (Filter(FilePath))
       if (auto Draft = DraftMgr.getDraft(FilePath)) // else disappeared in race?
-        Server->addDocument(FilePath, std::move(Draft->Contents),
+        Server->addDocument(FilePath, *Draft->Contents,
                             encodeVersion(Draft->Version),
                             WantDiagnostics::Auto);
 }
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to