njames93 updated this revision to Diff 316357.
njames93 added a comment.

Remove some unrelated changes

Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D94554/new/

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) {
@@ -121,7 +121,9 @@
   }
 
   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 +133,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,7 +496,8 @@
                  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;
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