[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

553

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

55

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

71

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.

71

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

72

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

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

415–416

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

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

214

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

55

We could, but it would messy code elsewhere.

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

214

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

584

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

584

Yes, I think so.

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

68

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

68

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

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),
                  Filter = std::move(Filter),
                  FeatureModules(this->FeatureModules)](
                     Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
-    auto Selections = tweakSelection(Sel, *InpAST);
+    auto Selections = tweakSelection(Sel, *InpAST, /*FS=*/nullptr);
     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, FeatureModules);
       if (T) {

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 is only populated when applying a tweak, not during prepare.
+    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 (auto &M : *Modules)
       M.contributeTweaks(All);
   }
   return All;
 }
 } // 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) {
   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,
               const FeatureModuleSet *Modules) {

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

if (AstNode.get<CompoundStmt>() &&
     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()))
+  assert(Sel.FS);
+  if (auto Source = getCorrespondingHeaderOrSource(FileName, Sel.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);
+
     if (!CCFile)
       return error("Couldn't find a suitable implementation file.");
-
-    auto &FS =
-        Sel.AST->getSourceManager().getFileManager().getVirtualFileSystem();
-    auto Buffer = FS.getBufferForFile(*CCFile);
+    assert(Sel.FS && "FS Must be set in apply");
+    auto Buffer = Sel.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) {
       if (!ShouldCheckLine(Pos))
         continue;
 
       // 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, Opts.FeatureModules)) {
         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/FeatureModulesTests.cpp

TEST(FeatureModulesTest, ContributesTweak) {
 
   FeatureModuleSet Set;
   Set.add(std::make_unique<TweakContributingModule>());
 
   auto AST = TestTU::withCode("").build();
   auto Tree =
       SelectionTree::createRight(AST.getASTContext(), AST.getTokens(), 0, 0);
   auto Actual = prepareTweak(
-      TweakID, Tweak::Selection(nullptr, AST, 0, 0, std::move(Tree)), &Set);
+      TweakID, Tweak::Selection(nullptr, AST, 0, 0, std::move(Tree), nullptr),
+      &Set);
   ASSERT_TRUE(bool(Actual));
   EXPECT_EQ(Actual->get()->id(), TweakID);
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

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

std::pair<unsigned, unsigned> rangeOrPoint(const Annotations &A) {
   return {cantFail(positionToOffset(A.code(), SelectionRng.start)),
           cantFail(positionToOffset(A.code(), SelectionRng.end))};
 }
 
 // Prepare and apply the specified tweak based on the selection in Input.
 // Returns None if and only if prepare() failed.
 llvm::Optional<llvm::Expected<Tweak::Effect>>
 applyTweak(ParsedAST &AST, const Annotations &Input, StringRef TweakID,
-           const SymbolIndex *Index) {
+           const SymbolIndex *Index, llvm::vfs::FileSystem *FS) {
   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),
+                                                 FS);
                               if (auto T = prepareTweak(TweakID, S, nullptr)) {
                                 Result = (*T)->apply(S);
                                 return true;
                               } else {
                                 llvm::consumeError(T.takeError());
                                 return false;
                               }
                             });
   return Result;
 }
 
 MATCHER_P7(TweakIsAvailable, TweakID, Ctx, Header, ExtraArgs, ExtraFiles, Index,
            FileName,
            (TweakID + (negation ? " is unavailable" : " is available")).str()) {
   std::string WrappedCode = wrap(Ctx, arg);
   Annotations Input(WrappedCode);
   TestTU TU;
   TU.Filename = std::string(FileName);
   TU.HeaderCode = Header;
   TU.Code = std::string(Input.code());
   TU.ExtraArgs = ExtraArgs;
   TU.AdditionalFiles = std::move(ExtraFiles);
   ParsedAST AST = TU.build();
-  auto Result = applyTweak(AST, Input, TweakID, Index);
+  auto Result = applyTweak(
+      AST, Input, TweakID, Index,
+      &AST.getSourceManager().getFileManager().getVirtualFileSystem());
   // We only care if prepare() succeeded, but must handle Errors.
   if (Result && !*Result)
     consumeError(Result->takeError());
   return Result.hasValue();
 }
 
 } // namespace
 
 std::string TweakTest::apply(llvm::StringRef MarkedCode,
                              llvm::StringMap<std::string> *EditedFiles) const {
   std::string WrappedCode = wrap(Context, MarkedCode);
   Annotations Input(WrappedCode);
   TestTU TU;
   TU.Filename = std::string(FileName);
   TU.HeaderCode = Header;
   TU.AdditionalFiles = std::move(ExtraFiles);
   TU.Code = std::string(Input.code());
   TU.ExtraArgs = ExtraArgs;
   ParsedAST AST = TU.build();
 
-  auto Result = applyTweak(AST, Input, TweakID, Index.get());
+  auto Result = applyTweak(
+      AST, Input, TweakID, Index.get(),
+      &AST.getSourceManager().getFileManager().getVirtualFileSystem());
   if (!Result)
     return "unavailable";
   if (!*Result)
     return "fail: " + llvm::toString(Result->takeError());
   const auto &Effect = **Result;
   if ((*Result)->ShowMessage)
     return "message:\n" + *Effect.ShowMessage;
   if (Effect.ApplyEdits.empty())