[clangd] Support multifile edits as output of Tweaks

[kadircet]

Summary

First patch for propogating multifile changes from tweak outputs to LSP
WorkspaceEdits.

Uses FS to convert tooling::Replacements to TextEdits.
Errors out if there are any inconsistencies between the draft version and the
version generated the edits.

[ilya-biryukov]

clang-tools-extra/clangd/ClangdLSPServer.cpp

674

This callback is called asynchronously and the version inside DraftMgr may be newer than the one we used to calculate the original offsets inside replacements at this point.

We should not rely on DraftMgr, doing conversions on the source buffers inside SourceManager is probably the most reliable option we have.

[kadircet]

clang-tools-extra/clangd/ClangdLSPServer.cpp

674

This callback is called asynchronously and the version inside DraftMgr may be newer than the one we used to calculate the original offsets inside replacements at this point.

That's exactly what we are checking though?

We want to bail out if user has a "different" version of the source code in their editor by the time we generated edits. Since editors will most likely try to apply the edits onto the version in editor.

[sammccall]

clang-tools-extra/clangd/ClangdLSPServer.cpp

674

I think you've misunderstood what this code is doing (it needs comments!)

The SourceManager contains the code that we've calculated the edits against, by definition. So checking against it does nothing.

When we send the edits to the client, it's going to apply them to the file if it's not open, but to the dirty buffer if it is open.
If the dirty buffer has different contents than what we saw (and calculated edits against) then the edits will be wrong - that's what this code guards against.

So looking at the latest available content in the draftmgr (at the end of the request, not the start!) is the right thing here, I think.

676

We're comparing here content in the source manager (which was read as bytes from disk, and converted to utf8) to that from LSP (which was sent as unicode text by the editor, and converted to utf8).

Are editors going to preserve the actual line endings from disk when sending LSP, or are they going to use the "array-of-lines" from their editor's native model? (e.g. join(lines, '\n'))

If the latter, we'd need to normalize line endings, possibly strip trailing line endings, etc.
I ask because clangd-vim definitely has an "array-of-lines" model :-)

[sammccall]

clang-tools-extra/clangd/SourceCode.h

43

Hmm, it is also attractive to make this something like

struct Edit {
  string CodeBefore;
  tooling::Replacements Reps;

  string codeAfter();
  std::vector<TextEdit> asTextEdits();
  bool canApplyTo(StringRef CodeBefore); // does line-ending normalization
}

if we were worried about abuse, we could make CodeBefore private

45

