[mlir][lsp] Enable registering dialects based on URI.

[jpienaar]

Previously the dialects registered were fixed per LSP binary. This works as long as all the dialects of interest from the different projects across which one uses the LSP, are disjoint. This expands this to support cases where there are dialects that overlap in dialect name but usage of these are separate wrt projects. The alternative is multiple binaries and switching LSP used in editor per project (there is some extra complexity in hosted instances).

This handles a simple (I believe common case) where one can determine based on path and have single binary - the cost of dynamically doing so based on path would be either keeping different registries to return or repopulating dialect & extension maps.

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Jacques Pienaar (jpienaar)

Changes

Previously the dialects registered were fixed per LSP binary. This works as long as all the dialects of interest from the different projects across which one uses the LSP, are disjoint. This expands this to support cases where there are dialects that overlap in dialect name but usage of these are separate wrt projects. The alternative is multiple binaries and switching LSP used in editor per project (there is some extra complexity in hosted instances).

This handles a simple (I believe common case) where one can determine based on path and have single binary - the cost of dynamically doing so based on path would be either keeping different registries to return or repopulating dialect & extension maps.

---

Full diff: https://github.com/llvm/llvm-project/pull/141331.diff

5 Files Affected:

• (added) mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h (+29)
• (modified) mlir/include/mlir/Tools/mlir-lsp-server/MlirLspServerMain.h (+9-2)
• (modified) mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp (+11-11)
• (modified) mlir/lib/Tools/mlir-lsp-server/MLIRServer.h (+3-3)
• (modified) mlir/lib/Tools/mlir-lsp-server/MlirLspServerMain.cpp (+11-2)

diff --git a/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h b/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h
new file mode 100644
index 0000000000000..c83b361a0e51c
--- /dev/null
+++ b/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h
@@ -0,0 +1,29 @@
+//===- MlirLspRegistryFunction.h - LSP registry functions -------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Registry function types for MLIR LSP.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPREGISTRYFUNCTION_H
+#define MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPREGISTRYFUNCTION_H
+
+namespace llvm {
+template <typename Fn> class function_ref;
+} // namespace llvm
+
+namespace mlir {
+class DialectRegistry;
+namespace lsp {
+class URIForFile;
+using DialectRegistryFn =
+    llvm::function_ref<DialectRegistry &(const URIForFile &uri)>;
+} // namespace lsp
+} // namespace mlir
+
+#endif // MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPREGISTRYFUNCTION_H
diff --git a/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspServerMain.h b/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspServerMain.h
index 66d5d40a6d28d..a461fc4702946 100644
--- a/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspServerMain.h
+++ b/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspServerMain.h
@@ -12,20 +12,27 @@
 
 #ifndef MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPSERVERMAIN_H
 #define MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPSERVERMAIN_H
+#include "mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h"
 
 namespace llvm {
 struct LogicalResult;
 } // namespace llvm
 
 namespace mlir {
-class DialectRegistry;
 
 /// Implementation for tools like `mlir-lsp-server`.
 /// - registry should contain all the dialects that can be parsed in source IR
-/// passed to the server.
+///   passed to the server.
 llvm::LogicalResult MlirLspServerMain(int argc, char **argv,
                                       DialectRegistry &registry);
 
+/// Implementation for tools like `mlir-lsp-server`.
+/// - registry should contain all the dialects that can be parsed in source IR
+///   passed to the server and may register different dialects depending on the
+///   input URI.
+llvm::LogicalResult MlirLspServerMain(int argc, char **argv,
+                                      lsp::DialectRegistryFn registry_fn);
+
 } // namespace mlir
 
 #endif // MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPSERVERMAIN_H
diff --git a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
index 4e19274c3da40..61987525a5ca5 100644
--- a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
+++ b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
@@ -997,7 +997,7 @@ namespace {
 class MLIRTextFile {
 public:
   MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents,
-               int64_t version, DialectRegistry &registry,
+               int64_t version, lsp::DialectRegistryFn registry_fn,
                std::vector<lsp::Diagnostic> &diagnostics);
 
   /// Return the current version of this text file.
@@ -1046,9 +1046,9 @@ class MLIRTextFile {
 } // namespace
 
 MLIRTextFile::MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents,
-                           int64_t version, DialectRegistry &registry,
+                           int64_t version, lsp::DialectRegistryFn registry_fn,
                            std::vector<lsp::Diagnostic> &diagnostics)
