[clangd] Implement the 'Organize Imports' source action. Fix include-cleaner findings in batch.

[kadircet]

there are unittests for tweaks under clang-tools-extra/clangd/unittests/tweaks/ can you create one for OrganizeImports there and test this in there?

i've got another concern around triggering of this interaction.
in theory code-action context has a only field, dictating which code actions should be surfaced, but when it's unset (as it's usually not set by clients and hard to understand semantics of code action kinds) clangd just provides all of the available code actions for a given selection.
since organize imports (and other source actions) are going to be available for all selections, it means we'll now have some code action response all the time and clients that don't know about semantics for "source" code action kind, will likely just show it in their UIs. so we need to be somewhat careful about triggering.

two alternatives that I can think of are:
1- show "source" kind code actions, only when they're explicitly requested.
2- make use of trigger kind and return these only when code actions are "explicitly" triggered.

1 has the downside of never showing up on editors that don't know about "source" code action kind, and also complicates behaviour on clangd side.
2 has the downside of being introduced very recently (trigger kind is made part of the spec at LSP 3.17), hence it'll likely be absent in old clients.

so my inclination is towards alternative 2), which at least limits the blast radius and gives clients a reasonably easy way to control the behavior. WDYT?

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

25–26

nit: namespace clang::clangd {

46–50

nit: i'd inline these into the class definition as well

52

s/OrganizeImports::Effect/Tweak::Effect

63

prefer llvm::StringSet<>

65–68

this will insert all providers for a symbol (e.g. if we have both foo.h and bar.h, we'll insert both).

76

78

unfortunately we actually have to consume the error here, as getStyle returns an Expected<T> and destroying an unchecked Expected is undefined behavior :/

something like elog("Couldn't get style for {0}: {1}", MainFilePath, FileStyle.takeError()); should do the trick.

80–84

nit:

[kadircet]

so as we were discussing offline (just to summarize it here) there's also the possibility of propagating context kind in selection, and making this code action fire when code actions are requested for "source code action kind" or the selection intersects with includes. which feels like the most useful alternative without annoying old clients.

[VitaNuo]

Address review comments. Update the action to trigger when source action is requested explicitly or from the preamble range. Add unit test.

[VitaNuo]

Thanks for the comments!

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

63

I think it should actually be llvm::DenseSet<inlude_cleaner::Header>. We should first de-duplicate providers and only then spell the resulting headers. In this version, we were potentially spelling the same provider multiple times.

65–68

you're right, thanks

[kadircet]

clang-tools-extra/clangd/ClangdServer.cpp

719–721

we're capturing ActionKinds as a reference, but Action is going to execute async. you can instead change the input to be llvm::ArrayRef<std::string> ActionKinds and capture it by value with ActionKinds = ActionKinds.vec()

nit: even better if you just changed the signature here to look like applyTweak(std::string TweakID, CodeActionInputs Inputs, Callback<Tweak::Effect> CB), as we want to consume all of these by value anyways. Then you can just move TweakID and Inputs into the Action.

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

50

can we also bail out for ObjC

57

relying on MainFileIncludes for preamble range is not correct in general. we can have includes way down inside the main file that'll be part of this vector, but not preamble (also it's not sorted explicitly).

we should instead use ComputePreambleBounds (nit: we can also store it in ParsedAST, instead of relexing the file one more time with each CodeAction request)

57

Having a comment for reason would also be helpful, something like To accommodate clients without knowledge of source actions, we trigger without checking code action kinds when inside the preamble region.

64

we should also make sure that RequestedActionKinds is empty in this case. Because we don't want to offer a code action of kind Source, if user is explicitly asking for some other kind(s).

73

nit: insertions of Replacements with offset UINT_MAX can't fail. so you just wrap this inside llvm::cantFail instead.

88

same here for never failing

93

i think we should perform this check on Final instead (as replacements can still go empty after cleanups, and empty set of replacements will result in an empty set anyways)

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

20

can you also add some availability tests ?

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

87

ATM this is passing SOURCE_KIND for all tweaks, which is conceptually wrong. can you instead take this in as a parameter?

for TweakTest::apply, we can default it to be an empty vector, and just update OrganizeImportsTests to pass {CodeAction::SOURCE_KIND};
for TweakTest::isAvailable, we can update EXPECT_AVAILABLE_ to pass in {} and call isAvailable directly in OrganizeImportsTests with SOURCE_KIND.

[VitaNuo]

Thank you for the comments!

clang-tools-extra/clangd/ClangdServer.cpp

719–721

Ok, I can change the signature, but it seems like I still need to move individual members of CodeActionInputs into the callback separately. File cannot be moved, since it's also needed after the callback as an argument to runWithAST.

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

57

nit: we can also store it in ParsedAST

Seems like the data is already there, I just need to expose it.

[kadircet]

FYI, i don't think you uploaded the new version of the patch

clang-tools-extra/clangd/ClangdServer.cpp

719–721

yeah, it's fine to make a copy

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

57

nit: we can also store it in ParsedAST

Seems like the data is already there, I just need to expose it.

not really, we calculate bounds when building a ParsedAST but we don't store it anywhere.

[VitaNuo]

I'm sorry I think the upload failed because I tried to upload without a message or something, and I didn't notice that it failed :(

[VitaNuo]

Address review comments.

[VitaNuo]

I think I also managed to update the wrong patch. Not sure how an alternative patch got created in the first place :( Sorry for the mess, should be fine now.

[VitaNuo]

Respond to comments.

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

57

I was referring to https://github.com/llvm/llvm-project/blob/main/clang/lib/Frontend/PrecompiledPreamble.cpp#L580. We store the PreambleData and that seems to store the PrecompiledPreamble which seems to know its bounds

[kadircet]

clang-tools-extra/clangd/ClangdLSPServer.cpp

810–815

this is kind of hard to read

clang-tools-extra/clangd/ParsedAST.h

124

Unfortunately this won't work in presence of stale preambles. You can use the "patched bounds" we calculate when building a parsedast in https://github.com/llvm/llvm-project/blob/main/clang-tools-extra/clangd/ParsedAST.cpp#L646; store it in ParsedAST in a new field, and return it here.

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

74

prefer llvm::ArrayRef<std:string> here, Selections are not passed in between threads. So no need to make an extra copy when creating one.

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

52

nit: llvm::find(Inputs.RequestedActionKinds, CodeAction::SOURCE_KIND)

56

nit:

if (!Inputs.RequestedActionKinds.empty())
  return llvm::find(Inputs.RequestedActionKinds, CodeAction::SOURCE_KIND) != Inputs.RequestedActionKinds.end();
// To accommodate clients without knowledge of source actions, we trigger
// without checking code action kinds when inside the preamble region.
return Inputs.SelectionEnd <= Inputs.AST->getPreambleBounds().Size;

68

we should respect Config::Diagnostics::Includes::IgnoreHeader here, similar to finding integration, similarly for missing include fixes.

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

1307

nit: Again extracting CodeActionInputs, setting filename explicitly and passing that into the call would be great here.

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

35–41

nit: you can use EXPECT_AVAILABLE and EXPECT_UNAVAILABLE directly for these two cases

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

101

type is not really giving any ideas what the parameter is about here, can you name it ?

[VitaNuo]

Address review comments.

[VitaNuo]

Thanks!

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

35–41

I think it's the other way around. I can use the macros directly when I don't need to pass specific requested action kinds, which is the last (third) case. So it's probably not worth it.

clang-tools-extra/clangd/ClangdLSPServer.cpp

 }
 
 const llvm::StringLiteral ApplyFixCommand = "clangd.applyFix";
 const llvm::StringLiteral ApplyTweakCommand = "clangd.applyTweak";
 
 /// 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) {
+                        Range Selection,
+                        const std::vector<std::string> &RequestedActionKinds) {
   CodeAction CA;
   CA.title = T.Title;
   CA.kind = T.Kind.str();
   // This tweak may have an expensive second stage, we only run it if the user
   // actually chooses it in the UI. We reply with a command that would run the
   // corresponding tweak.
   // FIXME: for some tweaks, computing the edits is cheap and we could send them
   //        directly.
   CA.command.emplace();
   CA.command->title = T.Title;
   CA.command->command = std::string(ApplyTweakCommand);
   TweakArgs Args;
   Args.file = File;
   Args.tweakID = T.ID;
   Args.selection = Selection;
+  Args.requestedActionKinds = RequestedActionKinds;
   CA.command->argument = std::move(Args);
   return CA;
 }
 
 /// Convert from Fix to LSP CodeAction.
 CodeAction toCodeAction(const Fix &F, const URIForFile &File,
                         const std::optional<int64_t> &Version,
                         bool SupportsDocumentChanges,
void ClangdLSPServer::onInitialize(const InitializeParams &Params,
   // Per LSP, codeActionProvider can be either boolean or CodeActionOptions.
   // CodeActionOptions is only valid if the client supports action literal
   // via textDocument.codeAction.codeActionLiteralSupport.
   ServerCaps["codeActionProvider"] =
       Params.capabilities.CodeActionStructure
           ? llvm::json::Object{{"codeActionKinds",
                                 {CodeAction::QUICKFIX_KIND,
                                  CodeAction::REFACTOR_KIND,
-                                 CodeAction::INFO_KIND}}}
+                                 CodeAction::INFO_KIND,
+                                 CodeAction::SOURCE_KIND}}}
           : llvm::json::Value(true);
 
   std::vector<llvm::StringRef> Commands;
   for (llvm::StringRef Command : Handlers.CommandHandlers.keys())
     Commands.push_back(Command);
   llvm::sort(Commands);
   ServerCaps["executeCommandProvider"] =
       llvm::json::Object{{"commands", Commands}};
