[clangd] Add a Filesystem that overlays Dirty files.

[njames93]

Summary

Create a ThreadsafeFS in the DraftStore that overlays the dirty file contents over another ThreadsafeFS.
This provides a nice thread-safe interface for using dirty file contents throughout the codebase, for example cross file refactoring.
Creating a Filesystem view will overlay a snapshot of the current contents, so if the draft store is updated while the view is being used, it will contain stale contents.

[njames93]

Remove some unrelated changes

[njames93]

Refactor RefCntString out of DraftStore class
Remove MTime from Draft, other users of DraftStore don't need to see it.

[njames93]

Start Splitting this patch up

[njames93]

Rebase

[njames93]

Ping?

[sammccall]

Sorry for the long delay here. I've gotten myself buried in too many different things.

Overall the idea of exposing drafts as an overlay FS is definitely growing on me - it makes it easier and clearer to control which feature is seeing which version of the files.

At a high level my main requests are:

• we should move maintenance of the drafts from ClangdLSPServer into ClangdServer. (This should maybe have been done a while ago, but this change makes it more complicated and therefore pressing)
• we should simplify the implementation as much as possible, and lean on existing infrastructure - there's a lot of custom pieces here like RefCntString and DraftStoreFileSystem that seem like small optimizations over existing types

(I might take a crack at moving DraftStore by itself, as there probably are issues and I want to understand what they are...)

clang-tools-extra/clangd/ClangdServer.h

176–181

Passing in FS + DirtyFS doesn't seem quite right, because we also pass in all the mutations that create the differences.

It seem rather that *ClangdServer* should maintain the overlay containing the dirty buffers, and expose it (for use in the few places that ClangdLSPServer currently uses DraftStore).

