[clangd] Enable modules to contribute tweaks.

[kadircet]

Summary

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.

[sammccall]

Should we have a test somewhere that tweaks defined in modules actually work?

clang-tools-extra/clangd/FeatureModule.h

98

nit: can you move the forward declaration up the top? Makes the (sort-of) dep clearer, and is less fragile re scopes when moving code around.

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

55

nit: the early return seems more confusing to me here than conditionally adding...

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

132

forward decl: and here

132

nit: move the modules to the end of the list?

[kadircet]

• Add test
• Move forward declarations out of signatures
• Move FeatureModuleSets to the last parameter in prepareTweaks

clang-tools-extra/clangd/ClangdServer.cpp

 
 void ClangdServer::enumerateTweaks(
     PathRef File, Range Sel, llvm::unique_function<bool(const Tweak &)> Filter,
     Callback<std::vector<TweakRef>> CB) {
   // Tracks number of times a tweak has been offered.
   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);
     if (!Selections)
       return CB(Selections.takeError());
     std::vector<TweakRef> Res;
     // Don't allow a tweak to fire more than once across ambiguous selections.
     llvm::DenseSet<llvm::StringRef> PreparedTweaks;
     auto DeduplicatingFilter = [&](const Tweak &T) {
       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());
       }
     }
 
     CB(std::move(Res));
   };
if (!InpAST)
       return CB(InpAST.takeError());
     auto Selections = tweakSelection(Sel, *InpAST);
     if (!Selections)
       return CB(Selections.takeError());
     llvm::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);
+      auto T = prepareTweak(TweakID, *Selection, FeatureModules);
       if (T) {
         Effect = (*T)->apply(*Selection);
         break;
       }
       Effect = T.takeError();
     }
     assert(Effect.hasValue() && "Expected at least one selection");
     if (*Effect) {

clang-tools-extra/clangd/FeatureModule.h

 #include <vector>
 
 namespace clang {
 namespace clangd {
 class LSPBinder;
 class SymbolIndex;
 class ThreadsafeFS;
 class TUScheduler;
+class Tweak;
 
 /// A FeatureModule contributes a vertical feature to clangd.
 ///
 /// The lifetime of a module is roughly:
 ///  - feature modules are created before the LSP server, in ClangdMain.cpp
 ///  - these modules are then passed to ClangdLSPServer in a FeatureModuleSet
 ///  - initializeLSP() is called when the editor calls initialize.
 //   - initialize() is then called by ClangdServer as it is constructed.
public:
 
   /// Waits until the module is idle (no background work) or a deadline expires.
   /// In general all modules should eventually go idle, though it may take a
   /// long time (e.g. background indexing).
   /// FeatureModules should go idle quickly if stop() has been called.
   /// 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();
   /// The scheduler is used to run tasks on worker threads and access ASTs.
   TUScheduler &scheduler() { return facilities().Scheduler; }
   /// The index is used to get information about the whole codebase.
   const SymbolIndex *index() { return facilities().Index; }
   /// The filesystem is used to read source files on disk.

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

 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include <string>
 
 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 {
 ///   public:
 ///     const char* id() const override final; // defined by REGISTER_TWEAK.
 ///     // implement other methods here.
 ///   };
 ///   REGISTER_TWEAK(MyTweak);

 #define REGISTER_TWEAK(Subclass)                                               \
   ::llvm::Registry<::clang::clangd::Tweak>::Add<Subclass>                      \
       TweakRegistrationFor##Subclass(#Subclass, /*Description=*/"");           \
   const char *Subclass::id() const { return #Subclass; }
 
 /// Calls prepare() on all tweaks that satisfy the filter, returning those that
 /// 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
 
 #endif

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

 //===--- Tweak.cpp -----------------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #include "Tweak.h"
+#include "FeatureModule.h"
 #include "SourceCode.h"
 #include "index/Index.h"
 #include "support/Logger.h"
 #include "support/Path.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/Registry.h"
 #include <functional>
 #include <memory>
 #include <utility>
+#include <vector>
 
 LLVM_INSTANTIATE_REGISTRY(llvm::Registry<clang::clangd::Tweak>)
 
 namespace clang {
 namespace clangd {
 
 /// A handy typedef to save some typing.
 typedef llvm::Registry<Tweak> TweakRegistry;
 
 namespace {
 /// Asserts invariants on TweakRegistry. No-op with assertion disabled.
 void validateRegistry() {
 #ifndef NDEBUG
   llvm::StringSet<> Seen;
   for (const auto &E : TweakRegistry::entries()) {
     // REGISTER_TWEAK ensures E.getName() is equal to the tweak class name.
     // We check that id() matches it.
     assert(E.instantiate()->id() == E.getName() &&
            "id should be equal to class name");
     assert(Seen.try_emplace(E.getName()).second && "duplicate check id");
   }
 #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,
                             unsigned RangeBegin, unsigned RangeEnd,
                             SelectionTree ASTSelection)
     : Index(Index), AST(&AST), SelectionBegin(RangeBegin),
       SelectionEnd(RangeEnd), ASTSelection(std::move(ASTSelection)) {
   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) {
+              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));
   }
   // Ensure deterministic order of the results.
   llvm::sort(Available,
              [](const std::unique_ptr<Tweak> &L,
                 const std::unique_ptr<Tweak> &R) { return L->id() < R->id(); });
   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,
                         tooling::Replacements Replacements) {
   Edit Ed(SM.getBufferData(FID), std::move(Replacements));
   if (auto FilePath = getCanonicalPath(SM.getFileEntryForID(FID), SM))
     return std::make_pair(*FilePath, std::move(Ed));
   return error("Failed to get absolute path for edited file: {0}",

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

for (const auto &Tok : SpelledTokens) {
         continue;
 
       // 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));
-      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());
           ++ErrCount;
         } else {
           vlog("    tweak: {0}", T->id());
         }
       }

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

add_unittest(ClangdUnitTests ClangdTests
   ConfigProviderTests.cpp
   ConfigYAMLTests.cpp
   DecisionForestTests.cpp
   DexTests.cpp
   DiagnosticsTests.cpp
   DraftStoreTests.cpp
   DumpASTTests.cpp
   ExpectedTypeTest.cpp
+  FeatureModulesTests.cpp
   FileDistanceTests.cpp
   FileIndexTests.cpp
   FindSymbolsTests.cpp
   FindTargetTests.cpp
   FormatTests.cpp
   FSTests.cpp
   FuzzyMatchTests.cpp
   GlobalCompilationDatabaseTests.cpp

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

+//===--- 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

 applyTweak(ParsedAST &AST, const Annotations &Input, StringRef TweakID,
            const SymbolIndex *Index) {
   auto Range = rangeOrPoint(Input);
   llvm::Optional<llvm::Expected<Tweak::Effect>> Result;
   SelectionTree::createEach(AST.getASTContext(), AST.getTokens(), Range.first,
                             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 {
                                 llvm::consumeError(T.takeError());
                                 return false;
                               }
                             });
   return Result;