auto Action = [this, Reply = std::move(Reply)](
     WE.changes.emplace();
     for (const auto &It : R->ApplyEdits) {
       (*WE.changes)[URI::createFile(It.first()).toString()] =
           It.second.asTextEdits();
     }
     // ApplyEdit will take care of calling Reply().
     return applyEdit(std::move(WE), "Tweak applied.", std::move(Reply));
   };
-  Server->applyTweak(Args.file.file(), Args.selection, Args.tweakID,
+  Server->applyTweak(Args.tweakID,
+                     {Args.file.file().str(),
+                      Args.selection,
+                      Args.requestedActionKinds,
+                      {},
+                      {}},
                      std::move(Action));
 }
 
 void ClangdLSPServer::applyEdit(WorkspaceEdit WE, llvm::json::Value Success,
                                 Callback<llvm::json::Value> Reply) {
   ApplyWorkspaceEditParams Edit;
   Edit.edit = std::move(WE);
   ApplyWorkspaceEdit(
for (const auto& LSPDiag : Params.context.diagnostics) {
     if (auto DiagRef = getDiagRef(File.file(), LSPDiag)) {
       ToLSPDiags[*DiagRef] = LSPDiag;
       Inputs.Diagnostics.push_back(*DiagRef);
     }
   }
   Inputs.File = File.file();
   Inputs.Selection = Params.range;
   Inputs.RequestedActionKinds = Params.context.only;
-  Inputs.TweakFilter = [this](const Tweak &T) {
-    return Opts.TweakFilter(T);
-  };
-  auto CB = [this,
-             Reply = std::move(Reply),
-             ToLSPDiags = std::move(ToLSPDiags), File,
-             Selection = Params.range](
+  Inputs.TweakFilter = [this](const Tweak &T) { return Opts.TweakFilter(T); };
+  auto CB = [this, Reply = std::move(Reply), ToLSPDiags = std::move(ToLSPDiags),
+             File, Selection = Params.range,
+             ActionKinds = Inputs.RequestedActionKinds](
                 llvm::Expected<ClangdServer::CodeActionResult> Fixits) mutable {
     if (!Fixits)
       return Reply(Fixits.takeError());
     std::vector<CodeAction> CAs;
     auto Version = decodeVersion(Fixits->Version);
     for (const auto &QF : Fixits->QuickFixes) {
       CAs.push_back(toCodeAction(QF.F, File, Version, SupportsDocumentChanges,
                                  SupportsChangeAnnotation));
-      if (auto It = ToLSPDiags.find(QF.Diag);
-          It != ToLSPDiags.end()) {
+      if (auto It = ToLSPDiags.find(QF.Diag); It != ToLSPDiags.end()) {
         CAs.back().diagnostics = {It->second};
       }
     }
     for (const auto &TR : Fixits->TweakRefs)
-      CAs.push_back(toCodeAction(TR, File, Selection));
+      CAs.push_back(toCodeAction(TR, File, Selection, ActionKinds));
 
     // If there's exactly one quick-fix, call it "preferred".
     // We never consider refactorings etc as preferred.
     CodeAction *OnlyFix = nullptr;
     for (auto &Action : CAs) {
       if (Action.kind && *Action.kind == CodeAction::QUICKFIX_KIND) {
         if (OnlyFix) {
           OnlyFix = nullptr;

clang-tools-extra/clangd/ClangdServer.h

struct CodeActionResult {
     std::vector<QuickFix> QuickFixes;
     std::vector<TweakRef> TweakRefs;
   };
   /// Surface code actions (quick-fixes for diagnostics, or available code
   /// tweaks) for a given range in a file.
   void codeAction(const CodeActionInputs &Inputs,
                   Callback<CodeActionResult> CB);
 
-  /// Apply the code tweak with a specified \p ID.
-  void applyTweak(PathRef File, Range Sel, StringRef ID,
+  /// Apply the code tweak with a specified \p tweakID.
+  void applyTweak(std::string TweakID, CodeActionInputs Inputs,
                   Callback<Tweak::Effect> CB);
 
   /// Called when an event occurs for a watched file in the workspace.
   void onFileEvent(const DidChangeWatchedFilesParams &Params);
 
   /// Get symbol info for given position.
   /// Clangd extension - not part of official LSP.
   void symbolInfo(PathRef File, Position Pos,

clang-tools-extra/clangd/ClangdServer.cpp

void ClangdServer::rename(PathRef File, Position Pos, llvm::StringRef NewName,
   };
   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,
-               llvm::vfs::FileSystem *FS) {
+               llvm::vfs::FileSystem *FS,
+               const std::vector<std::string> &RequestedActionKinds) {
   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), FS));
+  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), FS, RequestedActionKinds));
                               return false;
                             });
   assert(!Result.empty() && "Expected at least one SelectionTree");
   return std::move(Result);
 }
 
 void ClangdServer::codeAction(const CodeActionInputs &Params,
                               Callback<CodeActionResult> CB) {
   auto Action = [Params, CB = std::move(CB),
                  FeatureModules(this->FeatureModules)](
if (KindAllowed(CodeAction::QUICKFIX_KIND)) {
         return {};
       };
       for (const auto &Diag : Params.Diagnostics)
         for (const auto &Fix : FindMatchedFixes(Diag))
           Result.QuickFixes.push_back({Diag, Fix});
     }
 
     // Collect Tweaks
-    auto Selections = tweakSelection(Params.Selection, *InpAST, /*FS=*/nullptr);
+    auto Selections = tweakSelection(Params.Selection, *InpAST, /*FS=*/nullptr,
+                                     Params.RequestedActionKinds);
     if (!Selections)
       return CB(Selections.takeError());
     // Don't allow a tweak to fire more than once across ambiguous selections.
     llvm::DenseSet<llvm::StringRef> PreparedTweaks;
     auto DeduplicatingFilter = [&](const Tweak &T) {
       return KindAllowed(T.kind()) && Params.TweakFilter(T) &&
              !PreparedTweaks.count(T.id());
     };
     for (const auto &Sel : *Selections) {
       for (auto &T : prepareTweaks(*Sel, DeduplicatingFilter, FeatureModules)) {
         Result.TweakRefs.push_back(TweakRef{T->id(), T->title(), T->kind()});
         PreparedTweaks.insert(T->id());
         TweakAvailable.record(1, T->id());
       }
     }
     CB(std::move(Result));
   };
 
   WorkScheduler->runWithAST("codeAction", Params.File, std::move(Action),
                             Transient);
 }
 
