[clangd] Propogate FS into Tweak::Selection

[kadircet]

Summary

Incoming define out-of-line code action needs to access files that
weren't touched during compilation of the translation unit. This patch passes FS
into Tweak::Selection to make files accessible.

[ilya-biryukov]

This can be obtained from the file manager inside the AST.
I guess a helper function to do this would be helpful, but I don't think we need to strore another copy of it.
WDYT?

[kadircet]

As discussed offline, dropping the change to fetch the FS from filemanager, to prevent any consistency that might be caused due to having different contents for files in AST while creating it vs while generating refactorings.

clang-tools-extra/clangd/ClangdServer.cpp

 static llvm::Expected<Tweak::Selection>
 tweakSelection(const Range &Sel, const InputsAndAST &AST) {
   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();
-  return Tweak::Selection(AST.AST, *Begin, *End, AST.Inputs.Index);
+  return Tweak::Selection(AST.AST, *Begin, *End, AST.Inputs.Index,
+                          AST.Inputs.FS.get());
 }
 
 void ClangdServer::enumerateTweaks(PathRef File, Range Sel,
                                    Callback<std::vector<TweakRef>> CB) {
   auto Action = [File = File.str(), Sel, CB = std::move(CB),
                  this](Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());

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

 #include "Selection.h"
 #include "SourceCode.h"
 #include "index/Index.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/VirtualFileSystem.h"
 #include <string>
 
 namespace clang {
 namespace clangd {
 
 /// An interface base for small context-sensitive refactoring actions.
 /// To implement a new tweak use the following pattern in a .cpp file:
 ///   class MyTweak : public Tweak {
 ///   public:
 ///     const char* id() const override final; // defined by REGISTER_TWEAK.
 ///     // implement other methods here.
 ///   };
 ///   REGISTER_TWEAK(MyTweak);
 class Tweak {
 public:
   /// Input to prepare and apply tweaks.
   struct Selection {
     Selection(ParsedAST &AST, unsigned RangeBegin, unsigned RangeEnd,
-              const SymbolIndex *Index);
+              const SymbolIndex *Index, llvm::vfs::FileSystem *FS);
     /// The text of the active document.
     llvm::StringRef Code;
     /// Parsed AST of the active file.
     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;
     /// The Index being used by ClangdServer.
     const SymbolIndex *Index = nullptr;
+    /// The FS to be used by actions that will need access to files that are not
+    /// touched while building AST.
+    llvm::vfs::FileSystem *FS = nullptr;
     // FIXME: provide a way to get sources and ASTs for other files.
   };
 
   /// Output of a tweak.
   enum Intent {
     /// Apply changes that preserve the behavior of the code.
     Refactor,
     /// Provide information to the user.

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

 #include "index/Index.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/Registry.h"
+#include "llvm/Support/VirtualFileSystem.h"
 #include <functional>
 #include <memory>
 #include <utility>
 
 LLVM_INSTANTIATE_REGISTRY(llvm::Registry<clang::clangd::Tweak>)
 
 namespace clang {
 namespace clangd {
assert(E.instantiate()->id() == E.getName() &&
            "id should be equal to class name");
     assert(Seen.try_emplace(E.getName()).second && "duplicate check id");
   }
 #endif
 }
 } // namespace
 
 Tweak::Selection::Selection(ParsedAST &AST, unsigned RangeBegin,
-                            unsigned RangeEnd, const SymbolIndex *Index)
+                            unsigned RangeEnd, const SymbolIndex *Index,
+                            llvm::vfs::FileSystem *FS)
     : AST(AST), SelectionBegin(RangeBegin), SelectionEnd(RangeEnd),
       ASTSelection(AST.getASTContext(), AST.getTokens(), RangeBegin, RangeEnd),
-      Index(Index) {
+      Index(Index), 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) {

clang-tools-extra/clangd/unittests/TestTU.h

 //===---------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_UNITTESTS_CLANGD_TESTTU_H
 #define LLVM_CLANG_TOOLS_EXTRA_UNITTESTS_CLANGD_TESTTU_H
 
 #include "ParsedAST.h"
 #include "Path.h"
 #include "index/Index.h"
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/StringMap.h"
+#include "llvm/Support/VirtualFileSystem.h"
 #include "gtest/gtest.h"
 #include <string>
 #include <utility>
 #include <vector>
 
 namespace clang {
 namespace clangd {
 
struct TestTU {
   const SymbolIndex *ExternalIndex = nullptr;
 
   // Simulate a header guard of the header (using an #import directive).
   bool ImplicitHeaderGuard = true;
 
   ParsedAST build() const;
   SymbolSlab headerSymbols() const;
   std::unique_ptr<SymbolIndex> index() const;
+  llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> getVFS() const;
 };
 
 // Look up an index symbol by qualified name, which must be unique.
 const Symbol &findSymbol(const SymbolSlab &, llvm::StringRef QName);
 // Look up an AST symbol by qualified name, which must be unique and top-level.
 const NamedDecl &findDecl(ParsedAST &AST, llvm::StringRef QName);
 // Look up an AST symbol that satisfies \p Filter.
 const NamedDecl &findDecl(ParsedAST &AST,
                           std::function<bool(const NamedDecl &)> Filter);
 // Look up an AST symbol by unqualified name, which must be unique.
 const NamedDecl &findUnqualifiedDecl(ParsedAST &AST, llvm::StringRef Name);
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_UNITTESTS_CLANGD_TESTTU_H

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

 #include "index/MemIndex.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/Utils.h"
 
 namespace clang {
 namespace clangd {
 
-ParsedAST TestTU::build() const {
+llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> TestTU::getVFS() const {
   std::string FullFilename = testPath(Filename),
               FullHeaderName = testPath(HeaderFilename),
               ImportThunk = testPath("import_thunk.h");
   // We want to implicitly include HeaderFilename without messing up offsets.
   // -include achieves this, but sometimes we want #import (to simulate a header
   // guard without messing up offsets). In this case, use an intermediate file.
   std::string ThunkContents = "#import \"" + FullHeaderName + "\"\n";
 
   llvm::StringMap<std::string> Files(AdditionalFiles);
   Files[FullFilename] = Code;
   Files[FullHeaderName] = HeaderCode;
   Files[ImportThunk] = ThunkContents;
+  return buildTestFS(Files);
+}
+
+ParsedAST TestTU::build() const {
+  std::string FullFilename = testPath(Filename),
+              FullHeaderName = testPath(HeaderFilename),
+              ImportThunk = testPath("import_thunk.h");
 
   std::vector<const char *> Cmd = {"clang"};
   // FIXME: this shouldn't need to be conditional, but it breaks a
   // GoToDefinition test for some reason (getMacroArgExpandedLocation fails).
   if (!HeaderCode.empty()) {
     Cmd.push_back("-include");
     Cmd.push_back(ImplicitHeaderGuard ? ImportThunk.c_str()
                                       : FullHeaderName.c_str());
     // ms-compatibility changes the meaning of #import.
     // The default is OS-dependent (on on windows), ensure it's off.
     if (ImplicitHeaderGuard)
       Cmd.push_back("-fno-ms-compatibility");
   }
   Cmd.insert(Cmd.end(), ExtraArgs.begin(), ExtraArgs.end());
   // Put the file name at the end -- this allows the extra arg (-xc++) to
   // override the language setting.
   Cmd.push_back(FullFilename.c_str());
   ParseInputs Inputs;
   Inputs.CompileCommand.Filename = FullFilename;
   Inputs.CompileCommand.CommandLine = {Cmd.begin(), Cmd.end()};
   Inputs.CompileCommand.Directory = testRoot();
   Inputs.Contents = Code;
-  Inputs.FS = buildTestFS(Files);
+  Inputs.FS = getVFS();
   Inputs.Opts = ParseOptions();
   Inputs.Opts.ClangTidyOpts.Checks = ClangTidyChecks;
   Inputs.Opts.ClangTidyOpts.WarningsAsErrors = ClangTidyWarningsAsErrors;
   Inputs.Index = ExternalIndex;
   if (Inputs.Index)
     Inputs.Opts.SuggestMissingIncludes = true;
   StoreDiags Diags;
   auto CI = buildCompilerInvocation(Inputs, Diags);

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

MATCHER_P5(TweakIsAvailable, TweakID, Ctx, Header, ExtraFiles, Index,
   std::string WrappedCode = wrap(Ctx, arg);
   Annotations Input(WrappedCode);
   auto Selection = rangeOrPoint(Input);
   TestTU TU;
   TU.HeaderCode = Header;
   TU.Code = Input.code();
   TU.AdditionalFiles = std::move(ExtraFiles);
   ParsedAST AST = TU.build();
-  Tweak::Selection S(AST, Selection.first, Selection.second, Index);
+  auto FS = TU.getVFS();
+  Tweak::Selection S(AST, Selection.first, Selection.second, Index, FS.get());
   auto PrepareResult = prepareTweak(TweakID, S);
   if (PrepareResult)
     return true;
   llvm::consumeError(PrepareResult.takeError());
   return false;
 }
 
 } // namespace
 
 std::string TweakTest::apply(llvm::StringRef MarkedCode,
                              llvm::StringMap<std::string> *EditedFiles) const {
   std::string WrappedCode = wrap(Context, MarkedCode);
   Annotations Input(WrappedCode);
   auto Selection = rangeOrPoint(Input);
 
   TestTU TU;
   TU.HeaderCode = Header;
   TU.AdditionalFiles = std::move(ExtraFiles);
   TU.Code = Input.code();
   TU.ExtraArgs = ExtraArgs;
   ParsedAST AST = TU.build();
-  Tweak::Selection S(AST, Selection.first, Selection.second, Index);
+
+  auto FS = TU.getVFS();
+  Tweak::Selection S(AST, Selection.first, Selection.second, Index, FS.get());
 
   auto T = prepareTweak(TweakID, S);
   if (!T) {
     llvm::consumeError(T.takeError());
     return "unavailable";
   }
   llvm::Expected<Tweak::Effect> Result = (*T)->apply(S);
   if (!Result)