-    : context(registry, MLIRContext::Threading::DISABLED),
+    : context(registry_fn(uri), MLIRContext::Threading::DISABLED),
       contents(fileContents.str()), version(version) {
   context.allowUnregisteredDialects();
 
@@ -1263,11 +1263,11 @@ MLIRTextFileChunk &MLIRTextFile::getChunkFor(lsp::Position &pos) {
 //===----------------------------------------------------------------------===//
 
 struct lsp::MLIRServer::Impl {
-  Impl(DialectRegistry &registry) : registry(registry) {}
+  Impl(lsp::DialectRegistryFn registry_fn) : registry_fn(registry_fn) {}
 
-  /// The registry containing dialects that can be recognized in parsed .mlir
-  /// files.
-  DialectRegistry &registry;
+  /// The registry factory for containing dialects that can be recognized in
+  /// parsed .mlir files.
+  lsp::DialectRegistryFn registry_fn;
 
   /// The files held by the server, mapped by their URI file name.
   llvm::StringMap<std::unique_ptr<MLIRTextFile>> files;
@@ -1277,15 +1277,15 @@ struct lsp::MLIRServer::Impl {
 // MLIRServer
 //===----------------------------------------------------------------------===//
 
-lsp::MLIRServer::MLIRServer(DialectRegistry &registry)
-    : impl(std::make_unique<Impl>(registry)) {}
+lsp::MLIRServer::MLIRServer(lsp::DialectRegistryFn registry_fn)
+    : impl(std::make_unique<Impl>(registry_fn)) {}
 lsp::MLIRServer::~MLIRServer() = default;
 
 void lsp::MLIRServer::addOrUpdateDocument(
     const URIForFile &uri, StringRef contents, int64_t version,
     std::vector<Diagnostic> &diagnostics) {
   impl->files[uri.file()] = std::make_unique<MLIRTextFile>(
-      uri, contents, version, impl->registry, diagnostics);
+      uri, contents, version, impl->registry_fn, diagnostics);
 }
 
 std::optional<int64_t> lsp::MLIRServer::removeDocument(const URIForFile &uri) {
@@ -1348,7 +1348,7 @@ void lsp::MLIRServer::getCodeActions(const URIForFile &uri, const Range &pos,
 
 llvm::Expected<lsp::MLIRConvertBytecodeResult>
 lsp::MLIRServer::convertFromBytecode(const URIForFile &uri) {
-  MLIRContext tempContext(impl->registry);
+  MLIRContext tempContext(impl->registry_fn(uri));
   tempContext.allowUnregisteredDialects();
 
   // Collect any errors during parsing.
diff --git a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.h b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.h
index 979be615b82cc..85e69e69f6631 100644
--- a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.h
+++ b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.h
@@ -10,6 +10,7 @@
 #define LIB_MLIR_TOOLS_MLIRLSPSERVER_SERVER_H_
 
 #include "mlir/Support/LLVM.h"
+#include "mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h"
 #include "llvm/Support/Error.h"
 #include <memory>
 #include <optional>
@@ -28,15 +29,14 @@ struct Location;
 struct MLIRConvertBytecodeResult;
 struct Position;
 struct Range;
-class URIForFile;
 
 /// This class implements all of the MLIR related functionality necessary for a
 /// language server. This class allows for keeping the MLIR specific logic
 /// separate from the logic that involves LSP server/client communication.
 class MLIRServer {
 public:
-  /// Construct a new server with the given dialect regitstry.
-  MLIRServer(DialectRegistry &registry);
+  /// Construct a new server with the given dialect registry function.
+  MLIRServer(DialectRegistryFn registry_fn);
   ~MLIRServer();
 
   /// Add or update the document, with the provided `version`, at the given URI.
diff --git a/mlir/lib/Tools/mlir-lsp-server/MlirLspServerMain.cpp b/mlir/lib/Tools/mlir-lsp-server/MlirLspServerMain.cpp
index 259bd2613a6cc..b1bbf98ce769e 100644
--- a/mlir/lib/Tools/mlir-lsp-server/MlirLspServerMain.cpp
+++ b/mlir/lib/Tools/mlir-lsp-server/MlirLspServerMain.cpp
@@ -19,7 +19,7 @@ using namespace mlir;
 using namespace mlir::lsp;
 
 LogicalResult mlir::MlirLspServerMain(int argc, char **argv,
-                                      DialectRegistry &registry) {
+                                      DialectRegistryFn registry_fn) {
   llvm::cl::opt<JSONStreamStyle> inputStyle{
       "input-style",
       llvm::cl::desc("Input JSON stream encoding"),
@@ -72,6 +72,15 @@ LogicalResult mlir::MlirLspServerMain(int argc, char **argv,
   URIForFile::registerSupportedScheme("mlir.bytecode-mlir");
 
   // Configure the servers and start the main language server.
-  MLIRServer server(registry);
+  MLIRServer server(registry_fn);
   return runMlirLSPServer(server, transport);
 }
+
+llvm::LogicalResult mlir::MlirLspServerMain(int argc, char **argv,
+                                            DialectRegistry &registry) {
+  auto registry_fn =
+      [&registry](const lsp::URIForFile &uri) -> DialectRegistry & {
+    return registry;
+  };
+  return MlirLspServerMain(argc, argv, registry_fn);
+}

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[joker-eph]

I assume it's not easy to write tests for this?

[jpienaar]

I assume it's not easy to write tests for this?

Not too difficult, but would require either introducing a new lsp binary or adding flag/convention to existing one.

E.g., I could make one where no dialects are registered if the file contains _disable_lsp_registration in its name and MLIR_INCLUDE_TESTS is set, and then just validating that parsing of the file fails. That means the regular mlir-lsp-server binary (which is semi combination of testing tool and example TBF - its very small making own LSP and probably everybody does to get their own dialects) only gets this while testing and so should have rather small overhead only when building with testing added, no extra binary and verifying this works. WDYT?