-void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef TweakID,
+void ClangdServer::applyTweak(std::string TweakID, CodeActionInputs Inputs,
                               Callback<Tweak::Effect> CB) {
   // Tracks number of times a tweak has been attempted.
   static constexpr trace::Metric TweakAttempt(
       "tweak_attempt", trace::Metric::Counter, "tweak_id");
   // Tracks number of times a tweak has failed to produce edits.
   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),
+  auto Action = [File = Inputs.File, Sel = std::move(Inputs.Selection),
+                 TweakID = std::move(TweakID), CB = std::move(CB),
+                 ActionKinds = std::move(Inputs.RequestedActionKinds),
                  this](Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     auto FS = DirtyFS->view(std::nullopt);
-    auto Selections = tweakSelection(Sel, *InpAST, FS.get());
+    auto Selections = tweakSelection(Sel, *InpAST, FS.get(), ActionKinds);
     if (!Selections)
       return CB(Selections.takeError());
     std::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) {
if (*Effect && (*Effect)->FormatEdits) {
         if (llvm::Error Err = reformatEdit(E, Style))
           elog("Failed to format {0}: {1}", It.first(), std::move(Err));
       }
     } else {
       TweakFailed.record(1, TweakID);
     }
     return CB(std::move(*Effect));
   };
-  WorkScheduler->runWithAST("ApplyTweak", File, std::move(Action));
+  WorkScheduler->runWithAST("ApplyTweak", Inputs.File, std::move(Action));
 }
 
 void ClangdServer::locateSymbolAt(PathRef File, Position Pos,
                                   Callback<std::vector<LocatedSymbol>> CB) {
   auto Action = [Pos, CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());

clang-tools-extra/clangd/ParsedAST.h

 #include "CollectMacros.h"
 #include "Compiler.h"
 #include "Diagnostics.h"
 #include "Headers.h"
 #include "Preamble.h"
 #include "clang-include-cleaner/Record.h"
 #include "support/Path.h"
 #include "clang/Frontend/FrontendAction.h"
+#include "clang/Lex/Lexer.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
 #include <memory>
 #include <optional>
 #include <string>
 #include <vector>
public:
 
   /// Returns the path passed by the caller when building this AST.
   PathRef tuPath() const { return TUPath; }
 
   /// Returns the version of the ParseInputs used to build Preamble part of this
   /// AST. Might be std::nullopt if no Preamble is used.
   std::optional<llvm::StringRef> preambleVersion() const;
 
+  // Describes the bounds of the preamble.
+  PreambleBounds getPreambleBounds() const {
+    return Preamble->Preamble.getBounds();
+  }
+
   const HeuristicResolver *getHeuristicResolver() const {
     return Resolver.get();
   }
 
 private:
   ParsedAST(PathRef TUPath, llvm::StringRef Version,
             std::shared_ptr<const PreambleData> Preamble,
             std::unique_ptr<CompilerInstance> Clang,

clang-tools-extra/clangd/Protocol.h

 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_PROTOCOL_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_PROTOCOL_H
 
 #include "URI.h"
 #include "index/SymbolID.h"
 #include "support/MemoryTree.h"
 #include "clang/Index/IndexSymbol.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_ostream.h"
 #include <bitset>
 #include <memory>
 #include <optional>
 #include <string>
 #include <vector>
 

 /// tweak.
 struct TweakArgs {
   /// A file provided by the client on a textDocument/codeAction request.
   URIForFile file;
   /// A selection provided by the client on a textDocument/codeAction request.
   Range selection;
   /// ID of the tweak that should be executed. Corresponds to Tweak::id().
   std::string tweakID;
+  /// Code action kinds requested by the client via the `only` field in
+  /// the context.
+  std::vector<std::string> requestedActionKinds;
 };
 bool fromJSON(const llvm::json::Value &, TweakArgs &, llvm::json::Path);
 llvm::json::Value toJSON(const TweakArgs &A);
 
 struct ExecuteCommandParams {
   /// The identifier of the actual command handler.
   std::string command;
 
struct CodeAction {
   std::string title;
 
   /// The kind of the code action.
   /// Used to filter code actions.
   std::optional<std::string> kind;
   const static llvm::StringLiteral QUICKFIX_KIND;
   const static llvm::StringLiteral REFACTOR_KIND;
   const static llvm::StringLiteral INFO_KIND;
+  const static llvm::StringLiteral SOURCE_KIND;
 
   /// The diagnostics that this code action resolves.
   std::optional<std::vector<Diagnostic>> diagnostics;
 
   /// Marks this as a preferred action. Preferred actions are used by the
   /// `auto fix` command and can be targeted by keybindings.
   /// A quick fix should be marked preferred if it properly addresses the
   /// underlying error. A refactoring should be marked preferred if it is the

clang-tools-extra/clangd/Protocol.cpp

llvm::json::Value toJSON(const Command &C) {
   if (!C.argument.getAsNull())
     Cmd["arguments"] = llvm::json::Array{C.argument};
   return std::move(Cmd);
 }
 
 const llvm::StringLiteral CodeAction::QUICKFIX_KIND = "quickfix";
 const llvm::StringLiteral CodeAction::REFACTOR_KIND = "refactor";
 const llvm::StringLiteral CodeAction::INFO_KIND = "info";
+const llvm::StringLiteral CodeAction::SOURCE_KIND = "source";
 
 llvm::json::Value toJSON(const CodeAction &CA) {
   auto CodeAction = llvm::json::Object{{"title", CA.title}};
   if (CA.kind)
     CodeAction["kind"] = *CA.kind;
   if (CA.diagnostics)
     CodeAction["diagnostics"] = llvm::json::Array(*CA.diagnostics);
   if (CA.isPreferred)
llvm::json::Value toJSON(const WorkspaceEdit &WE) {
   }
   return Result;
 }
 
 bool fromJSON(const llvm::json::Value &Params, TweakArgs &A,
               llvm::json::Path P) {
   llvm::json::ObjectMapper O(Params, P);
   return O && O.map("file", A.file) && O.map("selection", A.selection) &&
-         O.map("tweakID", A.tweakID);
+         O.map("tweakID", A.tweakID) &&
+         O.map("requestedActionKinds", A.requestedActionKinds);
 }
 
 llvm::json::Value toJSON(const TweakArgs &A) {
-  return llvm::json::Object{
-      {"tweakID", A.tweakID}, {"selection", A.selection}, {"file", A.file}};
+  return llvm::json::Object{{"tweakID", A.tweakID},
+                            {"selection", A.selection},
+                            {"file", A.file},
+                            {"requestedActionKinds", A.requestedActionKinds}};
 }
 
 llvm::json::Value toJSON(const ApplyWorkspaceEditParams &Params) {
   return llvm::json::Object{{"edit", Params.edit}};
 }
 
 bool fromJSON(const llvm::json::Value &Response, ApplyWorkspaceEditResponse &R,
               llvm::json::Path P) {

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

 ///   };
 ///   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,
-              llvm::vfs::FileSystem *VFS);
+              llvm::vfs::FileSystem *VFS,
+              const std::vector<std::string> &RequestedActionKinds);
     /// 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;
+    /// Requested code action kinds from the `only` field of
+    /// code action request context.
+    std::vector<std::string> RequestedActionKinds;
     // FIXME: provide a way to get sources and ASTs for other files.
   };
 
   struct Effect {
     /// A message to be displayed to the user.
     std::optional<std::string> ShowMessage;
     FileEdits ApplyEdits;
     /// Whether the edits should be formatted before presenting to the client.

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

getAllTweaks(const FeatureModuleSet *Modules) {
   if (Modules) {
     for (auto &M : *Modules)
       M.contributeTweaks(All);
   }
   return All;
 }
 } // namespace
 
-Tweak::Selection::Selection(const SymbolIndex *Index, ParsedAST &AST,
-                            unsigned RangeBegin, unsigned RangeEnd,
-                            SelectionTree ASTSelection,
-                            llvm::vfs::FileSystem *FS)
+Tweak::Selection::Selection(
+    const SymbolIndex *Index, ParsedAST &AST, unsigned RangeBegin,
+    unsigned RangeEnd, SelectionTree ASTSelection, llvm::vfs::FileSystem *FS,
+    const std::vector<std::string> &RequestedActionKinds)
     : Index(Index), AST(&AST), SelectionBegin(RangeBegin),
-      SelectionEnd(RangeEnd), ASTSelection(std::move(ASTSelection)), FS(FS) {
+      SelectionEnd(RangeEnd), ASTSelection(std::move(ASTSelection)), FS(FS),
+      RequestedActionKinds(RequestedActionKinds) {
   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/CMakeLists.txt

add_clang_library(clangDaemonTweaks OBJECT
   DefineOutline.cpp
   ExpandDeducedType.cpp
   ExpandMacro.cpp
   ExtractFunction.cpp
   ExtractVariable.cpp
   MemberwiseConstructor.cpp
   ObjCLocalizeStringLiteral.cpp
   ObjCMemberwiseInitializer.cpp
+  OrganizeImports.cpp
   PopulateSwitch.cpp
   RawStringLiteral.cpp
   RemoveUsingNamespace.cpp
   SpecialMembers.cpp
   SwapIfBranches.cpp
 
   LINK_LIBS
   clangAST

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

+//===--- OrganizeImports.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 "IncludeCleaner.h"
+#include "Protocol.h"
+#include "SourceCode.h"
+#include "clang-include-cleaner/IncludeSpeller.h"
+#include "clang-include-cleaner/Types.h"
+#include "refactor/Tweak.h"
+#include "support/Logger.h"
+#include "clang/Basic/LLVM.h"
+#include "clang/Basic/SourceLocation.h"
+#include "clang/Format/Format.h"
+#include "clang/Tooling/Core/Replacement.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
+#include <algorithm>
+#include <cassert>
+#include <climits>
+#include <string>
+#include <vector>
+
+namespace clang::clangd {
+namespace {
+
+// Tweak for applying IWYU-related changes (removing unused and adding missing
+// includes) in batch. The tweak can be triggered via the "Organize Imports"
+// source action (VS Code). The set of changes applied fully corresponds to the
+// findings from the clang-include-cleaner tool.
+class OrganizeImports : public Tweak {
+public:
+  const char *id() const override;
+
+  bool prepare(const Selection &Inputs) override;
+  Expected<Effect> apply(const Selection &Inputs) override;
+  std::string title() const override {
+    return "Remove unused and add missing includes";
+  }
+  llvm::StringLiteral kind() const override { return CodeAction::SOURCE_KIND; }
+};
+REGISTER_TWEAK(OrganizeImports)
+
+bool OrganizeImports::prepare(const Tweak::Selection &Inputs) {
+  if (Inputs.AST->getLangOpts().ObjC)
+    return false;
+  if (std::find(Inputs.RequestedActionKinds.begin(),
+                Inputs.RequestedActionKinds.end(),
+                CodeAction::SOURCE_KIND) != Inputs.RequestedActionKinds.end())
+    return true;
+  if (Inputs.RequestedActionKinds.empty())
+    // To accommodate clients without knowledge of source actions, we trigger
+    // without checking code action kinds when inside the preamble region.
+    return offsetToPosition(Inputs.Code, Inputs.SelectionEnd) <=
+           offsetToPosition(Inputs.Code, Inputs.AST->getPreambleBounds().Size);
+  return false;
+}
+
+Expected<Tweak::Effect> OrganizeImports::apply(const Selection &Inputs) {
+  IncludeCleanerFindings Findings = computeIncludeCleanerFindings(*Inputs.AST);
+  const auto MainFilePath = Inputs.AST->tuPath();
+  tooling::Replacements Replacements;
+  for (const auto *Inc : Findings.UnusedIncludes)
+    llvm::cantFail(Replacements.add(
+        tooling::Replacement{MainFilePath, UINT_MAX, 1, Inc->Written}));
+
+  const auto &SM = Inputs.AST->getSourceManager();
+  llvm::DenseSet<include_cleaner::Header> Providers;
+  for (const auto &Missing : Findings.MissingIncludes) {
+    assert(!Missing.Providers.empty());
+    Providers.insert(Missing.Providers[0]);
+  }
+
+  for (const auto &P : Providers) {
+    std::string Spelling = include_cleaner::spellHeader(
+        {P, Inputs.AST->getPreprocessor().getHeaderSearchInfo(),
+         SM.getFileEntryForID(SM.getMainFileID())});
+    llvm::cantFail(Replacements.add(tooling::Replacement{
+        MainFilePath, UINT_MAX, 0, "#include " + Spelling}));
+  }
+
+  auto FileStyle =
+      format::getStyle(format::DefaultFormatStyle, MainFilePath,
+                       format::DefaultFallbackStyle, Inputs.Code, Inputs.FS);
+  if (!FileStyle) {
+    elog("Couldn't get style for {0}: {1}", MainFilePath,
+         FileStyle.takeError());
+    FileStyle = format::getLLVMStyle();
+  }
+  auto Final =
+      format::cleanupAroundReplacements(Inputs.Code, Replacements, *FileStyle);
+  if (!Final)
+    return Final.takeError();
+  if (Final->empty())
+    return Tweak::Effect{"No edits to apply.", {}};
+  return Effect::mainFileEdit(SM, *Final);
+}
+
+} // namespace
+} // namespace clang::clangd

clang-tools-extra/clangd/test/code-action-request.test

 }
 #      CHECK:  "id": 1,
 # CHECK-NEXT:  "jsonrpc": "2.0",
 # CHECK-NEXT:  "result": [
 # CHECK-NEXT:    {
 # CHECK-NEXT:      "arguments": [
 # CHECK-NEXT:        {
 # CHECK-NEXT:          "file": "file://{{.*}}/clangd-test/main.cpp",
+# CHECK-NEXT:          "requestedActionKinds": [],
 # CHECK-NEXT:          "selection": {
 # CHECK-NEXT:            "end": {
 # CHECK-NEXT:              "character": 4,
 # CHECK-NEXT:              "line": 0
 # CHECK-NEXT:            },
 # CHECK-NEXT:            "start": {
 # CHECK-NEXT:              "character": 0,
 # CHECK-NEXT:              "line": 0
"textDocument": { "uri": "test:///main.cpp" },
         }
     }
 }
 #      CHECK:  "id": 3,
 # CHECK-NEXT:  "jsonrpc": "2.0",
 # CHECK-NEXT:  "result": [
 # CHECK-NEXT:    {
 ---
-{"jsonrpc":"2.0","id":4,"method":"workspace/executeCommand","params":{"command":"clangd.applyTweak","arguments":[{"file":"test:///main.cpp","selection":{"end":{"character":4,"line":0},"start":{"character":0,"line":0}},"tweakID":"ExpandDeducedType"}]}}
+{"jsonrpc":"2.0","id":4,"method":"workspace/executeCommand","params":{"command":"clangd.applyTweak","arguments":[{"file":"test:///main.cpp","requestedActionKinds": [],"selection":{"end":{"character":4,"line":0},"start":{"character":0,"line":0}},"tweakID":"ExpandDeducedType"}]}}
 #      CHECK:    "newText": "int",
 # CHECK-NEXT:    "range": {
 # CHECK-NEXT:      "end": {
 # CHECK-NEXT:        "character": 4,
 # CHECK-NEXT:        "line": 0
 # CHECK-NEXT:      },
 # CHECK-NEXT:      "start": {
 # CHECK-NEXT:        "character": 0,
 # CHECK-NEXT:        "line": 0
 # CHECK-NEXT:      }
 # CHECK-NEXT:    }
 ---
 {"jsonrpc":"2.0","id":5,"method":"shutdown"}
 ---
 {"jsonrpc":"2.0","method":"exit"}
 ---

clang-tools-extra/clangd/test/include-cleaner-batch-fix.test

 # CHECK-NEXT:                "version": 0
 # CHECK-NEXT:              }
 # CHECK-NEXT:            }
 # CHECK-NEXT:          ]
 # CHECK-NEXT:        }
 # CHECK-NEXT:      ],
 # CHECK-NEXT:      "command": "clangd.applyFix",
 # CHECK-NEXT:      "title": "Apply fix: fix all includes"
+# CHECK-NEXT:    },
+# CHECK-NEXT: {
+# CHECK-NEXT:      "arguments": [
+# CHECK-NEXT:        {
+# CHECK-NEXT:          "file": "file:///clangd-test/simple.cpp",
+# CHECK-NEXT:          "requestedActionKinds": [],
+# CHECK-NEXT:          "selection": {
+# CHECK-NEXT:            "end": {
+# CHECK-NEXT:              "character": 17,
+# CHECK-NEXT:              "line": 0
+# CHECK-NEXT:            },
+# CHECK-NEXT:            "start": {
+# CHECK-NEXT:              "character": 0,
+# CHECK-NEXT:              "line": 0
+# CHECK-NEXT:            }
+# CHECK-NEXT:          },
+# CHECK-NEXT:          "tweakID": "OrganizeImports"
+# CHECK-NEXT:        }
+# CHECK-NEXT:      ],
+# CHECK-NEXT:      "command": "clangd.applyTweak",
+# CHECK-NEXT:      "title": "Remove unused and add missing includes"
 # CHECK-NEXT:    }
 # CHECK-NEXT:  ]
 ---
 {"jsonrpc":"2.0","id":4,"method":"shutdown"}
 ---
 {"jsonrpc":"2.0","method":"exit"}

clang-tools-extra/clangd/test/request-reply.test

 # RUN: clangd -log=verbose -lit-test < %s | FileCheck -strict-whitespace %s
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
 ---
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{"uri":"test:///main.cpp","languageId":"cpp","version":1,"text":"auto i = 0;"}}}
 ---