, if possible.
(Sometimes we're just sending these over the wire)

48

document that these are equivalent.
Why is Edits optional?

48

nit: reps -> replacements
ContentDigests -> ContentDigest

51

make this a constructor?

[ilya-biryukov]

clang-tools-extra/clangd/ClangdLSPServer.cpp

674

You are right, I'm sorry. I did not read into the code, just assumed it was doing something it doesn't.
Again, very sorry for the disturbance.

[kadircet]

• Address comments

[sammccall]

clang-tools-extra/clangd/SourceCode.cpp

601

you seem to be checking this both here and in clangdlspserver. Why?

clang-tools-extra/clangd/SourceCode.h

42

nit: drop "represents"

42

nit: this could also describe Replacements, vector<TextEdit>, etc. Motivate this class a little more?

195

not sure what "generates TextEdits from those" refers to.

Could this function be called "reformatEdits" or "formatAroundEdits"?

200

This signature is confusing: we pass code in *three* different ways (FS, Edits, and MainFileCode). Much of this is because we put the loop (and therefore all special cases) inside this function.
The logic around the way the VFS is mapped for the main file vs others doesn't really belong in this layer. Neither does "please save"...

It seems this wants to be something simpler/smaller like reformatEdit(const Edit&, const Style&) that could be called from ClangdServer. There's probably another helper like checkApplies(const Edit&, VFS*) that would go in ClangdLSPServer.

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

76

what's the difference between None and an empty map?

78–79

(if the null map case goes away, this can too)

83

This greatly increases the burden of callers of this function, who mostly want to edit the current file.

• need to provide a fully-qualified name (I don't think the implementations in this patch actually do that)
• need to construct an Edit

can we provide an Effect mainFileEdit(const SourceManager&, Replacements)? and maybe Effect fileEdit(const SourceManager&, FileID, Replacements) though maybe the latter doesn't cover enough cases to pull its weight.

[kadircet]

clang-tools-extra/clangd/SourceCode.cpp

601

the contents in here are coming from disk, whereas the one in clangdlspserver is coming from drafts. I suppose it would make more sense to merge the two in clangdlspserver, use draft manager as content source backed by disk in case user hasn't opened that file.

clang-tools-extra/clangd/SourceCode.h

195

not sure what "generates TextEdits from those" refers to.

Leftover from an old-version.

200

as discussed offline for reformatEdit to work, made InitialCode in Edit public. As both formatting and style deduction requires access to file contents.

for checkApplies, it requires access to both VFS and DraftMgrand pretty special use-case for clangdlspserver, so don't think the function would carry its weight. Therefore left the check inlined.

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

76

None :P

[kadircet]

• Address comments

[kadircet]

• Handle formatting error in ClangdServer rather than just printing it

[sammccall]

clang-tools-extra/clangd/SourceCode.cpp

601

OK, so you're checking here whether the contents-on-disk have changed between clang reading the file as part of parse, and the tweak completing?

I don't think this is worth doing. Sure, there's potentially a race condition, but it's relatively small (bounded on the order of a few seconds, vs unbounded for unsaved drafts) and we can't eliminate it entirely (file on disk can change between when we create edits and when they're applied).

[kadircet]

• Get rid off compatibility check for files that are not open in the editor.

[sammccall]

clang-tools-extra/clangd/ClangdLSPServer.cpp

678

please pull out a method for this part
e.g. Error validateEdits(...)

685

As these are absolute paths, we probably want to invert the error message for readability, and at least count the files.

e.g.
"Files must be saved first: path/to/Foo.cpp (and 3 others)"

clang-tools-extra/clangd/ClangdServer.cpp

414

isn't this just
elog("Failed to format {0}: {1}", std::move(Err))

clang-tools-extra/clangd/SourceCode.cpp

856

documentation

870

why?

clang-tools-extra/clangd/SourceCode.h

42

move this class near the related code? (replacementToEdit etc)

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

92

how do you know this is the correct/absolute path for the file? Can we add tests?

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

75–77

is it possible to be more specific than "absolute"?

133

what will this be used for?

You can use at most one of these Effect factories (unless we're going to do some convoluted merge thing), so this seems useful if you've got exactly one edit to a non-main file.

It seems more useful to return a pair<string, Edit> which could be inserted into a map

133

why are these moved out of the Effect class?

136

nit: convenience

[kadircet]

clang-tools-extra/clangd/ClangdServer.cpp

414

IIUC, just printing the error doesn't actually mark it as checked, therefore causes an assertion failure (see getPtr vs getPayload in llvm::Error). And I couldn't see any special handling in elog, but not sure if formatv object has this.

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

92

changed it to use FileInfo::getName rather then FileEntry::getName, as the former says:

/// Returns the name of the file that was used when the file was loaded from
/// the underlying file system.

These also get tested through any lit tests we have for tweaks, adding more unittests.

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

133

need to provide a fully-qualified name (I don't think the implementations in this patch actually do that)

I've thought you were also requesting these to be helper functions rather than static methods so that callers don't need to qualify the call.
Can move it back into class if I misunderstood your point, but this actually seems a lot cleaner.

133

what will this be used for?

you are right it doesn't really compose well, changed it to return the pair instead.

[kadircet]

• Address comments

[kadircet]

• Define more strict semantics around filename

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

133

moved back into class

[kadircet]

• Rebase

[sammccall]

clang-tools-extra/clangd/ClangdServer.cpp

414

The special handling is on Logger.h:38. Do you see crashes?

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

93

nit: "for edited file: "...

[kadircet]

• Address comments

clang-tools-extra/clangd/ClangdServer.cpp

414

ah OK, didn't notice the fmt_consume previously. no, I wasn't seeing any crashes.

clang-tools-extra/clangd/ClangdLSPServer.cpp

 //===--- ClangdLSPServer.cpp - LSP server ------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "ClangdLSPServer.h"
 #include "Diagnostics.h"
+#include "DraftStore.h"
 #include "FormattedString.h"
 #include "GlobalCompilationDatabase.h"
 #include "Protocol.h"
 #include "SemanticHighlighting.h"
 #include "SourceCode.h"
 #include "Trace.h"
 #include "URI.h"
 #include "refactor/Tweak.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
+#include "llvm/Support/SHA1.h"
 #include "llvm/Support/ScopedPrinter.h"
+#include <cstddef>
+#include <string>
 
 namespace clang {
 namespace clangd {
 namespace {
 /// Transforms a tweak into a code action that would apply it if executed.
 /// EXPECTS: T.prepare() was called and returned true.
 CodeAction toCodeAction(const ClangdServer::TweakRef &T, const URIForFile &File,
                         Range Selection) {
std::vector<std::vector<std::string>> buildHighlightScopeLookupTable() {
   std::vector<std::vector<std::string>> LookupTable;
   // HighlightingKind is using as the index.
   for (int KindValue = 0; KindValue < (int)HighlightingKind::NumKinds;
        ++KindValue)
     LookupTable.push_back({toTextMateScope((HighlightingKind)(KindValue))});
   return LookupTable;
 }
 
+// Makes sure edits in \p E are applicable to latest file contents reported by
+// editor. If not generates an error message containing information about files
+// that needs to be saved.
+llvm::Error validateEdits(const DraftStore &DraftMgr, const Tweak::Effect &E) {
+  size_t InvalidFileCount = 0;
+  llvm::StringRef LastInvalidFile;
+  for (const auto &It : E.ApplyEdits) {
+    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)) {
+        ++InvalidFileCount;
+        LastInvalidFile = It.first();
+      }
+    }
+  }
+  if (!InvalidFileCount)
+    return llvm::Error::success();
+  if (InvalidFileCount == 1)
+    return llvm::createStringError(llvm::inconvertibleErrorCode(),
+                                   "File must be saved first: " +
+                                       LastInvalidFile);
+  return llvm::createStringError(
+      llvm::inconvertibleErrorCode(),
+      "Files must be saved first: " + LastInvalidFile + " (and " +
+          llvm::to_string(InvalidFileCount - 1) + " others)");
+}
+
 } // namespace
 
 // MessageHandler dispatches incoming LSP messages.
 // It handles cross-cutting concerns:
 //  - serializes/deserializes protocol objects to JSON
 //  - logging of inbound messages
 //  - cancellation handling
 //  - basic call tracing
if (!Code)
           "trying to apply a code action for a non-added file"));
 
     auto Action = [this, ApplyEdit, Reply = std::move(Reply),
                    File = Params.tweakArgs->file, Code = std::move(*Code)](
                       llvm::Expected<Tweak::Effect> R) mutable {
       if (!R)
         return Reply(R.takeError());
 
-      assert(R->ShowMessage || (R->ApplyEdit && "tweak has no effect"));
+      assert(R->ShowMessage ||
+             (!R->ApplyEdits.empty() && "tweak has no effect"));
 
       if (R->ShowMessage) {
         ShowMessageParams Msg;
         Msg.message = *R->ShowMessage;
         Msg.type = MessageType::Info;
         notify("window/showMessage", Msg);
       }
-      if (R->ApplyEdit) {
-        WorkspaceEdit WE;
-        WE.changes.emplace();
-        (*WE.changes)[File.uri()] = replacementsToEdits(Code, *R->ApplyEdit);
-        // ApplyEdit will take care of calling Reply().
-        return ApplyEdit(std::move(WE), "Tweak applied.", std::move(Reply));
-      }
-      // When no edit is specified, make sure we Reply().
-      return Reply("Tweak applied.");
+      // When no edit is specified, make sure we Reply().
+      if (R->ApplyEdits.empty())
+        return Reply("Tweak applied.");
+
+      if (auto Err = validateEdits(DraftMgr, *R))
+        return Reply(std::move(Err));
+
+      for (const auto &It : R->ApplyEdits) {
+        (*WE.changes)[URI::create(It.first()).toString()] =
+            It.second.asTextEdits();
+      }
     };
     Server->applyTweak(Params.tweakArgs->file.file(),
                        Params.tweakArgs->selection, Params.tweakArgs->tweakID,
                        std::move(Action));
   } else {
     // We should not get here because ExecuteCommandParams would not have
     // parsed in the first place and this handler should not be called. But if
     // more commands are added, this will be here has a safe guard.

clang-tools-extra/clangd/ClangdServer.cpp

 //===--- ClangdServer.cpp - Main clangd server code --------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===-------------------------------------------------------------------===//
 
 #include "ClangdServer.h"
 #include "ClangdUnit.h"
 #include "CodeComplete.h"
 #include "FindSymbols.h"
 #include "Format.h"
 #include "FormattedString.h"
 #include "Headers.h"
+#include "Logger.h"
 #include "Protocol.h"
 #include "SemanticHighlighting.h"
 #include "SourceCode.h"
 #include "TUScheduler.h"
 #include "Trace.h"
 #include "XRefs.h"
 #include "index/CanonicalIncludes.h"
 #include "index/FileIndex.h"
void ClangdServer::enumerateTweaks(PathRef File, Range Sel,
   };
 
   WorkScheduler.runWithAST("EnumerateTweaks", File, std::move(Action));
 }
 
 void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef TweakID,
                               Callback<Tweak::Effect> CB) {
   auto Action = [File = File.str(), Sel, TweakID = TweakID.str(),
-                 CB = std::move(CB)](Expected<InputsAndAST> InpAST) mutable {
+                 CB = std::move(CB), FS = FSProvider.getFileSystem()](
+                    Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     auto Selection = tweakSelection(Sel, *InpAST);
     if (!Selection)
       return CB(Selection.takeError());
     auto A = prepareTweak(TweakID, *Selection);
     if (!A)
       return CB(A.takeError());
     auto Effect = (*A)->apply(*Selection);
     if (!Effect)
       return CB(Effect.takeError());
-    if (Effect->ApplyEdit) {
-      // FIXME: this function has I/O operations (find .clang-format file),
-      // figure out a way to cache the format style.
-      auto Style = getFormatStyleForFile(File, InpAST->Inputs.Contents,
-                                         InpAST->Inputs.FS.get());
-      if (auto Formatted = cleanupAndFormat(InpAST->Inputs.Contents,
-                                            *Effect->ApplyEdit, Style))
-        Effect->ApplyEdit = std::move(*Formatted);
-      else
-        elog("Failed to format replacements: {0}", Formatted.takeError());
+    for (auto &It : Effect->ApplyEdits) {
+      Edit &E = It.second;
+      format::FormatStyle Style =
+          getFormatStyleForFile(File, E.InitialCode, FS.get());
+      if (llvm::Error Err = reformatEdit(E, Style)) {
+        llvm::handleAllErrors(std::move(Err), [&](llvm::ErrorInfoBase &EIB) {
+          elog("Failed to format {0}: {1}", It.first(), EIB.message());
+        });
+      }
     }
     return CB(std::move(*Effect));
   };
   WorkScheduler.runWithAST("ApplyTweak", File, std::move(Action));
 }
 
 void ClangdServer::dumpAST(PathRef File,
                            llvm::unique_function<void(std::string)> Callback) {

clang-tools-extra/clangd/SourceCode.h

 //===--- SourceCode.h - Manipulating source code as strings -----*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Various code that examines C++ source code without using heavy AST machinery
 // (and often not even the lexer). To be used sparingly!
 //
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_SOURCECODE_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_SOURCECODE_H
+
 #include "Context.h"
 #include "Protocol.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Format/Format.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/SHA1.h"
+#include <string>
 
 namespace clang {
 class SourceManager;
 
 namespace clangd {
 
 // We tend to generate digests for source codes in a lot of different places.
 // This represents the type for those digests to prevent us hard coding details
 // of hashing function at every place that needs to store this information.
 using FileDigest = std::array<uint8_t, 8>;
 FileDigest digest(StringRef Content);
 Optional<FileDigest> digestFile(const SourceManager &SM, FileID FID);
 
 // This context variable controls the behavior of functions in this file
 // that convert between LSP offsets and native clang byte offsets.
 // If not set, defaults to UTF-16 for backwards-compatibility.
 extern Key<OffsetEncoding> kCurrentOffsetEncoding;
 
 // Counts the number of UTF-16 code units needed to represent a string (LSP
 // specifies string lengths in UTF-16 code units).
 // Use of UTF-16 may be overridden by kCurrentOffsetEncoding.
 size_t lspLength(StringRef Code);
 
 /// Turn a [line, column] pair into an offset in Code.
 ///
 /// If P.character exceeds the line length, returns the offset at end-of-line.
 /// (If !AllowColumnsBeyondLineLength, then returns an error instead).
 /// If the line number is out of range, returns an error.
 ///
 /// The returned value is in the range [0, Code.size()].
 llvm::Expected<size_t>

 /// This will usually use FS to look for .clang-format directories.
 /// FIXME: should we be caching the .clang-format file search?
 /// This uses format::DefaultFormatStyle and format::DefaultFallbackStyle,
 /// though the latter may have been overridden in main()!
 format::FormatStyle getFormatStyleForFile(llvm::StringRef File,
                                           llvm::StringRef Content,
                                           llvm::vfs::FileSystem *FS);
 
-// Cleanup and format the given replacements.
+/// Cleanup and format the given replacements.
 llvm::Expected<tooling::Replacements>
 cleanupAndFormat(StringRef Code, const tooling::Replacements &Replaces,
                  const format::FormatStyle &Style);
 
+/// A set of edits generated for a single file. Can verify whether it is safe to
+/// apply these edits to a code block.
+struct Edit {
+  tooling::Replacements Replacements;
+  std::string InitialCode;
+
+  Edit(llvm::StringRef Code, tooling::Replacements Reps)
+      : Replacements(std::move(Reps)), InitialCode(Code) {}
+
+  /// Returns the file contents after changes are applied.
+  llvm::Expected<std::string> apply() const;
+
+  /// Represents Replacements as TextEdits that are available for use in LSP.
+  std::vector<TextEdit> asTextEdits() const;
+
+  /// Checks whether the Replacements are applicable to given Code.
+  bool canApplyTo(llvm::StringRef Code) const;
+};
+
+/// Formats the edits and code around it according to Style. Changes
+/// Replacements to formatted ones if succeeds.
+llvm::Error reformatEdit(Edit &E, const format::FormatStyle &Style);
+
 /// Collects identifiers with counts in the source code.
 llvm::StringMap<unsigned> collectIdentifiers(llvm::StringRef Content,
                                              const format::FormatStyle &Style);
 
 /// Collects words from the source code.
 /// Unlike collectIdentifiers:
 /// - also finds text in comments:
 /// - splits text into words

clang-tools-extra/clangd/SourceCode.cpp

 //===--- SourceCode.h - Manipulating source code as strings -----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #include "SourceCode.h"
 
 #include "Context.h"
 #include "FuzzyMatch.h"
 #include "Logger.h"
 #include "Protocol.h"
+#include "refactor/Tweak.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Format/Format.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/Preprocessor.h"
+#include "clang/Tooling/Core/Replacement.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/LineIterator.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
+#include "llvm/Support/SHA1.h"
+#include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Support/xxhash.h"
 #include <algorithm>
 
 namespace clang {
 namespace clangd {
 
 // Here be dragons. LSP positions use columns measured in *UTF-16 code units*!
 // Clangd uses UTF-8 and byte-offsets internally, so conversion is nontrivial.

 llvm::StringMap<unsigned> collectIdentifiers(llvm::StringRef Content,
                                              const format::FormatStyle &Style) {
   llvm::StringMap<unsigned> Identifiers;
   lex(Content, Style, [&](const clang::Token &Tok) {
     switch (Tok.getKind()) {
     case tok::identifier:
       ++Identifiers[Tok.getIdentifierInfo()->getName()];
       break;
     case tok::raw_identifier:
       ++Identifiers[Tok.getRawIdentifier()];
       break;
     default:
       break;
     }
   });
   return Identifiers;
 }
SourceLocation BeforeSearchedLocation =
                                          .getLocWithOffset(DecLoc.second - 1));
   MacroDefinition MacroDef =
       PP.getMacroDefinitionAtLoc(IdentifierInfo, BeforeSearchedLocation);
   if (auto *MI = MacroDef.getMacroInfo())
     return DefinedMacro{IdentifierInfo->getName(), MI};
   return None;
 }
 
+llvm::Expected<std::string> Edit::apply() const {
+  return tooling::applyAllReplacements(InitialCode, Replacements);
+}
+
+std::vector<TextEdit> Edit::asTextEdits() const {
+  return replacementsToEdits(InitialCode, Replacements);
+}
+
+bool Edit::canApplyTo(llvm::StringRef Code) const {
+  // Create line iterators, since line numbers are important while applying our
+  // edit we cannot skip blank lines.
+  auto LHS = llvm::MemoryBuffer::getMemBuffer(Code);
+  llvm::line_iterator LHSIt(*LHS, /*SkipBlanks=*/false);
+
+  auto RHS = llvm::MemoryBuffer::getMemBuffer(InitialCode);
+  llvm::line_iterator RHSIt(*RHS, /*SkipBlanks=*/false);
+
+  // Compare the InitialCode we prepared the edit for with the Code we received
+  // line by line to make sure there are no differences.
+  // FIXME: This check is too conservative now, it should be enough to only
+  // check lines around the replacements contained inside the Edit.
+  while (!LHSIt.is_at_eof() && !RHSIt.is_at_eof()) {
+    if (*LHSIt != *RHSIt)
+      return false;
+    ++LHSIt;
+    ++RHSIt;
+  }
+
+  // After we reach EOF for any of the files we make sure the other one doesn't
+  // contain any additional content except empty lines, they should not
+  // interfere with the edit we produced.
+  while (!LHSIt.is_at_eof()) {
+    if (!LHSIt->empty())
+      return false;
+    ++LHSIt;
+  }
+  while (!RHSIt.is_at_eof()) {
+    if (!RHSIt->empty())
+      return false;
+    ++RHSIt;
+  }
+  return true;
+}
+
+llvm::Error reformatEdit(Edit &E, const format::FormatStyle &Style) {
+  if (auto NewEdits = cleanupAndFormat(E.InitialCode, E.Replacements, Style))
+    E.Replacements = std::move(*NewEdits);
+  else
+    return NewEdits.takeError();
+  return llvm::Error::success();
+}
+
 } // namespace clangd
 } // namespace clang

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

 //   - Stage 2 is performed after stage 1 and can be more expensive to compute.
 //     It is performed when the user actually chooses the action.
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_REFACTOR_ACTIONS_TWEAK_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_REFACTOR_ACTIONS_TWEAK_H
 
 #include "ClangdUnit.h"
+#include "Path.h"
 #include "Protocol.h"
 #include "Selection.h"
+#include "SourceCode.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/Optional.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.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.
enum Intent {
     /// Apply changes that preserve the behavior of the code.
     Refactor,
     /// Provide information to the user.
     Info,
   };
   struct Effect {
     /// A message to be displayed to the user.
     llvm::Optional<std::string> ShowMessage;
-    /// An edit to apply to the input file.
-    llvm::Optional<tooling::Replacements> ApplyEdit;
+    /// A mapping from file path(the one used for accessing the underlying VFS)
+    /// to edits.
+    llvm::StringMap<Edit> ApplyEdits;
 
-    static Effect applyEdit(tooling::Replacements R) {
-      Effect E;
-      E.ApplyEdit = std::move(R);
-      return E;
-    }
     static Effect showMessage(StringRef S) {
       Effect E;
       E.ShowMessage = S;
       return E;
     }
   };
 
   virtual ~Tweak() = default;
   /// A unique id of the action, it is always equal to the name of the class
   /// defining the Tweak. Definition is provided automatically by
   /// REGISTER_TWEAK.
   virtual const char *id() const = 0;
   /// Run the first stage of the action. Returns true indicating that the
prepareTweaks(const Tweak::Selection &S,
               llvm::function_ref<bool(const Tweak &)> Filter);
 
 // Calls prepare() on the tweak with a given ID.
 // If prepare() returns false, returns an error.
 // If prepare() returns true, returns the corresponding tweak.
 llvm::Expected<std::unique_ptr<Tweak>> prepareTweak(StringRef TweakID,
                                                     const Tweak::Selection &S);
 
+/// Returns path to the FID as known by SM and an Edit containing Replacements
+/// as the ApplyEdits.
+std::pair<Path, Edit> fileEdit(const SourceManager &SM, FileID FID,
+                               tooling::Replacements Replacements);
+
+/// Creates an effect with an Edit for the main file.
+Tweak::Effect mainFileEdit(const SourceManager &SM,
+                           tooling::Replacements Replacements);
 } // namespace clangd
 } // namespace clang
 
 #endif

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

 //===--- Tweak.cpp -----------------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #include "Tweak.h"
 #include "Logger.h"
+#include "Path.h"
+#include "SourceCode.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/Registry.h"
 #include <functional>
 #include <memory>
 
 LLVM_INSTANTIATE_REGISTRY(llvm::Registry<clang::clangd::Tweak>)
return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    "id of the tweak is invalid");
   std::unique_ptr<Tweak> T = It->instantiate();
   if (!T->prepare(S))
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    "failed to prepare() a check");
   return std::move(T);
 }
 
