[clangd] Modules can have a public API. NFC

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

clang-tools-extra/clangd/CMakeLists.txt

@@ -75,6 +75,7 @@ add_clang_library(clangDaemon
   Hover.cpp
   IncludeFixer.cpp
   JSONTransport.cpp
+  Module.cpp
   PathMapping.cpp
   Protocol.cpp
   Quality.cpp

clang-tools-extra/clangd/ClangdServer.cpp

@@ -127,7 +127,7 @@ ClangdServer::Options::operator TUScheduler::Options() const {
 ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
                            const ThreadsafeFS &TFS, const Options &Opts,
                            Callbacks *Callbacks)
-    : CDB(CDB), TFS(TFS),
+    : Modules(Opts.Modules), CDB(CDB), TFS(TFS),
       DynamicIdx(Opts.BuildDynamicSymbolIndex ? new FileIndex() : nullptr),
       ClangTidyProvider(Opts.ClangTidyProvider),
       WorkspaceRoot(Opts.WorkspaceRoot),

clang-tools-extra/clangd/ClangdServer.h

@@ -161,6 +161,15 @@ public:
   ClangdServer(const GlobalCompilationDatabase &CDB, const ThreadsafeFS &TFS,
                const Options &Opts, Callbacks *Callbacks = nullptr);
 
+  /// Gets the installed module of a given type, if any.
+  /// This exposes access the public interface of modules that have one.
+  template <typename Mod> Mod *getModule() {
+    return Modules ? Modules->get<Mod>() : nullptr;
+  }
+  template <typename Mod> const Mod *getModule() const {
+    return Modules ? Modules->get<Mod>() : nullptr;
+  }
+
   /// Add a \p File to the list of tracked C++ files or update the contents if
   /// \p File is already tracked. Also schedules parsing of the AST for it on a
   /// separate thread. When the parsing is complete, DiagConsumer passed in
@@ -337,6 +346,7 @@ private:
                   ArrayRef<tooling::Range> Ranges,
                   Callback<tooling::Replacements> CB);
 
+  ModuleSet *Modules;
   const GlobalCompilationDatabase &CDB;
   const ThreadsafeFS &TFS;
 

clang-tools-extra/clangd/Module.cpp

@@ -0,0 +1,18 @@
+#include "Module.h"
+
+namespace clang {
+namespace clangd {
+
+bool ModuleSet::addImpl(void *Key, std::unique_ptr<Module> M,
+                        const char *Source) {
+  if (!Map.try_emplace(Key, M.get()).second) {
+    // Source should (usually) include the name of the concrete module type.
+    elog("Tried to register duplicate modules via {0}", Source);
+    return false;
+  }
+  Modules.push_back(std::move(M));
+  return true;
+}
+
+} // namespace clangd
+} // namespace clang

clang-tools-extra/clangd/Module.h

@@ -3,8 +3,10 @@
 
 #include "LSPBinder.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Compiler.h"
 #include "llvm/Support/JSON.h"
 #include <memory>
+#include <type_traits>
 #include <vector>
 
 namespace clang {
@@ -41,12 +43,28 @@ public:
                              llvm::json::Object &ServerCaps) {}
 };
 
+/// A ModuleSet is a collection of modules installed in clangd.
+///
+/// Modules can be looked up by type, or used through the Module interface.
+/// This allows individual modules to expose a public API.
+/// For this reason, there can be only one module of each type.
+///
+/// ModuleSet owns the modules. It is itself owned by main, not ClangdServer.
 class ModuleSet {
   std::vector<std::unique_ptr<Module>> Modules;
+  llvm::DenseMap<void *, Module *> Map;
+
+  template <typename Mod> struct ID {
+    static_assert(std::is_base_of<Module, Mod>::value &&
+                      std::is_final<Mod>::value,
+                  "Modules must be final classes derived from clangd::Module");
+    static int Key;
+  };
+
+  bool addImpl(void *Key, std::unique_ptr<Module>, const char *Source);
 
 public:
-  explicit ModuleSet(std::vector<std::unique_ptr<Module>> Modules)
-      : Modules(std::move(Modules)) {}
+  ModuleSet() = default;
 
   using iterator = llvm::pointee_iterator<decltype(Modules)::iterator>;
   using const_iterator =
@@ -55,7 +73,20 @@ public:
   iterator end() { return iterator(Modules.end()); }
   const_iterator begin() const { return const_iterator(Modules.begin()); }
   const_iterator end() const { return const_iterator(Modules.end()); }
+
+  template <typename Mod> bool add(std::unique_ptr<Mod> M) {
+    return addImpl(&ID<Mod>::Key, std::move(M), LLVM_PRETTY_FUNCTION);
+  }
+  template <typename Mod> Mod *get() {
+    return static_cast<Mod *>(Map.lookup(&ID<Mod>::Key));
+  }
+  template <typename Mod> const Mod *get() const {
+    return const_cast<ModuleSet *>(this)->get<Mod>();
+  }
 };
+
+template <typename Mod> int ModuleSet::ID<Mod>::Key;
+
 } // namespace clangd
 } // namespace clang
 #endif

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

@@ -222,25 +222,26 @@ CompileFlags:
 }
 
 TEST_F(LSPTest, ModulesTest) {
-  class MathModule : public Module {
+  class MathModule final : public Module {
     void initializeLSP(LSPBinder &Bind, const llvm::json::Object &ClientCaps,
                        llvm::json::Object &ServerCaps) override {
       Bind.notification("add", this, &MathModule::add);
       Bind.method("get", this, &MathModule::get);
     }
 
+    int Value = 0;
+
+  public:
     void add(const int &X) { Value += X; }
     void get(const std::nullptr_t &, Callback<int> Reply) { Reply(Value); }
-    int Value = 0;
   };
-  std::vector<std::unique_ptr<Module>> Mods;
-  Mods.push_back(std::make_unique<MathModule>());
-  ModuleSet ModSet(std::move(Mods));
-  Opts.Modules = &ModSet;
+  ModuleSet Mods;
+  Mods.add(std::make_unique<MathModule>());
+  Opts.Modules = &Mods;
 
   auto &Client = start();
   Client.notify("add", 2);
-  Client.notify("add", 8);
+  Mods.get<MathModule>()->add(8);
   EXPECT_EQ(10, Client.call("get", nullptr).takeValue());
 }