-{"jsonrpc":"2.0","id":4,"method":"workspace/executeCommand","params":{"command":"clangd.applyTweak","arguments":[{"file":"test:///main.cpp","selection":{"end":{"character":4,"line":0},"start":{"character":0,"line":0}},"tweakID":"ExpandDeducedType"}]}}
+{"jsonrpc":"2.0","id":4,"method":"workspace/executeCommand","params":{"command":"clangd.applyTweak","arguments":[{"file":"test:///main.cpp","requestedActionKinds":[],"selection":{"end":{"character":4,"line":0},"start":{"character":0,"line":0}},"tweakID":"ExpandDeducedType"}]}}
 #      CHECK:  "id": 0,
 #      CHECK:  "method": "workspace/applyEdit",
 #      CHECK:  "newText": "int",
 # CHECK-NEXT:  "range": {
 # CHECK-NEXT:    "end": {
 # CHECK-NEXT:      "character": 4,
 # CHECK-NEXT:      "line": 0
 # CHECK-NEXT:    },
 # CHECK-NEXT:    "start": {
 # CHECK-NEXT:      "character": 0,
 # CHECK-NEXT:      "line": 0
 # CHECK-NEXT:    }
 # CHECK-NEXT:  }
 ---
 {"jsonrpc":"2.0","id":0,"result":{"applied":false}}
 #      CHECK:  "error": {
 # CHECK-NEXT:    "code": -32001,
 # CHECK-NEXT:    "message": "edits were not applied: unknown reason"
 # CHECK-NEXT:  },
 # CHECK-NEXT:  "id": 4,
 ---