+std::pair<Path, Edit> fileEdit(const SourceManager &SM, FileID FID,
+                               tooling::Replacements Replacements) {
+  bool IsInvalid = false;
+  llvm::StringRef FilePath =
+      SM.getSLocEntry(FID, &IsInvalid).getFile().getName();
+  assert(!IsInvalid && "Generating edits for unknown file!");
+  Edit Ed(SM.getBufferData(FID), std::move(Replacements));
+  return {FilePath.str(), std::move(Ed)};
+}
+
+Tweak::Effect mainFileEdit(const SourceManager &SM,
+                           tooling::Replacements Replacements) {
+  Tweak::Effect E;
+  auto PathAndEdit = fileEdit(SM, SM.getMainFileID(), std::move(Replacements));
+  E.ApplyEdits.try_emplace(PathAndEdit.first, PathAndEdit.second);
+  return E;
+}
+
 } // namespace clangd
 } // namespace clang

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

 //===--- AnnotateHighlightings.cpp -------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #include "SemanticHighlighting.h"
 #include "refactor/Tweak.h"
+#include "llvm/ADT/StringRef.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
 /// Annotate all highlighting tokens in the current file. This is a hidden tweak
 /// which is used to debug semantic highlightings.
 /// Before:
if (!CommonDecl) {
     Inputs.AST.getASTContext().setTraversalScope(
         {const_cast<Decl *>(CommonDecl)});
     HighlightingTokens = getSemanticHighlightings(Inputs.AST);
     // Restore the traversal scope.
     Inputs.AST.getASTContext().setTraversalScope(BackupScopes);
   }
   auto &SM = Inputs.AST.getSourceManager();
   tooling::Replacements Result;
