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

[VitaNuo]

Address comments.

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;