-{"jsonrpc":"2.0","id":5,"method":"workspace/executeCommand","params":{"command":"clangd.applyTweak","arguments":[{"file":"test:///main.cpp","selection":{"end":{"character":4,"line":0},"start":{"character":0,"line":0}},"tweakID":"ExpandDeducedType"}]}}
+{"jsonrpc":"2.0","id":5,"method":"workspace/executeCommand","params":{"command":"clangd.applyTweak","arguments":[{"file":"test:///main.cpp","requestedActionKinds":[],"selection":{"end":{"character":4,"line":0},"start":{"character":0,"line":0}},"tweakID":"ExpandDeducedType"}]}}
 #      CHECK:  "id": 1,
 #      CHECK:  "method": "workspace/applyEdit",
 ---
 {"jsonrpc":"2.0","id":1,"result":{"applied":true}}
 #      CHECK:  "id": 5,
 # CHECK-NEXT:  "jsonrpc": "2.0",
 # CHECK-NEXT:  "result": "Tweak applied."
 ---
 # a reply with an invalid id.
 {"jsonrpc":"2.0","id":"invalid","result":{"applied":true}}
 # clangd doesn't reply, just emits an elog.
 ---
 {"jsonrpc":"2.0","id":6,"method":"shutdown"}
 ---
 {"jsonrpc":"2.0","method":"exit"}