+  llvm::StringRef FilePath = SM.getFilename(Inputs.Cursor);
   for (const auto &Token : HighlightingTokens) {
     assert(Token.R.start.line == Token.R.end.line &&
            "Token must be at the same line");
     auto InsertOffset = positionToOffset(Inputs.Code, Token.R.start);
     if (!InsertOffset)
       return InsertOffset.takeError();
 
     auto InsertReplacement = tooling::Replacement(
-        SM.getFileEntryForID(SM.getMainFileID())->getName(), *InsertOffset, 0,
+        FilePath, *InsertOffset, 0,
         ("/* " + toTextMateScope(Token.Kind) + " */").str());
     if (auto Err = Result.add(InsertReplacement))
       return std::move(Err);
   }
-  return Effect::applyEdit(Result);
+  return mainFileEdit(SM, std::move(Result));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

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

Expected<Tweak::Effect> ExpandAutoType::apply(const Selection& Inputs) {
 
   std::string PrettyTypeName = printType(*DeducedType,
       Inputs.ASTSelection.commonAncestor()->getDeclContext());
 
   tooling::Replacement
       Expansion(SrcMgr, CharSourceRange(CachedLocation->getSourceRange(), true),
                 PrettyTypeName);
 
-  return Tweak::Effect::applyEdit(tooling::Replacements(Expansion));
+  return mainFileEdit(SrcMgr, tooling::Replacements(Expansion));
 }
 
 llvm::Error ExpandAutoType::createErrorMessage(const std::string& Message,
                                                const Selection& Inputs) {
   auto& SrcMgr = Inputs.AST.getASTContext().getSourceManager();
   std::string ErrorMessage =
       Message + ": " +
           SrcMgr.getFilename(Inputs.Cursor).str() + " Line " +

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

bool ExpandMacro::prepare(const Selection &Inputs) {
   if (!Expansion)
     return false;
   this->MacroName = T->text(Inputs.AST.getSourceManager());
   this->Expansion = *Expansion;
   return true;
 }
 
 Expected<Tweak::Effect> ExpandMacro::apply(const Selection &Inputs) {
-  auto &SM = Inputs.AST.getASTContext().getSourceManager();
+  auto &SM = Inputs.AST.getSourceManager();
 
   std::string Replacement;
   for (const syntax::Token &T : Expansion.Expanded) {
     Replacement += T.text(SM);
     Replacement += " ";
   }
   if (!Replacement.empty()) {
     assert(Replacement.back() == ' ');
     Replacement.pop_back();
   }
 
   CharSourceRange MacroRange =
       CharSourceRange::getCharRange(Expansion.Spelled.front().location(),
                                     Expansion.Spelled.back().endLocation());
 
-  Tweak::Effect E;
-  E.ApplyEdit.emplace();
-  llvm::cantFail(
-      E.ApplyEdit->add(tooling::Replacement(SM, MacroRange, Replacement)));
-  return E;
+  tooling::Replacements Reps;
+  llvm::cantFail(Reps.add(tooling::Replacement(SM, MacroRange, Replacement)));
+  return mainFileEdit(SM, std::move(Reps));
 }
 
 std::string ExpandMacro::title() const {
   return llvm::formatv("Expand macro '{0}'", MacroName);
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

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

Expected<Tweak::Effect> ExtractVariable::apply(const Selection &Inputs) {
   std::string VarName = "dummy";
   SourceRange Range = Target->getExtractionChars();
   // insert new variable declaration
   if (auto Err = Result.add(Target->insertDeclaration(VarName, Range)))
     return std::move(Err);
   // replace expression with variable name
   if (auto Err = Result.add(Target->replaceWithVar(Range, VarName)))
     return std::move(Err);
-  return Effect::applyEdit(Result);
+  return mainFileEdit(Inputs.AST.getSourceManager(), std::move(Result));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

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

if (!N)
     return false;
   Str = dyn_cast_or_null<StringLiteral>(N->ASTNode.get<Stmt>());
   return Str &&
          isNormalString(*Str, Inputs.Cursor, Inputs.AST.getSourceManager()) &&
          needsRaw(Str->getBytes()) && canBeRaw(Str->getBytes());
 }
 
 Expected<Tweak::Effect> RawStringLiteral::apply(const Selection &Inputs) {
-  return Effect::applyEdit(tooling::Replacements(
+  auto &SM = Inputs.AST.getSourceManager();
+  auto Reps = tooling::Replacements(
       tooling::Replacement(Inputs.AST.getSourceManager(), Str,
                            ("R\"(" + Str->getBytes() + ")\"").str(),
-                           Inputs.AST.getASTContext().getLangOpts())));
+                           Inputs.AST.getASTContext().getLangOpts()));
+  return mainFileEdit(SM, std::move(Reps));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

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

Expected<Tweak::Effect> SwapIfBranches::apply(const Selection &Inputs) {
   if (auto Err = Result.add(tooling::Replacement(Ctx.getSourceManager(),
                                                  ThenRng->getBegin(),
                                                  ThenCode.size(), ElseCode)))
     return std::move(Err);
   if (auto Err = Result.add(tooling::Replacement(Ctx.getSourceManager(),
                                                  ElseRng->getBegin(),
                                                  ElseCode.size(), ThenCode)))
     return std::move(Err);
-  return Effect::applyEdit(Result);
+  return mainFileEdit(SrcMgr, std::move(Result));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

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

 //===-- TweakTesting.cpp ------------------------------------------------*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "TweakTesting.h"
 
 #include "Annotations.h"
-#include "refactor/Tweak.h"
 #include "SourceCode.h"
+#include "refactor/Tweak.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/Support/Error.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 using Context = TweakTest::CodeContext;
 
if (!T) {
     llvm::consumeError(T.takeError());
     return "unavailable";
   }
   llvm::Expected<Tweak::Effect> Result = (*T)->apply(S);
   if (!Result)
     return "fail: " + llvm::toString(Result.takeError());
   if (Result->ShowMessage)
     return "message:\n" + *Result->ShowMessage;
-  if (Result->ApplyEdit) {
-    if (auto NewText =
-            tooling::applyAllReplacements(Input.code(), *Result->ApplyEdit))
-      return unwrap(Context, *NewText);
-    else
-      return "bad edits: " + llvm::toString(NewText.takeError());
-  }
-  return "no effect";
-}
+  if (Result->ApplyEdits.empty())
+    return "no effect";
+  if (Result->ApplyEdits.size() > 1)
+    return "received multi-file edits";
+
+  auto ApplyEdit = Result->ApplyEdits.begin()->second;
+  if (auto NewText = ApplyEdit.apply())
 
 ::testing::Matcher<llvm::StringRef> TweakTest::isAvailable() const {
   return TweakIsAvailable(llvm::StringRef(TweakID), Context, Header); 
 }
 
 std::vector<std::string> TweakTest::expandCases(llvm::StringRef MarkedCode) {
   Annotations Test(MarkedCode);
   llvm::StringRef Code = Test.code();

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

 //===-- TweakTests.cpp ------------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Annotations.h"
+#include "ClangdUnit.h"
 #include "SourceCode.h"
+#include "TestFS.h"
 #include "TestTU.h"
 #include "TweakTesting.h"
 #include "refactor/Tweak.h"
 #include "clang/AST/Expr.h"
 #include "clang/Basic/LLVM.h"
 #include "clang/Rewrite/Core/Rewriter.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/StringRef.h"
if (!T)
     return T.takeError();
   return (*T)->apply(S);
 }
 
 llvm::Expected<std::string> applyEdit(StringRef ID, llvm::StringRef Input) {
   auto Effect = apply(ID, Input);
   if (!Effect)
     return Effect.takeError();
-  if (!Effect->ApplyEdit)
+  if (Effect->ApplyEdits.empty())
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    "No replacements");
+  auto Edits = Effect->ApplyEdits;
+  if (Edits.size() > 1)
+    return llvm::createStringError(llvm::inconvertibleErrorCode(),
+                                   "Multi-file edits");
   Annotations Code(Input);
-  return applyAllReplacements(Code.code(), *Effect->ApplyEdit);
+  return applyAllReplacements(Code.code(), Edits.begin()->second.Replacements);
 }
 
 void checkTransform(llvm::StringRef ID, llvm::StringRef Input,
                     std::string Output) {
   auto Result = applyEdit(ID, Input);
   ASSERT_TRUE(bool(Result)) << llvm::toString(Result.takeError()) << Input;
   EXPECT_EQ(Output, std::string(*Result)) << Input;
 }
 
+TEST(FileEdits, AbsolutePath) {
+  TestTU TU;
+  TU.Filename = "foo.cpp";
+  TU.Code = R"cpp(
+    #include "a.h"
+    #include "b.h"
+    void bar();
+  )cpp";
+  TU.AdditionalFiles = {
+      {"a.h", "void baz();"},
+      {"test/b.h", "void foo();"},
+  };
+  auto T = "-I" + testPath("test/");
+  TU.ExtraArgs = {T.c_str()};
+
+  const ParsedAST AST = TU.build();
+  ASSERT_TRUE(AST.getDiagnostics().empty());
+
+  auto &SM = AST.getSourceManager();
+  const auto Files = {testPath("foo.cpp"), testPath("test/b.h"),
+                      testPath("a.h")};
+
+  for (const Decl *D : AST.getASTContext().getTranslationUnitDecl()->decls()) {
+    if (D->getLocation().isInvalid() ||
+        SM.isWrittenInBuiltinFile(D->getLocation()))
+      continue;
+
+    auto FID = SM.getFileID(D->getLocation());
+    auto Res = fileEdit(SM, FID, tooling::Replacements());
+    EXPECT_THAT(Files, testing::Contains(Res.first));
+  }
+}
+
 TWEAK_TEST(SwapIfBranches);
 TEST_F(SwapIfBranchesTest, Test) {
   Context = Function;
   EXPECT_EQ(apply("^if (true) {return 100;} else {continue;}"),
             "if (true) {continue;} else {return 100;}");
   EXPECT_EQ(apply("^if () {return 100;} else {continue;}"),
             "if () {continue;} else {return 100;}") << "broken condition";
   EXPECT_AVAILABLE("^i^f^^(^t^r^u^e^) { return 100; } ^e^l^s^e^ { continue; }");