[clangd] Use dirty filesystem when performing cross file tweaks

[njames93]

Summary

Cross file tweaks can now use the dirty buffer contents easily when performing cross file effects.
This can be noted on the DefineOutline tweak, now working when the target file is unsaved

[sammccall]

Ah, thanks for working on this!

A few thoughts:

• when we're pseudoparsing the file we're going to modify as we do here, using the new content is strictly better, no downside :-)
• the old method here of accessing the content through the FS attached to the AST is a hack
• however I think DraftStore is the wrong abstraction here and VFS is the right one.
  • each tweak shouldn't have to deal with falling back from DraftStore to VFS (and rename also has an implementation of this!)
  • C++ embedders don't have a DraftStore lying around
  • bug: this code won't find a named but unsaved implementation file (since getCorrespondingHeaderOrSource uses VFS)
• (also Tweak::Selection isn't a pure abstraction, we can jam the FS in there instead of passing it around as a separate param everywhere)

So I think the right way to expose this to tweaks is to add a vfs::FileSystem * member to Tweak::Selection, that allows accessing the current FS (as opposed to the one used for building).
The simplest implementation uses OverlayVFS + InMemoryFS + TUScheduler::getAllFileContents(). This isn't as expensive as it sounds as we need only set the FS pointer for apply().
(We can also implement a FS that copies more lazily, though I suspect it's not worth it)

[njames93]

In D93978#2479440, @sammccall wrote:

Ah, thanks for working on this!

A few thoughts:

• when we're pseudoparsing the file we're going to modify as we do here, using the new content is strictly better, no downside :-)
• the old method here of accessing the content through the FS attached to the AST is a hack
• however I think DraftStore is the wrong abstraction here and VFS is the right one.
  • each tweak shouldn't have to deal with falling back from DraftStore to VFS (and rename also has an implementation of this!)
  • C++ embedders don't have a DraftStore lying around
  • bug: this code won't find a named but unsaved implementation file (since getCorrespondingHeaderOrSource uses VFS)
• (also Tweak::Selection isn't a pure abstraction, we can jam the FS in there instead of passing it around as a separate param everywhere)

So I think the right way to expose this to tweaks is to add a vfs::FileSystem * member to Tweak::Selection, that allows accessing the current FS (as opposed to the one used for building).
The simplest implementation uses OverlayVFS + InMemoryFS + TUScheduler::getAllFileContents(). This isn't as expensive as it sounds as we need only set the FS pointer for apply().
(We can also implement a FS that copies more lazily, though I suspect it's not worth it)

I'm getting a little stuck here. I've made a OverlayFS which has the a TFS::view as the base, and an InMemoryFS containing the contents of the TUScheduler. I Passed that in the Selection and it all worked.
However once I changed the getSourceFile function to also use that same FS, it would find the file ok, but then clangd would crash(with no visible trace) when it tried to open the file.

auto Buffer = FS.getBufferForFile(*CCFile); // Here

From what I can getCorrespondingSourceOrHeader is querying the InMemoryFS and maybe those calls to FS::exists are somehow corrupting something.
Whatever it is, it all seems a bit messed up to me.

I'm gonna try with a debug build (and possibly asan) to see If I could possibly get anything more.

scratch that, asan told me that the OverlayFS was getting destroyed as I had it stored as a unique_ptr then it was copied into a RefCntPtr and subsequently destroyed in there

[njames93]

Use VFS instead of DraftStore

[sammccall]

Thanks, this looks much better!

clang-tools-extra/clangd/ClangdServer.cpp

431

nit: pass nullptr explicitly rather than as a default arg.

clang-tools-extra/clangd/TUScheduler.cpp

1462 ↗

(On Diff #315070)

nit: say why it's needless: usually 0-1 buffer is ever read.

clang-tools-extra/clangd/TUScheduler.h

239 ↗

(On Diff #315070)

FIXME to remove this?

243 ↗

(On Diff #315070)

I do think this is the best option, but just wanted to spell out the implications to make sure we're on the same page...

• edits that occur while the action is queued will not be seen. (This is bad for the use cases we know of, but minor)
• the optimization of deferring buffer copies (hinted to in the implementation) would actually break the contract, as you may see a *newer* version of the file
• we don't have to worry about locking inside TUScheduler to grab the buffers

246 ↗

(On Diff #315070)

there's no need for both of these to be methods on TUScheduler, pick one?

(I'd lean towards exposing this one only as it simplifies callers and has a better name)

248 ↗

(On Diff #315070)

nit: take const TFS& as arg, as this function is designed for caller convenience and that's always how it'll be used

248 ↗

(On Diff #315070)

nit: return unique_ptr<FileSystem> (which can implicitly convert)

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

52

VFS is a public field, and optional, so we could omit it from the constructor and assign it explicitly instead.

68

let's be a bit more explicit here.

/// File system used to access source code (for cross-file tweaks).
/// This overlays 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.

68

mention that this is only populated for apply(), not prepare()?

69

nit: FS rather than VFS
(the "V" is already implied by the presence of the object)

clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp

405–409

It's not actually optional for apply, right? We shouldn't need this fallback.

clang-tools-extra/clangd/tool/Check.cpp

207

hmm, I see.

What if we make the Selection constructor fill in the default FS from AST, and then just overwrote the public field in ClangdServer::applyTweak?

Then we get the FS always being set (good for tweaks), simple usage here or in tests, and the expensive copy is still explicit...

[njames93]

clang-tools-extra/clangd/TUScheduler.h

248 ↗

(On Diff #315070)

I think the flexibility of passing a Filesystem is better than the ThreadsafeFS.

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

52

We could, but it would messy code elsewhere.

clang-tools-extra/clangd/tool/Check.cpp

207

I've gone with an approach that the FS gets passed in the constructor, if nullptr is passed then we use the fallback, helpful for tests

[njames93]

Address most comments.

[njames93]

So I did some tweaking and I have a snapshot filesystem that will hold shared references to the contents, thereby avoiding the need to copy the file contents.
It works for rename and tweaks with no visible issues.
Reasoning behind this a long term goal of letting a user choose between using on disk or dirty buffers when building preambles.
Obviously there, if using dirty buffers, we wouldn't want to be copying every open file (even unused files) during preamble building/checking.

[njames93]

I've decided I much prefer the implementation in D94554, So once thats cleaned up and split up, I'll merge changes in here with there.

[njames93]

Change implementation to use DirtyFS instead.

[njames93]

Fixed last series being a little corrupted.

[njames93]

Rebased.

[njames93]

clang-tools-extra/clangd/ClangdServer.cpp

461–463

Regarding this, Is it wise to set the contract so that Tweak::prepare isn't allowed to do IO so we should just pass a nullptr here, inline with how prepareRename works.

[sammccall]

Sorry, I thought this had already gone in!

LG, sorry for the back and forth.

clang-tools-extra/clangd/ClangdServer.cpp

461–463

Yes, I think so.

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

54

I'm no longer convinced this fallback is a good idea.
We want prepare() to do no IO, ideally we'd assert if we do. But by falling back to the AST FS, we'll mask that problem.

There are only two callsites that need this fallback (TweakTesting & Check) - I think we should just inline what we mean there.

clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp

67

why take this as an explicit parameter rather than just accessing (and asserting) Sel.FS?

[njames93]

Address comments.

[njames93]

Fix up compilation failures and test failures.

clang-tools-extra/clangd/ClangdLSPServer.cpp

 llvm::Error validateEdits(const DraftStore &DraftMgr, const FileEdits &FE) {
   size_t InvalidFileCount = 0;
   llvm::StringRef LastInvalidFile;
   for (const auto &It : FE) {
     if (auto Draft = DraftMgr.getDraft(It.first())) {
       // 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();
       }
     }
   }
   if (!InvalidFileCount)
     return llvm::Error::success();
   if (InvalidFileCount == 1)
void ClangdLSPServer::onInitialize(const InitializeParams &Params,
   {
     // Switch caller's context with LSPServer's background context. Since we
     // rather want to propagate information from LSPServer's context into the
     // Server, CDB, etc.
     WithContext MainContext(BackgroundContext.clone());
     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);
 
   Opts.CodeComplete.EnableSnippets = Params.capabilities.CompletionSnippets;
   Opts.CodeComplete.IncludeFixIts = Params.capabilities.CompletionFixes;
   if (!Opts.CodeComplete.BundleOverloads.hasValue())
     Opts.CodeComplete.BundleOverloads = Params.capabilities.HasSignatureHelp;
if (!Draft) {
     // the client.  It is better to remove the draft and let further operations
     // fail rather than giving wrong results.
     DraftMgr.removeDraft(File);
     Server->removeDocument(File);
     elog("Failed to update {0}: {1}", File, Draft.takeError());
     return;
   }
 
-  Server->addDocument(File, Draft->Contents, encodeVersion(Draft->Version),
+  Server->addDocument(File, *Draft->Contents, encodeVersion(Draft->Version),
                       WantDiags, Params.forceRebuild);
 }
 
 void ClangdLSPServer::onDocumentDidSave(
     const DidSaveTextDocumentParams &Params) {
   reparseOpenFilesIfNeeded([](llvm::StringRef) { return true; });
 }
 
void ClangdLSPServer::onDocumentOnTypeFormatting(
     Callback<std::vector<TextEdit>> Reply) {
   auto File = Params.textDocument.uri.file();
   auto Code = DraftMgr.getDraft(File);
   if (!Code)
     return Reply(llvm::make_error<LSPError>(
         "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));
 }
 
 void ClangdLSPServer::onDocumentRangeFormatting(
     const DocumentRangeFormattingParams &Params,
     Callback<std::vector<TextEdit>> Reply) {
   auto File = Params.textDocument.uri.file();
   auto Code = DraftMgr.getDraft(File);
   if (!Code)
     return Reply(llvm::make_error<LSPError>(
         "onDocumentRangeFormatting called for non-added file",
         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());
       });
 }
 
 void ClangdLSPServer::onDocumentFormatting(
     const DocumentFormattingParams &Params,
     Callback<std::vector<TextEdit>> Reply) {
   auto File = Params.textDocument.uri.file();
   auto Code = DraftMgr.getDraft(File);
   if (!Code)
     return Reply(llvm::make_error<LSPError>(
         "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());
                      });
 }
 
 /// The functions constructs a flattened view of the DocumentSymbol hierarchy.
 /// Used by the clients that do not support the hierarchical view.
 static std::vector<SymbolInformation>
ClangdLSPServer::ClangdLSPServer(class Transport &Transp,
                                  const ClangdLSPServer::Options &Opts)
     : ShouldProfile(/*Period=*/std::chrono::minutes(5),
                     /*Delay=*/std::chrono::minutes(1)),
       ShouldCleanupMemory(/*Period=*/std::chrono::minutes(1),
                           /*Delay=*/std::chrono::minutes(1)),
       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);
   MsgHandler->bind("initialized", &ClangdLSPServer::onInitialized);
   MsgHandler->bind("shutdown", &ClangdLSPServer::onShutdown);
   MsgHandler->bind("sync", &ClangdLSPServer::onSync);
   MsgHandler->bind("textDocument/rangeFormatting", &ClangdLSPServer::onDocumentRangeFormatting);
   MsgHandler->bind("textDocument/onTypeFormatting", &ClangdLSPServer::onDocumentOnTypeFormatting);

 // of simplicity here.
 bool ClangdLSPServer::shouldRunCompletion(
     const CompletionParams &Params) const {
   if (Params.context.triggerKind != CompletionTriggerKind::TriggerCharacter)
     return true;
   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(
     PathRef File, llvm::StringRef Version,
     std::vector<HighlightingToken> Highlightings) {
   std::vector<HighlightingToken> Old;
   std::vector<HighlightingToken> HighlightingsCopy = Highlightings;
   {

 }
 
 void ClangdLSPServer::reparseOpenFilesIfNeeded(
     llvm::function_ref<bool(llvm::StringRef File)> Filter) {
   // Reparse only opened files that were modified.
   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);
 }
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/ClangdServer.h

public:
   /// Creates a new ClangdServer instance.
   ///
   /// ClangdServer uses \p CDB to obtain compilation arguments for parsing. Note
   /// that ClangdServer only obtains compilation arguments once for each newly
   /// added file (i.e., when processing a first call to addDocument) and reuses
   /// 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
   /// separate thread. When the parsing is complete, DiagConsumer passed in
   /// constructor will receive onDiagnosticsReady callback.
   /// Version identifies this snapshot and is propagated to ASTs, preambles,
   /// diagnostics etc built from it.
   void addDocument(PathRef File, StringRef Contents,
private:
   /// The empty string means no particular file is the target.
   /// Rather than called by each feature, this is exposed to the components
   /// that control worker threads, like TUScheduler and BackgroundIndex.
   /// This means it's OK to do some IO here, and it cuts across all features.
   Context createProcessingContext(PathRef) const;
   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;
 
   Path ResourceDir;
   // The index used to look up symbols. This could be:
   //   - null (all index functionality is optional)
   //   - the dynamic index owned by ClangdServer (DynamicIdx)
   //   - the static index passed to the constructor

clang-tools-extra/clangd/ClangdServer.cpp

ClangdServer::Options::operator TUScheduler::Options() const {
   Opts.RetentionPolicy = RetentionPolicy;
   Opts.StorePreamblesInMemory = StorePreamblesInMemory;
   Opts.UpdateDebounce = UpdateDebounce;
   Opts.AsyncPreambleBuilds = AsyncPreambleBuilds;
   return Opts;
 }
 
 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)
                      : nullptr),
       ClangTidyProvider(Opts.ClangTidyProvider),
       SuggestMissingIncludes(Opts.SuggestMissingIncludes),
       BuildRecoveryAST(Opts.BuildRecoveryAST),
       PreserveRecoveryASTType(Opts.PreserveRecoveryASTType),
auto Action = [File = File.str(), NewName = NewName.str(), Pos, Opts,
     return CB(*R);
   };
   WorkScheduler.runWithAST("Rename", File, std::move(Action));
 }
 
 // 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();
   auto End = positionToOffset(AST.Inputs.Contents, Sel.end);
   if (!End)
     return End.takeError();
   std::vector<std::unique_ptr<Tweak::Selection>> Result;
   SelectionTree::createEach(
       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");
   return std::move(Result);
 }
 
 void ClangdServer::enumerateTweaks(
     PathRef File, Range Sel, llvm::unique_function<bool(const Tweak &)> Filter,
     Callback<std::vector<TweakRef>> CB) {
   // 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;
     // Don't allow a tweak to fire more than once across ambiguous selections.
     llvm::DenseSet<llvm::StringRef> PreparedTweaks;
     auto DeduplicatingFilter = [&](const Tweak &T) {
       return Filter(T) && !PreparedTweaks.count(T.id());
     };
void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef TweakID,
   static constexpr trace::Metric TweakFailed(
       "tweak_failed", trace::Metric::Counter, "tweak_id");
   TweakAttempt.record(1, TweakID);
   auto Action = [File = File.str(), Sel, TweakID = TweakID.str(),
                  CB = std::move(CB),
                  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) {
       auto T = prepareTweak(TweakID, *Selection);
       if (T) {

clang-tools-extra/clangd/DraftStore.h

 //===--- DraftStore.h - File contents container -----------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_DRAFTSTORE_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_DRAFTSTORE_H
 
 #include "Protocol.h"
 #include "support/Path.h"
+#include "support/ThreadsafeFS.h"
 #include "clang/Basic/LLVM.h"
 #include "llvm/ADT/StringMap.h"
 #include <mutex>
 #include <string>
 #include <vector>
 
 namespace clang {
 namespace clangd {
 
+/// A Reference counted string that is used to hold the contents of open files.
+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;
+};
+
 /// A thread-safe container for files opened in a workspace, addressed by
 /// filenames. The contents are owned by the DraftStore. This class supports
 /// both whole and incremental updates of the documents.
 /// Each time a draft is updated, it is assigned a version number. This can be
 /// specified by the caller or incremented from the previous version.
 class DraftStore {
 public:
   struct Draft {
-    std::string Contents;
+    llvm::IntrusiveRefCntPtr<RefCntString> Contents;
     int64_t Version = -1;
   };
 
+  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;
 
   /// \return List of names of the drafts in this store.
   std::vector<Path> getActiveFiles() const;
 
   /// Replace contents of the draft for \p File with \p Contents.
public:
   /// changes couldn't be applied.
   llvm::Expected<Draft>
   updateDraft(PathRef File, llvm::Optional<int64_t> Version,
               llvm::ArrayRef<TextDocumentContentChangeEvent> Changes);
 
   /// 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:
+  struct DraftAndTime {
+    Draft Draft;
+    llvm::sys::TimePoint<> MTime;
+  };
+  class DraftstoreFS;
+  std::unique_ptr<ThreadsafeFS> DraftFS;
   mutable std::mutex Mutex;
-  llvm::StringMap<Draft> Drafts;
+  llvm::StringMap<DraftAndTime> Drafts;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/clangd/DraftStore.cpp

 //===--- DraftStore.cpp - File contents container ---------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "DraftStore.h"
 #include "SourceCode.h"
 #include "support/Logger.h"
+#include "llvm/Support/Chrono.h"
 #include "llvm/Support/Errc.h"
+#include <cstdint>
 
 namespace clang {
 namespace clangd {
 
 llvm::Optional<DraftStore::Draft> DraftStore::getDraft(PathRef File) const {
   std::lock_guard<std::mutex> Lock(Mutex);
 
   auto It = Drafts.find(File);
   if (It == Drafts.end())
     return None;
 
-  return It->second;
+  return It->second.Draft;
 }
 
 std::vector<Path> DraftStore::getActiveFiles() const {
   std::lock_guard<std::mutex> Lock(Mutex);
   std::vector<Path> ResultVector;
 
   for (auto DraftIt = Drafts.begin(); DraftIt != Drafts.end(); DraftIt++)
     ResultVector.push_back(std::string(DraftIt->getKey()));
if (*Version <= D.Version)
       log("File version went from {0} to {1}", D.Version, Version);
     D.Version = *Version;
   } else {
     // Note that if D was newly-created, this will bump D.Version from -1 to 0.
     ++D.Version;
   }
 }
 
+static void updateTime(llvm::sys::TimePoint<> &Time) {
+  Time = llvm::sys::TimePoint<>::clock::now();
+}
+
 int64_t DraftStore::addDraft(PathRef File, llvm::Optional<int64_t> Version,
                          llvm::StringRef Contents) {
   std::lock_guard<std::mutex> Lock(Mutex);
 
-  Draft &D = Drafts[File];
-  updateVersion(D, Version);
-  D.Contents = Contents.str();
-  return D.Version;
+  DraftAndTime &D = Drafts[File];
+  updateVersion(D.Draft, Version);
+  updateTime(D.MTime);
+  D.Draft.Contents = llvm::makeIntrusiveRefCnt<RefCntString>(Contents.str());
+  return D.Draft.Version;
 }
 
 llvm::Expected<DraftStore::Draft> DraftStore::updateDraft(
     PathRef File, llvm::Optional<int64_t> Version,
     llvm::ArrayRef<TextDocumentContentChangeEvent> Changes) {
   std::lock_guard<std::mutex> Lock(Mutex);
 
   auto EntryIt = Drafts.find(File);
   if (EntryIt == Drafts.end()) {
     return error(llvm::errc::invalid_argument,
                  "Trying to do incremental update on non-added document: {0}",
                  File);
   }
-  Draft &D = EntryIt->second;
-  std::string Contents = EntryIt->second.Contents;
+  DraftAndTime &D = EntryIt->second;
+  std::string Contents = *D.Draft.Contents;
 
   for (const TextDocumentContentChangeEvent &Change : Changes) {
     if (!Change.range) {
       Contents = Change.text;
       continue;
     }
 
     const Position &Start = Change.range->start;
for (const TextDocumentContentChangeEvent &Change : Changes) {
 
     NewContents = Contents.substr(0, *StartIndex);
     NewContents += Change.text;
     NewContents += Contents.substr(*EndIndex);
 
     Contents = std::move(NewContents);
   }
 
-  updateVersion(D, Version);
-  D.Contents = std::move(Contents);
-  return D;
+  updateVersion(D.Draft, Version);
+  updateTime(D.MTime);
+  D.Draft.Contents =
+      llvm::makeIntrusiveRefCnt<RefCntString>(std::move(Contents));
+  return D.Draft;
 }
 
 void DraftStore::removeDraft(PathRef File) {
   std::lock_guard<std::mutex> Lock(Mutex);
 
   Drafts.erase(File);
 }
 
+namespace {
+class RefCntStringMemBuffer final : public llvm::MemoryBuffer {
+public:
+  static std::unique_ptr<RefCntStringMemBuffer>
+  create(IntrusiveRefCntPtr<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<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<RefCntString> BufferContents;
+  size_t NameSize;
+};
+class DraftStoreFile : public llvm::vfs::File {
+public:
+  struct FileData {
+    IntrusiveRefCntPtr<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().Draft.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

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

 ///     // implement other methods here.
 ///   };
 ///   REGISTER_TWEAK(MyTweak);
 class Tweak {
 public:
   /// 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.
     const SymbolIndex *Index = nullptr;
     /// The parsed active file. Never null. (Pointer so Selection is movable).
     ParsedAST *AST;
     /// A location of the cursor in the editor.
     // FIXME: Cursor is redundant and should be removed
     SourceLocation Cursor;
     /// The begin offset of the selection
     unsigned SelectionBegin;
     /// The end offset of the selection
     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.
   };
 
   struct Effect {
     /// A message to be displayed to the user.
     llvm::Optional<std::string> ShowMessage;
     FileEdits ApplyEdits;
 

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

for (const auto &E : TweakRegistry::entries()) {
     assert(Seen.try_emplace(E.getName()).second && "duplicate check id");
   }
 #endif
 }
 } // namespace
 
 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);
 }
 
 std::vector<std::unique_ptr<Tweak>>
 prepareTweaks(const Tweak::Selection &S,
               llvm::function_ref<bool(const Tweak &)> Filter) {

clang-tools-extra/clangd/refactor/tweaks/DefineOutline.cpp

if (AstNode.get<CompoundStmt>() &&
       SelNode->Selected == SelectionTree::Complete) {
     if (const SelectionTree::Node *P = SelNode->Parent)
       return P->ASTNode.get<FunctionDecl>();
   }
   return nullptr;
 }
 
 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);
 }
 
 // Synthesize a DeclContext for TargetNS from CurContext. TargetNS must be empty
 // for global namespace, and endwith "::" otherwise.
 // Returns None if TargetNS is not a prefix of CurContext.
 llvm::Optional<const DeclContext *>
public:
 
   Expected<Effect> apply(const Selection &Sel) override {
     const SourceManager &SM = Sel.AST->getSourceManager();
     auto MainFileName =
         getCanonicalPath(SM.getFileEntryForID(SM.getMainFileID()), SM);
     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)
       return llvm::errorCodeToError(Buffer.getError());
     auto Contents = Buffer->get()->getBuffer();
     auto InsertionPoint = getInsertionPoint(
         Contents, Source->getQualifiedNameAsString(), Sel.AST->getLangOpts());
     if (!InsertionPoint)

clang-tools-extra/clangd/tool/Check.cpp

for (const auto &Tok : SpelledTokens) {
       unsigned Start = AST->getSourceManager().getFileOffset(Tok.location());
       unsigned End = Start + Tok.length();
       Position Pos = offsetToPosition(Inputs.Contents, Start);
       // FIXME: dumping the tokens may leak sensitive code into bug reports.
       // Add an option to turn this off, once we decide how options work.
       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) {
           elog("    tweak: {0} ==> FAIL: {1}", T->id(), Result.takeError());
           ++ErrCount;
         } else {
           vlog("    tweak: {0}", T->id());
         }

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

TextDocumentContentChangeEvent Event{
         SrcBefore.range(),
         rangeLength(SrcBefore.code(), SrcBefore.range()),
         Contents.str(),
     };
 
     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));
   }
 }
 
 /// Send all the changes at once to updateDraft, check only the final result.
 void allAtOnce(llvm::ArrayRef<IncrementalTestStep> Steps) {
   DraftStore DS;
   Annotations InitialSrc(Steps.front().Src);
void allAtOnce(llvm::ArrayRef<IncrementalTestStep> Steps) {
 
   // Set the initial content.
   EXPECT_EQ(0, DS.addDraft(Path, llvm::None, InitialSrc.code()));
 
   llvm::Expected<DraftStore::Draft> Result =
       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);
 }
 
 TEST(DraftStoreIncrementalUpdateTest, Simple) {
   // clang-format off
   IncrementalTestStep Steps[] =
     {
       // Replace a range
TEST(DraftStoreIncrementalUpdateTest, InvalidRangeInASequence) {
   auto Result = DS.updateDraft(File, llvm::None, {Change1, Change2});
 
   EXPECT_TRUE(!Result);
   EXPECT_EQ(toString(Result.takeError()),
             "utf-16 offset 100 is invalid for line 0");
 
   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);
 }
 
 TEST(DraftStore, Version) {
   DraftStore DS;
   Path File = "foo.cpp";
 
   EXPECT_EQ(25, DS.addDraft(File, 25, ""));

clang-tools-extra/clangd/unittests/tweaks/TweakTesting.cpp

 llvm::Optional<llvm::Expected<Tweak::Effect>>
 applyTweak(ParsedAST &AST, const Annotations &Input, StringRef TweakID,
            const SymbolIndex *Index) {
   auto Range = rangeOrPoint(Input);
   llvm::Optional<llvm::Expected<Tweak::Effect>> Result;
   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;
                               } else {
                                 llvm::consumeError(T.takeError());
                                 return false;
                               }
                             });