clang-tools-extra/clangd/test/tweaks-format.test

 # RUN: clangd -lit-test < %s | FileCheck %s
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{},"trace":"off"}}
 ---
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{"uri":"test:///main.cc","languageId":"cpp","version":1,"text":"int f() { if (true) { return 1; } else {} }"}}}
 ---
 {"jsonrpc":"2.0","id":5,"method":"textDocument/codeAction","params":{"textDocument":{"uri":"test:///main.cc"},"range":{"start":{"line":0,"character":11},"end":{"line":0,"character":11}},"context":{"diagnostics":[]}}}
 ---
-{"jsonrpc":"2.0","id":6,"method":"workspace/executeCommand","params":{"command":"clangd.applyTweak","arguments":[{"file":"test:///main.cc","selection":{"end":{"character":11,"line":0},"start":{"character":11,"line":0}},"tweakID":"SwapIfBranches"}]}}
+{"jsonrpc":"2.0","id":6,"method":"workspace/executeCommand","params":{"command":"clangd.applyTweak","arguments":[{"file":"test:///main.cc","requestedActionKinds":[],"selection":{"end":{"character":11,"line":0},"start":{"character":11,"line":0}},"tweakID":"SwapIfBranches"}]}}
 #      CHECK:    "newText": "\n  ",
 # CHECK-NEXT:    "range": {
 # CHECK-NEXT:      "end": {
 # CHECK-NEXT:        "character": 10,
 # CHECK-NEXT:        "line": 0
 # CHECK-NEXT:      },
 # CHECK-NEXT:      "start": {
 # CHECK-NEXT:        "character": 9,

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

