[clangd] Enable modules to contribute tweaks.

First patch to enable diagnostic fix generation through modules. The
workflow will look like:
- ASTWorker letting modules know about diagnostics while building AST,
modules can read clang::Diagnostic and mutate clangd::Diagnostic through
that hook.
- Modules can implement and expose tweaks to fix diagnostics or act as
general refactorings.
- Tweak::Selection will contain information about the diagnostic
associated with the codeAction request to enable modules to fail their
diagnostic fixing tweakson prepare if need be.

Differential Revision: https://reviews.llvm.org/D98498

clang-tools-extra/clangd/ClangdServer.cpp

@@ -573,8 +573,9 @@ void ClangdServer::enumerateTweaks(
   static constexpr trace::Metric TweakAvailable(
       "tweak_available", trace::Metric::Counter, "tweak_id");
   auto Action = [File = File.str(), Sel, CB = std::move(CB),
-                 Filter =
-                     std::move(Filter)](Expected<InputsAndAST> InpAST) mutable {
+                 Filter = std::move(Filter),
+                 FeatureModules(this->FeatureModules)](
+                    Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
     auto Selections = tweakSelection(Sel, *InpAST);
@@ -587,7 +588,7 @@ void ClangdServer::enumerateTweaks(
       return Filter(T) && !PreparedTweaks.count(T.id());
     };
     for (const auto &Sel : *Selections) {
-      for (auto &T : prepareTweaks(*Sel, DeduplicatingFilter)) {
+      for (auto &T : prepareTweaks(*Sel, DeduplicatingFilter, FeatureModules)) {
         Res.push_back({T->id(), T->title(), T->kind()});
         PreparedTweaks.insert(T->id());
         TweakAvailable.record(1, T->id());
@@ -622,7 +623,7 @@ void ClangdServer::applyTweak(PathRef File, Range Sel, StringRef TweakID,
     // 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);
+      auto T = prepareTweak(TweakID, *Selection, FeatureModules);
       if (T) {
         Effect = (*T)->apply(*Selection);
         break;

clang-tools-extra/clangd/FeatureModule.h

@@ -25,6 +25,7 @@ class LSPBinder;
 class SymbolIndex;
 class ThreadsafeFS;
 class TUScheduler;
+class Tweak;
 
 /// A FeatureModule contributes a vertical feature to clangd.
 ///
@@ -91,6 +92,10 @@ public:
   /// Called by the server when shutting down, and also by tests.
   virtual bool blockUntilIdle(Deadline) { return true; }
 
+  /// Tweaks implemented by this module. Can be called asynchronously when
+  /// enumerating or applying code actions.
+  virtual void contributeTweaks(std::vector<std::unique_ptr<Tweak>> &Out) {}
+
 protected:
   /// Accessors for modules to access shared server facilities they depend on.
   Facilities &facilities();

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

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 #include "Tweak.h"
+#include "FeatureModule.h"
 #include "SourceCode.h"
 #include "index/Index.h"
 #include "support/Logger.h"
@@ -20,6 +21,7 @@
 #include <functional>
 #include <memory>
 #include <utility>
+#include <vector>
 
 LLVM_INSTANTIATE_REGISTRY(llvm::Registry<clang::clangd::Tweak>)
 
@@ -43,6 +45,18 @@ void validateRegistry() {
   }
 #endif
 }
+
+std::vector<std::unique_ptr<Tweak>>
+getAllTweaks(const FeatureModuleSet *Modules) {
+  std::vector<std::unique_ptr<Tweak>> All;
+  for (const auto &E : TweakRegistry::entries())
+    All.emplace_back(E.instantiate());
+  if (Modules) {
+    for (auto &M : *Modules)
+      M.contributeTweaks(All);
+  }
+  return All;
+}
 } // namespace
 
 Tweak::Selection::Selection(const SymbolIndex *Index, ParsedAST &AST,
@@ -57,12 +71,12 @@ Tweak::Selection::Selection(const SymbolIndex *Index, ParsedAST &AST,
 
 std::vector<std::unique_ptr<Tweak>>
 prepareTweaks(const Tweak::Selection &S,
-              llvm::function_ref<bool(const Tweak &)> Filter) {
+              llvm::function_ref<bool(const Tweak &)> Filter,
+              const FeatureModuleSet *Modules) {
   validateRegistry();
 
   std::vector<std::unique_ptr<Tweak>> Available;
-  for (const auto &E : TweakRegistry::entries()) {
-    std::unique_ptr<Tweak> T = E.instantiate();
+  for (auto &T : getAllTweaks(Modules)) {
     if (!Filter(*T) || !T->prepare(S))
       continue;
     Available.push_back(std::move(T));
@@ -74,18 +88,18 @@ prepareTweaks(const Tweak::Selection &S,
   return Available;
 }
 
-llvm::Expected<std::unique_ptr<Tweak>> prepareTweak(StringRef ID,
-                                                    const Tweak::Selection &S) {
-  auto It = llvm::find_if(
-      TweakRegistry::entries(),
-      [ID](const TweakRegistry::entry &E) { return E.getName() == ID; });
-  if (It == TweakRegistry::end())
-    return error("tweak ID {0} is invalid", ID);
-  std::unique_ptr<Tweak> T = It->instantiate();
+llvm::Expected<std::unique_ptr<Tweak>>
+prepareTweak(StringRef ID, const Tweak::Selection &S,
+             const FeatureModuleSet *Modules) {
+  for (auto &T : getAllTweaks(Modules)) {
+    if (T->id() != ID)
+      continue;
     if (!T->prepare(S))
       return error("failed to prepare() tweak {0}", ID);
     return std::move(T);
   }
+  return error("tweak ID {0} is invalid", ID);
+}
 
 llvm::Expected<std::pair<Path, Edit>>
 Tweak::Effect::fileEdit(const SourceManager &SM, FileID FID,

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

@@ -35,6 +35,8 @@
 namespace clang {
 namespace clangd {
 
+class FeatureModuleSet;
+
 /// An interface base for small context-sensitive refactoring actions.
 /// To implement a new tweak use the following pattern in a .cpp file:
 ///   class MyTweak : public Tweak {
@@ -128,13 +130,15 @@ public:
 /// can run on the selection.
 std::vector<std::unique_ptr<Tweak>>
 prepareTweaks(const Tweak::Selection &S,
-              llvm::function_ref<bool(const Tweak &)> Filter);
+              llvm::function_ref<bool(const Tweak &)> Filter,
+              const FeatureModuleSet *Modules);
 
 // Calls prepare() on the tweak with a given ID.
 // If prepare() returns false, returns an error.
 // If prepare() returns true, returns the corresponding tweak.
-llvm::Expected<std::unique_ptr<Tweak>> prepareTweak(StringRef TweakID,
-                                                    const Tweak::Selection &S);
+llvm::Expected<std::unique_ptr<Tweak>>
+prepareTweak(StringRef ID, const Tweak::Selection &S,
+             const FeatureModuleSet *Modules);
 } // namespace clangd
 } // namespace clang
 

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

@@ -211,7 +211,8 @@ public:
       auto Tree = SelectionTree::createRight(AST->getASTContext(),
                                              AST->getTokens(), Start, End);
       Tweak::Selection Selection(&Index, *AST, Start, End, std::move(Tree));
-      for (const auto &T : prepareTweaks(Selection, Opts.TweakFilter)) {
+      for (const auto &T :
+           prepareTweaks(Selection, Opts.TweakFilter, Opts.FeatureModules)) {
         auto Result = T->apply(Selection);
         if (!Result) {
           elog("    tweak: {0} ==> FAIL: {1}", T->id(), Result.takeError());

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

@@ -55,6 +55,7 @@ add_unittest(ClangdUnitTests ClangdTests
   DraftStoreTests.cpp
   DumpASTTests.cpp
   ExpectedTypeTest.cpp
+  FeatureModulesTests.cpp
   FileDistanceTests.cpp
   FileIndexTests.cpp
   FindSymbolsTests.cpp

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

@@ -0,0 +1,55 @@
+//===--- FeatureModulesTests.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 "FeatureModule.h"
+#include "Selection.h"
+#include "TestTU.h"
+#include "refactor/Tweak.h"
+#include "support/Logger.h"
+#include "llvm/Support/Error.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include <memory>
+
+namespace clang {
+namespace clangd {
+namespace {
+
+TEST(FeatureModulesTest, ContributesTweak) {
+  static constexpr const char *TweakID = "ModuleTweak";
+  struct TweakContributingModule final : public FeatureModule {
+    struct ModuleTweak final : public Tweak {
+      const char *id() const override { return TweakID; }
+      bool prepare(const Selection &Sel) override { return true; }
+      Expected<Effect> apply(const Selection &Sel) override {
+        return error("not implemented");
+      }
+      std::string title() const override { return id(); }
+      llvm::StringLiteral kind() const override { return ""; };
+    };
+
+    void contributeTweaks(std::vector<std::unique_ptr<Tweak>> &Out) override {
+      Out.emplace_back(new ModuleTweak);
+    }
+  };
+
+  FeatureModuleSet Set;
+  Set.add(std::make_unique<TweakContributingModule>());
+
+  auto AST = TestTU::withCode("").build();
+  auto Tree =
+      SelectionTree::createRight(AST.getASTContext(), AST.getTokens(), 0, 0);
+  auto Actual = prepareTweak(
+      TweakID, Tweak::Selection(nullptr, AST, 0, 0, std::move(Tree)), &Set);
+  ASSERT_TRUE(bool(Actual));
+  EXPECT_EQ(Actual->get()->id(), TweakID);
+}
+
+} // namespace
+} // namespace clangd
+} // namespace clang

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

@@ -75,7 +75,7 @@ applyTweak(ParsedAST &AST, const Annotations &Input, StringRef TweakID,
                             Range.second, [&](SelectionTree ST) {
                               Tweak::Selection S(Index, AST, Range.first,
                                                  Range.second, std::move(ST));
-                              if (auto T = prepareTweak(TweakID, S)) {
+                              if (auto T = prepareTweak(TweakID, S, nullptr)) {
                                 Result = (*T)->apply(S);
                                 return true;
                               } else {