(the fact that DraftStore isn't accessible from ClangdServer did also come up when trying to split the *Server classes into Modules, so this seems like a helpful direction from that perspective too)

clang-tools-extra/clangd/DraftStore.cpp

130

These classes need comments describing why they exist :-)
(I think you're right that buffers may outlive the VFS)

133

Is there any reason to think we need this micro-optimization? It's quite invasive.

198

can't we use InMemoryFileSystem for this?

272

doing IO in view() doesn't seem obviously OK.

what's the practical consequence of the overlay's inodes not matching that of the underlying FS? (It seems reasonable to say that the files always have different identity, and may/may not have the same content)

clang-tools-extra/clangd/DraftStore.h

25

shared_ptr<const string> seems like a cheap alternative here.

However if the point is ultimately to obtain MemoryBuffers that we can hand out with the underlying data being refcounted, can we do that directly?

class SharedBuffer : public MemoryBuffer {
  std::shared_ptr<const std::string> Data;
public:
  Shared(StringRef Name, std::string Data);
  Shared(const RefcountBuffer &); // copyable
};

31

providing non-explicit string conversion operators is IME very likely to cause ambiguous overload errors when using this in variaous places :-(

81

avoid assert in headers for ODR reasons (ooh, we have a couple of violations i should clean up)

[njames93]

clang-tools-extra/clangd/ClangdServer.h

176–181

I was never a fan of this tbh. But moving DraftStore would result in this being too big,

clang-tools-extra/clangd/DraftStore.cpp

133

Its basically copied from llvm/support/MemoryBuffer.cpp

198

InMemoryFileSystem Would work in theory. However using that requires building a lot of unnecessary nodes for thing like intermediate directories.

272

It's probably not necessary, but I know preambles use the UniqueID. It may just be safer to create a uniqueID for each item in the DraftStore and leave it at that.

[njames93]

Address some comments, though I have a feeling @sammccall, I may have to wait until DraftStore has been refactored first before continuing on with this.

[sammccall]

In D94554#2595625, @njames93 wrote:

Address some comments, though I have a feeling @sammccall, I may have to wait until DraftStore has been refactored first before continuing on with this.

D97738 is an attempt at this. It is... sprawly, but it's an overdue cleanup that should fit together well with this patch, I think.

There may be some conflicts but I don't think it'll be terrible actually... I think we should keep going here, it looks really good.

clang-tools-extra/clangd/DraftStore.cpp

133

Sure, but this doesn't really answer the question! Can we keep this trivial until we have a reason to optimize it? (Clang typically holds tens of thousands of MemoryBuffers, but we're only likely to have a few hundred of these)

198

Probably, but it would keep this code simpler, supports directory iteration and status (used in code completion), and I'm fairly confident its path handling is solid.

It's certainly possible to add this optimization later if it seems important, but in this patch I think it distracts from the design questions.

272

The biggest risks I can see with this approach:

• in build-preamble-from-dirty-FS mode, would opening a draft cause the preamble to be considered invalid and needing rebuilding? My read of PrecompiledPreamble::CanReuse is no (UniqueIDs are not stored)
• when parsing, if the file can be read through the overlay *and* underlying via different paths (e.g. hardlinks), these won't be deduplicated. However giving them the same inode only changes the nature of the bug: the file would get whichever content was read first.

So I think it's as good as we can expect.

clang-tools-extra/clangd/DraftStore.h

49

having DraftStore sit on top of TFS seems a bit inside-out, giving it bigger scope than necessary.

What about giving DraftStore an asVFS() method that returns an in-memory filesystem?

then separately we can have a separate DirtyFS : TFS, that has a DraftStore& and a TFS &Base and implements viewImpl() on top of their public interfaces. Having the dependency in that direction seems more natural to me.

86

alternatively you can define whichever internal struct you like, like

struct FileData {
  std::string Filename;
  std::string Contents;
  // and MTime and UID
}

and store DenseMap<StringRef, shared_ptr<FileData>> (with the keys pointing into the values).

You can still hand out shared_ptr<const std::string> using the aliasing constructor of shared_ptr.

The advantages of this:

• it avoids the internal representation from being too constrained by the public API
• SharedStringBuffer can easily wrap the whole shared_ptr<FileData>, and so you doesn't need a *copy* of the filename

this is a little bit of complexity though, up to you

[njames93]

Getting there.

[njames93]

clang-tools-extra/clangd/DraftStore.h

49

How I had it was a symptom of using the underlying TFS for getting the UniqueID of files. But yes now that's gone this approach works.

[njames93]

clang-tools-extra/clangd/ClangdServer.h

415

Probably needs moving to its own place, but wasn't sure where best to put it.

clang-tools-extra/clangd/DraftStore.cpp

272

Unfortunately in the current implementation, opening a draft will invalidate and preambles that use the underlying file.
This is due to the last modification times not matching.
This could be addressed by querying the TFS on the textDocument/didOpen request to get the actual last modification time, but we shouldn't really be doing any kind of IO on that thread.

[sammccall]

Thanks, LG!

clang-tools-extra/clangd/ClangdServer.h

415

I think ClangdServer is fine, but this can be a unique_ptr<ThreadsafeFS> with the class moved to the implementation file

clang-tools-extra/clangd/DraftStore.cpp

57

we don't need try_emplace here, we overwrite the whole value regardless.
Can we use Draft &D = Drafts[File] again?

60

nit: std::time() (and std::time_t)

147

nit: /*RequiresNullTerminator=*/

[njames93]

clang-tools-extra/clangd/ClangdServer.h

415

How about privately inheriting from ThreadsafeFS :P
I agree, hiding the impl in the cpp is probably a much nicer way.

clang-tools-extra/clangd/DraftStore.cpp

57

I changed this because in an older version I wanted to set the UniqueID on the first time the file was added, however now we don't do that (or check the return value for insertion at all) its safe to go back to the operator[] access.

272

@sammccall Any further comments with regard to this. Is calling IO to get modification time in didOpen ok, or just accept opening a file would invalidate a preamble if the DirtyFS was used for preambles.

[njames93]

Address comments.

[sammccall]

clang-tools-extra/clangd/DraftStore.cpp

272

There's no urgency in resolving this:

• preambles-from-dirty-fs doesn't exist in this patch
• even once it does, we can introduce the problem first (in this experimental mode) and then solve it

That said, my suspicion is that having ClangdLSPServer stat the file, conditional on being in preambles-from-dirty-fs mode, would be OK.
(If the filesize matches I think it's safe to say the content does, and copy the timestamp. If the filesize mismatches... well, there are a number of scenarios involving encodings, and it's hard to say which will matter in practice, so we'd probably want to log and wait for the bug reports)

[sammccall]

Still LG btw!

clang-tools-extra/clangd/ClangdServer.h

350–351

Maybe add an explicit comment that this returns nullptr if the file is not open.
(When the type was Optional this was fairly self-explanatory)

415

nit: we don't use const on pointer members, just on the pointee where appropriate.
Maybe not the best style in a vacuum but it's fairly consistent and diverging from it is a bit confusing.

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

struct Options {
     std::vector<std::string> QueryDriverGlobs;
 
     /// Enable notification-based semantic highlighting.
     bool TheiaSemanticHighlighting = false;
 
     /// 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.
   // Features like indexing must be enabled if desired.
   static Options optsForTest();
 
   /// 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,
public:
   /// Returns estimated memory usage and other statistics for each of the
   /// currently open files.
   /// Overall memory usage of clangd may be significantly more than reported
   /// here, as this metric does not account (at least) for:
   ///   - memory occupied by static and dynamic index,
   ///   - memory required for in-flight requests,
   /// FIXME: those metrics might be useful too, we should add them.
   llvm::StringMap<TUScheduler::FileStats> fileStats() const;
 
   // Blocks the main thread until the server is idle. Only for use in tests.
   // Returns false if the timeout expires.
   LLVM_NODISCARD bool
   blockUntilIdleForTest(llvm::Optional<double> TimeoutSeconds = 10);
 
   /// Builds a nested representation of memory used by components.
   void profile(MemoryTree &MT) const;
 
 private:
   void formatCode(PathRef File, llvm::StringRef Code,
                   ArrayRef<tooling::Range> Ranges,
                   Callback<tooling::Replacements> CB);
 
   /// Derives a context for a task processing the specified source file.
   /// This includes the current configuration (see Options::ConfigProvider).
   /// 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
private:
   // can be caused by missing includes (e.g. member access in incomplete type).
   bool SuggestMissingIncludes = false;
 
   // If true, preserve expressions in AST for broken code.
   bool BuildRecoveryAST = true;
   // If true, preserve the type for recovery AST.
   bool PreserveRecoveryASTType = false;
 
+  bool UseDirtyPreambles = false;
+
   // GUARDED_BY(CachedCompletionFuzzyFindRequestMutex)
   llvm::StringMap<llvm::Optional<FuzzyFindRequest>>
       CachedCompletionFuzzyFindRequestByFile;
   mutable std::mutex CachedCompletionFuzzyFindRequestMutex;
 
   llvm::Optional<std::string> WorkspaceRoot;
   // WorkScheduler has to be the last member, because its destructor has to be
   // called before all other members to stop the worker thread that references
   // ClangdServer.
   TUScheduler WorkScheduler;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

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),
+      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
       // is parsed.
       // FIXME(ioeric): this can be slow and we may be able to index on less
       // critical paths.
       WorkScheduler(
           CDB,

 void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
                                llvm::StringRef Version,
                                WantDiagnostics WantDiags, bool ForceRebuild) {
   ParseOptions Opts;
   Opts.SuggestMissingIncludes = SuggestMissingIncludes;
 
   // 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;
   Inputs.Opts = std::move(Opts);
   Inputs.Index = Index;
   Inputs.ClangTidyProvider = ClangTidyProvider;
   Inputs.Opts.BuildRecoveryAST = BuildRecoveryAST;
   Inputs.Opts.PreserveRecoveryASTType = PreserveRecoveryASTType;
if (!IP->Preamble) {
         {
           std::lock_guard<std::mutex> Lock(
               CachedCompletionFuzzyFindRequestMutex);
           SpecFuzzyFind->CachedReq =
               CachedCompletionFuzzyFindRequestByFile[File];
         }
       }
     }
-    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;
 
     // FIXME(ibiryukov): even if Preamble is non-null, we may want to check
     // both the old and the new version in case only one of them matches.
     CodeCompleteResult Result = clangd::codeComplete(
         File, Pos, IP->Preamble, ParseInput, CodeCompleteOpts,
auto Action = [Pos, File = File.str(), CB = std::move(CB),
                  this](llvm::Expected<InputsAndPreamble> IP) mutable {
     if (!IP)
       return CB(IP.takeError());
 
     const auto *PreambleData = IP->Preamble;
     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;
     CB(clangd::signatureHelp(File, Pos, *PreambleData, ParseInput));
   };
 
   // Unlike code completion, we wait for a preamble here.
   WorkScheduler.runWithPreamble("SignatureHelp", File, TUScheduler::Stale,
auto Action = [Pos, File = File.str(), CB = std::move(CB),
     //  - for single-file rename, performing rename isn't substantially more
     //    expensive than doing an AST-based check (the index is used to see if
     //    the rename is complete);
     //  - for cross-file rename, we deliberately pass a nullptr index to save
     //    the cost, thus the result may be incomplete as it only contains
     //    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
       // failure message. If we send an LSP error response, clients can surface
       // the message to users (VSCode does).
       return CB(Results.takeError());
     }
     return CB(*Results);
   };
   WorkScheduler.runWithAST("PrepareRename", File, std::move(Action));
 }
 
 void ClangdServer::rename(PathRef File, Position Pos, llvm::StringRef NewName,
                           const RenameOptions &Opts,
                           Callback<RenameResult> CB) {
-  // A snapshot of all file dirty buffers.
-  llvm::StringMap<std::string> Snapshot = WorkScheduler.getAllFileContents();
   auto Action = [File = File.str(), NewName = NewName.str(), Pos, Opts,
-                 CB = std::move(CB), Snapshot = std::move(Snapshot),
+                 CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     // Tracks number of files edited per invocation.
     static constexpr trace::Metric RenameFiles("rename_files",
                                                trace::Metric::Distribution);
     if (!InpAST)
       return CB(InpAST.takeError());
-    auto GetDirtyBuffer =
-        [&Snapshot](PathRef AbsPath) -> llvm::Optional<std::string> {
-      auto It = Snapshot.find(AbsPath);
-      if (It == Snapshot.end())
-        return llvm::None;
-      return It->second;
-    };
-    auto R = clangd::rename(
-        {Pos, NewName, InpAST->AST, File, Index, Opts, GetDirtyBuffer});
+    auto R = clangd::rename({Pos, NewName, InpAST->AST, File,
+                             DirtyFS.view(llvm::None), Index, Opts});
     if (!R)
       return CB(R.takeError());
 
     if (Opts.WantFormat) {
       auto Style = getFormatStyleForFile(File, InpAST->Inputs.Contents,
                                          *InpAST->Inputs.TFS);
       llvm::Error Err = llvm::Error::success();
       for (auto &E : R->GlobalChanges)
         Err =
             llvm::joinErrors(reformatEdit(E.getValue(), Style), std::move(Err));
 
       if (Err)
         return CB(std::move(Err));
     }
     RenameFiles.record(R->GlobalChanges.size());
     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) {
+    for (auto &Selection : *Selections) {
       auto T = prepareTweak(TweakID, *Selection);
       if (T) {
         Effect = (*T)->apply(*Selection);
         break;
       }
       Effect = T.takeError();
     }
     assert(Effect.hasValue() && "Expected at least one selection");

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 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;
+    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;
 
   /// \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:
+  class DraftstoreFS;
+  std::unique_ptr<ThreadsafeFS> DraftFS;
   mutable std::mutex Mutex;
   llvm::StringMap<Draft> Drafts;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/clangd/DraftStore.cpp

 }
 
 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();
+  D.Contents = llvm::makeIntrusiveRefCnt<Draft::RefCntString>(Contents.str());
   return D.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;
+  std::string Contents = *EntryIt->second.Contents;
 
   for (const TextDocumentContentChangeEvent &Change : Changes) {
     if (!Change.range) {
       Contents = Change.text;
       continue;
     }
 
     const Position &Start = Change.range->start;
ssize_t ComputedRangeLength =
         lspLength(Contents.substr(*StartIndex, *EndIndex - *StartIndex));
 
     if (Change.rangeLength && ComputedRangeLength != *Change.rangeLength)
       return error(llvm::errc::invalid_argument,
                    "Change's rangeLength ({0}) doesn't match the "
                    "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;
 }
 
 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<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

clang-tools-extra/clangd/TUScheduler.h

public:
   /// Returns true if the file was not previously tracked.
   bool update(PathRef File, ParseInputs Inputs, WantDiagnostics WD);
 
   /// Remove \p File from the list of tracked files and schedule removal of its
   /// resources. Pending diagnostics for closed files may not be delivered, even
   /// if requested with WantDiags::Auto or WantDiags::Yes.
   void remove(PathRef File);
 
-  /// Returns a snapshot of all file buffer contents, per last update().
-  llvm::StringMap<std::string> getAllFileContents() const;
-
   /// Schedule an async task with no dependencies.
   /// Path may be empty (it is used only to set the Context).
   void run(llvm::StringRef Name, llvm::StringRef Path,
            llvm::unique_function<void()> Action);
 
   /// Defines how a runWithAST action is implicitly cancelled by other actions.
   enum ASTActionInvalidation {
     /// The request will run unless explicitly cancelled.

clang-tools-extra/clangd/TUScheduler.cpp

 
 void TUScheduler::remove(PathRef File) {
   bool Removed = Files.erase(File);
   if (!Removed)
     elog("Trying to remove file from TUScheduler that is not tracked: {0}",
          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) {
     WithContext WithProvidedContext(Opts.ContextProvider(Path));
     return Action();
   }
   PreambleTasks->runAsync(Name, [this, Ctx = Context::current().clone(),
                                  Path(Path.str()),

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

 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/Support/Error.h"
 
 namespace clang {
 namespace clangd {
 class ParsedAST;
 class SymbolIndex;
 
-/// Gets dirty buffer for a given file \p AbsPath.
-/// Returns None if there is no dirty buffer for the given file.
-using DirtyBufferGetter =
-    llvm::function_ref<llvm::Optional<std::string>(PathRef AbsPath)>;
-
 struct RenameOptions {
   /// If true, enable cross-file rename; otherwise, only allows to rename a
   /// symbol that's only used in the current file.
   bool AllowCrossFile = false;
   /// The maximum number of affected files (0 means no limit), only meaningful
   /// when AllowCrossFile = true.
   /// If the actual number exceeds the limit, rename is forbidden.
   size_t LimitFiles = 50;
   /// If true, format the rename edits, only meaningful in ClangdServer layer.
   bool WantFormat = false;
 };
 
 struct RenameInputs {
   Position Pos; // the position triggering the rename
   llvm::StringRef NewName;
 
   ParsedAST &AST;
   llvm::StringRef MainFilePath;
 
+  // The filesystem to query when performing cross file renames.
+  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 {
   // The range of the symbol that the user can attempt to rename.
   Range Target;
   // Rename occurrences for the current main file.
   std::vector<Range> LocalChanges;
   // Complete edits for the rename, including LocalChanges.

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

 // We choose to trade off some correctness for performance and scalability.
 //
 // Clangd builds a dynamic index for all opened files on top of the static
 // index of the whole codebase. Dynamic index is up-to-date (respects dirty
 // buffers) as long as clangd finishes processing opened files, while static
 // index (background index) is relatively stale. We choose the dirty buffers
 // as the file content we rename on, and fallback to file content on disk if
 // there is no dirty buffer.
-llvm::Expected<FileEdits> renameOutsideFile(
-    const NamedDecl &RenameDecl, llvm::StringRef MainFilePath,
-    llvm::StringRef NewName, const SymbolIndex &Index, size_t MaxLimitFiles,
-    llvm::function_ref<llvm::Expected<std::string>(PathRef)> GetFileContent) {
+llvm::Expected<FileEdits>
+renameOutsideFile(const NamedDecl &RenameDecl, llvm::StringRef MainFilePath,
+                  llvm::StringRef NewName, const SymbolIndex &Index,
+                  size_t MaxLimitFiles, llvm::vfs::FileSystem &FS) {
   trace::Span Tracer("RenameOutsideFile");
   auto AffectedFiles = findOccurrencesOutsideFile(RenameDecl, MainFilePath,
                                                   Index, MaxLimitFiles);
   if (!AffectedFiles)
     return AffectedFiles.takeError();
   FileEdits Results;
   for (auto &FileAndOccurrences : *AffectedFiles) {
     llvm::StringRef FilePath = FileAndOccurrences.first();
 
-    auto AffectedFileCode = GetFileContent(FilePath);
-    if (!AffectedFileCode) {
-      elog("Fail to read file content: {0}", AffectedFileCode.takeError());
+    auto ExpBuffer = FS.getBufferForFile(FilePath);
+    if (!ExpBuffer) {
+      elog("Fail to read file content: Fail to open file {0}: {1}", FilePath,
+           ExpBuffer.getError().message());
       continue;
     }
+
+    auto AffectedFileCode = (*ExpBuffer)->getBuffer();
     auto RenameRanges =
-        adjustRenameRanges(*AffectedFileCode, RenameDecl.getNameAsString(),
+        adjustRenameRanges(AffectedFileCode, RenameDecl.getNameAsString(),
                            std::move(FileAndOccurrences.second),
                            RenameDecl.getASTContext().getLangOpts());
     if (!RenameRanges) {
       // Our heuristics fails to adjust rename ranges to the current state of
       // the file, it is most likely the index is stale, so we give up the
       // entire rename.
       return error("Index results don't match the content of file {0} "
                    "(the index may be stale)",
                    FilePath);
     }
     auto RenameEdit =
-        buildRenameEdit(FilePath, *AffectedFileCode, *RenameRanges, NewName);
+        buildRenameEdit(FilePath, AffectedFileCode, *RenameRanges, NewName);
     if (!RenameEdit)
       return error("failed to rename in file {0}: {1}", FilePath,
                    RenameEdit.takeError());
     if (!RenameEdit->Replacements.empty())
       Results.insert({FilePath, std::move(*RenameEdit)});
   }
   return Results;
 }

 } // namespace
 
 llvm::Expected<RenameResult> rename(const RenameInputs &RInputs) {
   trace::Span Tracer("Rename flow");
   const auto &Opts = RInputs.Opts;
   ParsedAST &AST = RInputs.AST;
   const SourceManager &SM = AST.getSourceManager();
   llvm::StringRef MainFileCode = SM.getBufferData(SM.getMainFileID());
-  auto GetFileContent = [&RInputs,
-                         &SM](PathRef AbsPath) -> llvm::Expected<std::string> {
-    llvm::Optional<std::string> DirtyBuffer;
-    if (RInputs.GetDirtyBuffer &&
-        (DirtyBuffer = RInputs.GetDirtyBuffer(AbsPath)))
-      return std::move(*DirtyBuffer);
-
-    auto Content =
-        SM.getFileManager().getVirtualFileSystem().getBufferForFile(AbsPath);
-    if (!Content)
-      return error("Fail to open file {0}: {1}", AbsPath,
-                   Content.getError().message());
-    if (!*Content)
-      return error("Got no buffer for file {0}", AbsPath);
-
-    return (*Content)->getBuffer().str();
-  };
   // Try to find the tokens adjacent to the cursor position.
   auto Loc = sourceLocationInMainFile(SM, RInputs.Pos);
   if (!Loc)
     return Loc.takeError();
   const syntax::Token *IdentifierToken =
       spelledIdentifierTouching(*Loc, AST.getTokens());
 
   // Renames should only triggered on identifiers.
if (!RInputs.Index) {
     assert(Result.GlobalChanges.empty() && Opts.AllowCrossFile);
     return Result;
   }
 
   auto OtherFilesEdits = renameOutsideFile(
       RenameDecl, RInputs.MainFilePath, RInputs.NewName, *RInputs.Index,
       Opts.LimitFiles == 0 ? std::numeric_limits<size_t>::max()
                            : Opts.LimitFiles,
-      GetFileContent);
+      *RInputs.FS);
   if (!OtherFilesEdits)
     return OtherFilesEdits.takeError();
   Result.GlobalChanges = *OtherFilesEdits;
   // Attach the rename edits for the main file.
   Result.GlobalChanges.try_emplace(RInputs.MainFilePath,
                                    std::move(MainFileEdits));
   return Result;
 }

clang-tools-extra/clangd/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/tool/ClangdMain.cpp

 
 opt<bool> CollectMainFileRefs{
     "collect-main-file-refs",
     cat(Misc),
     desc("Store references to main-file-only symbols in the index"),
     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",
     cat(Misc),
     desc("Release memory periodically via malloc_trim(3)."),
     init(true),
 };
 
if (!ClangTidyChecks.empty())
       Providers.push_back(addTidyChecks(ClangTidyChecks));
     else
       Providers.push_back(provideDefaultChecks());
     Providers.push_back(disableUnusableChecks());
     ClangTidyOptProvider = combine(std::move(Providers));
     Opts.ClangTidyProvider = ClangTidyOptProvider;
   }
   Opts.AsyncPreambleBuilds = AsyncPreamble;
+  Opts.UseDirtyPreambles = UseDirtyPreambles;
   Opts.SuggestMissingIncludes = SuggestMissingIncludes;
   Opts.QueryDriverGlobs = std::move(QueryDriverGlobs);
   Opts.TweakFilter = [&](const Tweak &T) {
     if (T.hidden() && !HiddenFeatures)
       return false;
     if (TweakList.getNumOccurrences())
       return llvm::is_contained(TweakList, T.id());
     return true;

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/RenameTests.cpp

 using testing::ElementsAre;
 using testing::Eq;
 using testing::IsEmpty;
 using testing::Pair;
 using testing::SizeIs;
 using testing::UnorderedElementsAre;
 using testing::UnorderedElementsAreArray;
 
+llvm::IntrusiveRefCntPtr<llvm::vfs::OverlayFileSystem>
+createOverlay(llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> Base,
+              llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> Overlay) {
+  auto OFS =
+      llvm::makeIntrusiveRefCnt<llvm::vfs::OverlayFileSystem>(std::move(Base));
+  OFS->pushOverlay(std::move(Overlay));
+  return OFS;
+}
+
+llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> getVFSFromAST(ParsedAST &AST) {
+  return &AST.getSourceManager().getFileManager().getVirtualFileSystem();
+}
+
 // Convert a Range to a Ref.
 Ref refWithRange(const clangd::Range &Range, const std::string &URI) {
   Ref Result;
   Result.Kind = RefKind::Reference | RefKind::Spelled;
   Result.Location.Start.setLine(Range.start.line);
   Result.Location.Start.setColumn(Range.start.character);
   Result.Location.End.setLine(Range.end.line);
   Result.Location.End.setColumn(Range.end.character);
TEST(RenameTest, WithinFileRename) {
   for (llvm::StringRef T : Tests) {
     SCOPED_TRACE(T);
     Annotations Code(T);
     auto TU = TestTU::withCode(Code.code());
     TU.ExtraArgs.push_back("-fno-delayed-template-parsing");
     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(
           applyEdits(std::move(RenameResult->GlobalChanges)).front().second,
           expectedResult(Code, NewName));
     }
   }
 }
for (const auto& Case : Cases) {
     if (Case.IsHeaderFile) {
       // We open the .h file as the main file.
       TU.Filename = "test.h";
       // Parsing the .h file as C++ include.
       TU.ExtraArgs.push_back("-xobjective-c++-header");
     }
     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;
     if (!WantRename) {
       assert(Case.ErrorMessage && "Error message must be set!");
       EXPECT_FALSE(Results)
           << "expected rename returned an error: " << T.code();
       auto ActualMessage = llvm::toString(Results.takeError());
TEST(RenameTest, MainFileReferencesOnly) {
   TU.AdditionalFiles["foo.inc"] = R"cpp(
       #define Macro(X) X
       &Macro(foo);
       &foo;
     )cpp";
   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,
             expectedResult(Code, NewName));
 }
 
 TEST(RenameTest, ProtobufSymbolIsExcluded) {
   Annotations Code("Prot^obuf buf;");
   auto TU = TestTU::withCode(Code.code());
   TU.HeaderCode =
       R"cpp(// Generated by the protocol buffer compiler.  DO NOT EDIT!
       class Protobuf {};
       )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"));
 }
 
 TEST(RenameTest, PrepareRename) {
   Annotations FooH("void func();");
   Annotations FooCC(R"cpp(
TEST(CrossFileRenameTests, DirtyBuffer) {
   FSymbols.update(FooPath, nullptr, buildRefSlab(FooCode, "Foo", FooPath),
                   nullptr, false);
   FSymbols.update(BarPath, nullptr, buildRefSlab(BarCode, "Bar", BarPath),
                   nullptr, false);
   auto Index = FSymbols.buildIndex(IndexType::Light);
 
   Annotations MainCode("class  [[Fo^o]] {};");
   auto MainFilePath = testPath("main.cc");
-  // Dirty buffer for foo.cc.
-  auto GetDirtyBuffer = [&](PathRef Path) -> llvm::Optional<std::string> {
-    if (Path == FooPath)
-      return FooDirtyBuffer.code().str();
-    return llvm::None;
-  };
+  llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> InMemFS =
+      new llvm::vfs::InMemoryFileSystem;
+  InMemFS->addFile(FooPath, 0,
+                   llvm::MemoryBuffer::getMemBuffer(FooDirtyBuffer.code()));
 
   // Run rename on Foo, there is a dirty buffer for foo.cc, rename should
   // respect the dirty buffer.
   TestTU TU = TestTU::withCode(MainCode.code());
   auto AST = TU.build();
   llvm::StringRef NewName = "newName";
   auto Results = rename({MainCode.point(),
                          NewName,
                          AST,
                          MainFilePath,
+                         createOverlay(getVFSFromAST(AST), InMemFS),
                          Index.get(),
-                         {/*CrossFile=*/true},
-                         GetDirtyBuffer});
+                         {/*CrossFile=*/true}});
   ASSERT_TRUE(bool(Results)) << Results.takeError();
   EXPECT_THAT(
       applyEdits(std::move(Results->GlobalChanges)),
       UnorderedElementsAre(
           Pair(Eq(FooPath), Eq(expectedResult(FooDirtyBuffer, NewName))),
           Pair(Eq(MainFilePath), Eq(expectedResult(MainCode, NewName)))));
 
   // Run rename on Bar, there is no dirty buffer for the affected file bar.cc,
   // so we should read file content from VFS.
   MainCode = Annotations("void [[Bar]]() { [[B^ar]](); }");
   TU = TestTU::withCode(MainCode.code());
   // Set a file "bar.cc" on disk.
   TU.AdditionalFiles["bar.cc"] = std::string(BarCode.code());
   AST = TU.build();
   Results = rename({MainCode.point(),
                     NewName,
                     AST,
                     MainFilePath,
+                    createOverlay(getVFSFromAST(AST), InMemFS),
                     Index.get(),
-                    {/*CrossFile=*/true},
-                    GetDirtyBuffer});
+                    {/*CrossFile=*/true}});
   ASSERT_TRUE(bool(Results)) << Results.takeError();
   EXPECT_THAT(
       applyEdits(std::move(Results->GlobalChanges)),
       UnorderedElementsAre(
           Pair(Eq(BarPath), Eq(expectedResult(BarCode, NewName))),
           Pair(Eq(MainFilePath), Eq(expectedResult(MainCode, NewName)))));
 
   // Run rename on a pagination index which couldn't return all refs in one
class PaginationIndex : public SymbolIndex {
     }
 
     size_t estimateMemoryUsage() const override { return 0; }
   } PIndex;
   Results = rename({MainCode.point(),
                     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"));
 }
 
 TEST(CrossFileRenameTests, DeduplicateRefsFromIndex) {
   auto MainCode = Annotations("int [[^x]] = 2;");
   auto MainFilePath = testPath("main.cc");
public:
     size_t estimateMemoryUsage() const override { return 0; }
     Ref ReturnedRef;
   } DIndex(XRefInBarCC);
   llvm::StringRef NewName = "newName";
   auto Results = rename({MainCode.point(),
                          NewName,
                          AST,
                          MainFilePath,
+                         getVFSFromAST(AST),
                          &DIndex,
                          {/*CrossFile=*/true}});
   ASSERT_TRUE(bool(Results)) << Results.takeError();
   EXPECT_THAT(
       applyEdits(std::move(Results->GlobalChanges)),
       UnorderedElementsAre(
           Pair(Eq(BarPath), Eq(expectedResult(BarCode, NewName))),
           Pair(Eq(MainFilePath), Eq(expectedResult(MainCode, NewName)))));

 TEST(CrossFileRenameTests, CrossFileOnLocalSymbol) {
   // cross-file rename should work for function-local symbols, even there is no
   // index provided.
   Annotations Code("void f(int [[abc]]) { [[a^bc]] = 3; }");
   auto TU = TestTU::withCode(Code.code());
   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)),
       UnorderedElementsAre(Pair(Eq(Path), Eq(expectedResult(Code, NewName)))));
 }
 
 TEST(CrossFileRenameTests, BuildRenameEdits) {
   Annotations Code("[[😂]]");

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;
                               }
                             });