for (const auto &Tok : SpelledTokens) {
       SPAN_ATTACH(Trace, "text", Tok.text(AST->getSourceManager()));
 
       // 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),
-                                 nullptr);
+                                 nullptr, {});
       // FS is only populated when applying a tweak, not during prepare as
       // prepare should not do any I/O to be fast.
       auto Tweaks =
           prepareTweaks(Selection, Opts.TweakFilter, Opts.FeatureModules);
       Selection.FS =
           &AST->getSourceManager().getFileManager().getVirtualFileSystem();
       for (const auto &T : Tweaks) {
         auto Result = T->apply(Selection);

clang-tools-extra/clangd/unittests/CMakeLists.txt

add_unittest(ClangdUnitTests ClangdTests
   tweaks/DumpSymbolTests.cpp
   tweaks/ExpandDeducedTypeTests.cpp
   tweaks/ExpandMacroTests.cpp
   tweaks/ExtractFunctionTests.cpp
   tweaks/ExtractVariableTests.cpp
   tweaks/MemberwiseConstructorTests.cpp
   tweaks/ObjCLocalizeStringLiteralTests.cpp
   tweaks/ObjCMemberwiseInitializerTests.cpp
+  tweaks/OrganizeImportsTests.cpp
   tweaks/PopulateSwitchTests.cpp
   tweaks/RawStringLiteralTests.cpp
   tweaks/RemoveUsingNamespaceTests.cpp
   tweaks/ShowSelectionTreeTests.cpp
   tweaks/SpecialMembersTests.cpp
   tweaks/SwapIfBranchesTests.cpp
   tweaks/TweakTesting.cpp
   tweaks/TweakTests.cpp

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

TEST(ClangdServer, RespectsTweakFormatting) {
   ClangdServer Server(CDB, FS, Opts);
 
   auto FooCpp = testPath("foo.cpp");
   Server.addDocument(FooCpp, "");
   ASSERT_TRUE(Server.blockUntilIdleForTest());
 
   // Ensure that disabled formatting is respected.
   Notification N;
-  Server.applyTweak(FooCpp, {}, TweakID, [&](llvm::Expected<Tweak::Effect> E) {
+  Server.applyTweak(
+      TweakID, {FooCpp, {}, {}, {}, {}}, [&](llvm::Expected<Tweak::Effect> E) {
         ASSERT_TRUE(static_cast<bool>(E));
         EXPECT_THAT(llvm::cantFail(E->ApplyEdits.lookup(FooCpp).apply()),
                     NewContents);
         N.notify();
       });
   N.wait();
 }
 
 TEST(ClangdServer, InactiveRegions) {
   struct InactiveRegionsCallback : ClangdServer::Callbacks {
     std::vector<std::vector<Range>> FoundInactiveRegions;
 
     void onInactiveRegionsReady(PathRef FIle,

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), nullptr),
-      &Set);
+      TweakID,
+      Tweak::Selection(nullptr, AST, 0, 0, std::move(Tree), nullptr, {}), &Set);
   ASSERT_TRUE(bool(Actual));
   EXPECT_EQ(Actual->get()->id(), TweakID);
 }
 
 TEST(FeatureModulesTest, SuppressDiags) {
   struct DiagModifierModule final : public FeatureModule {
     struct Listener : public FeatureModule::ASTListener {
       void sawDiagnostic(const clang::Diagnostic &Info,

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

+//===--OrganizeImportsTest.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 "Protocol.h"
+#include "TweakTesting.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Testing/Annotations/Annotations.h"
+#include "gtest/gtest.h"
+#include <string>
+#include <vector>
+
+namespace clang {
+namespace clangd {
+namespace {
+
+TWEAK_TEST(OrganizeImports);
+
+TEST_F(OrganizeImportsTest, Prepare) {
+  struct {
+    llvm::StringRef Code;
+    bool IsAvailable;
+    const std::vector<std::string> RequestedActionKinds;
+  } Cases[] = {{
+                   R"cpp(
+#include "Te^stTU.h"
+)cpp",
+                   true,
+                   {CodeAction::SOURCE_KIND.str()}},
+               {"void foo(^) {}", true, {CodeAction::SOURCE_KIND.str()}},
+               {
+                   R"cpp(
+#include "Te^stTU.h"
+)cpp",
+                   true,
+                   {}},
+               {"void foo(^) {}", false, {}}};
+
+  for (auto Case : Cases) {
+    llvm::Annotations A{Case.Code};
+    auto AST = build(A.code());
+    for (const auto &P : A.points())
+      EXPECT_EQ(Case.IsAvailable,
+                isAvailable(AST, {P, P}, Case.RequestedActionKinds))
+          << decorate(A.code(), P);
+  }
+}
+
+TEST_F(OrganizeImportsTest, Apply) {
+  Header = "void foo();";
+  struct {
+    llvm::StringRef Code;
+    llvm::StringRef Header;
+    llvm::StringRef Expected;
+  } Cases[] = {{// Remove unused include.
+                R"cpp(
+#include "TestTU.h"
+void foo() {}
+void b^ar() {
+    foo();
+}
+)cpp",
+                R"cpp(
+#pragma once
+)cpp",
+                R"cpp(
+void foo() {}
+void bar() {
+    foo();
+}
+)cpp"},
+               {// Add missing include.
+                R"cpp(
+void b^ar() {
+    foo();
+}
+)cpp",
+                R"cpp(
+#pragma once
+void foo();
+)cpp",
+                R"cpp(
+#include "TestTU.h"
+void bar() {
+    foo();
+}
+)cpp"},
+               {// Replace unused include with missing.
+                R"cpp(
+#include "foo.h"
+void b^ar() {
+    foo();
+}
+)cpp",
+                R"cpp(
+#pragma once
+void foo();
+)cpp",
+                R"cpp(
+#include "TestTU.h"
+void bar() {
+    foo();
+}
+)cpp"}};
+  for (auto C : Cases) {
+    Header = C.Header;
+    ExtraFiles["foo.h"] = "#pragma once";
+    EXPECT_EQ(C.Expected,
+              apply(C.Code, nullptr, {CodeAction::SOURCE_KIND.str()}))
+        << C.Expected;
+  }
+}
+
+} // namespace
+} // namespace clangd
+} // namespace clang

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

 #include "index/Index.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Testing/Annotations/Annotations.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <memory>
 #include <string>
+#include <vector>
 
 namespace clang {
 namespace clangd {
 
 // Fixture base for testing tweaks. Intended to be subclassed for each tweak.
 //
 // Usage:
 // TWEAK_TEST(ExpandDeducedType);
protected:
   //    it is non-null. It is a mapping from absolute path of the edited file to
   //    its new contents. Passing a nullptr to \p EditedFiles when there are
   //    changes, will result in a failure.
   //    The context added to MarkedCode will be stripped away before returning,
   //    unless the tweak edited it.
   //  - if the tweak produces a message, returns "message:\n<message>"
   //  - if prepare() returns false, returns "unavailable"
   //  - if apply() returns an error, returns "fail: <message>"
-  std::string apply(llvm::StringRef MarkedCode,
-                    llvm::StringMap<std::string> *EditedFiles = nullptr) const;
+  std::string
+  apply(llvm::StringRef MarkedCode,
+        llvm::StringMap<std::string> *EditedFiles = nullptr,
+        const std::vector<std::string> &RequestedActionKinds = {}) const;
 
   // Helpers for EXPECT_AVAILABLE/EXPECT_UNAVAILABLE macros.
   using WrappedAST = std::pair<ParsedAST, /*WrappingOffset*/ unsigned>;
   WrappedAST build(llvm::StringRef) const;
-  bool isAvailable(WrappedAST &, llvm::Annotations::Range) const;
+  bool isAvailable(WrappedAST &, llvm::Annotations::Range,
+                   const std::vector<std::string> & = {}) const;
   // Return code re-decorated with a single point/range.
   static std::string decorate(llvm::StringRef, unsigned);
   static std::string decorate(llvm::StringRef, llvm::Annotations::Range);
 };
 
 MATCHER_P2(FileWithContents, FileName, Contents, "") {
   return arg.first() == FileName && arg.second == Contents;
 }
 
 #define TWEAK_TEST(TweakID)                                                    \
   class TweakID##Test : public ::clang::clangd::TweakTest {                    \
   protected:                                                                   \
     TweakID##Test() : TweakTest(#TweakID) {}                                   \
   }
 
 #define EXPECT_AVAILABLE_(MarkedCode, Available)                               \
   do {                                                                         \
     llvm::Annotations A{llvm::StringRef(MarkedCode)};                          \
     auto AST = build(A.code());                                                \
     assert(!A.points().empty() || !A.ranges().empty());                        \
     for (const auto &P : A.points())                                           \
-      EXPECT_EQ(Available, isAvailable(AST, {P, P})) << decorate(A.code(), P); \
+      EXPECT_EQ(Available, isAvailable(AST, {P, P}, {}))                       \
+          << decorate(A.code(), P);                                            \
     for (const auto &R : A.ranges())                                           \
-      EXPECT_EQ(Available, isAvailable(AST, R)) << decorate(A.code(), R);      \
+      EXPECT_EQ(Available, isAvailable(AST, R, {})) << decorate(A.code(), R);  \
   } while (0)
 #define EXPECT_AVAILABLE(MarkedCode) EXPECT_AVAILABLE_(MarkedCode, true)
 #define EXPECT_UNAVAILABLE(MarkedCode) EXPECT_AVAILABLE_(MarkedCode, false)
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/clangd/unittests/tweaks/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 "Protocol.h"
 #include "SourceCode.h"
 #include "TestTU.h"
 #include "refactor/Tweak.h"
 #include "llvm/Support/Error.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <optional>
 #include <string>
+#include <vector>
 
 namespace clang {
 namespace clangd {
 namespace {
 using Context = TweakTest::CodeContext;
 
 std::pair<llvm::StringRef, llvm::StringRef> wrapping(Context Ctx) {
   switch (Ctx) {
llvm::Annotations::Range rangeOrPoint(const llvm::Annotations &A) {
   }
   return A.range();
 }
 
 // Prepare and apply the specified tweak based on the selection in Input.
 // Returns std::nullopt if and only if prepare() failed.
 std::optional<llvm::Expected<Tweak::Effect>>
 applyTweak(ParsedAST &AST, llvm::Annotations::Range Range, StringRef TweakID,
-           const SymbolIndex *Index, llvm::vfs::FileSystem *FS) {
+           const SymbolIndex *Index, llvm::vfs::FileSystem *FS,
+           const std::vector<std::string> &RequestedActionKinds) {
   std::optional<llvm::Expected<Tweak::Effect>> Result;
   SelectionTree::createEach(AST.getASTContext(), AST.getTokens(), Range.Begin,
                             Range.End, [&](SelectionTree ST) {
                               Tweak::Selection S(Index, AST, Range.Begin,
-                                                 Range.End, std::move(ST), FS);
+                                                 Range.End, std::move(ST), FS,
+                                                 RequestedActionKinds);
                               if (auto T = prepareTweak(TweakID, S, nullptr)) {
                                 Result = (*T)->apply(S);
                                 return true;
                               } else {
                                 llvm::consumeError(T.takeError());
                                 return false;
                               }
                             });
   return Result;
 }
 
 } // namespace
 
-std::string TweakTest::apply(llvm::StringRef MarkedCode,
-                             llvm::StringMap<std::string> *EditedFiles) const {
+std::string
+TweakTest::apply(llvm::StringRef MarkedCode,
+                 llvm::StringMap<std::string> *EditedFiles,
+                 const std::vector<std::string> &RequestedActionKinds) const {
   std::string WrappedCode = wrap(Context, MarkedCode);
   llvm::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, rangeOrPoint(Input), TweakID, Index.get(),
-      &AST.getSourceManager().getFileManager().getVirtualFileSystem());
+      &AST.getSourceManager().getFileManager().getVirtualFileSystem(),
+      RequestedActionKinds);
   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())
else {
                          "provided a nullptr for EditedFiles.";
       else
         EditedFiles->insert_or_assign(It.first(), Unwrapped.str());
     }
   }
   return EditedMainFile;
 }
 
-bool TweakTest::isAvailable(WrappedAST &AST,
-                            llvm::Annotations::Range Range) const {
+bool TweakTest::isAvailable(
+    WrappedAST &AST, llvm::Annotations::Range Range,
+    const std::vector<std::string> &RequestedActionKinds) const {
   // Adjust range for wrapping offset.
   Range.Begin += AST.second;
   Range.End += AST.second;
   auto Result = applyTweak(
       AST.first, Range, TweakID, Index.get(),
-      &AST.first.getSourceManager().getFileManager().getVirtualFileSystem());
+      &AST.first.getSourceManager().getFileManager().getVirtualFileSystem(),
+      RequestedActionKinds);
   // We only care if prepare() succeeded, but must handle Errors.
   if (Result && !*Result)
     consumeError(Result->takeError());
   return Result.has_value();
 }
 
 TweakTest::WrappedAST TweakTest::build(llvm::StringRef Code) const {
   TestTU TU;