[mlir][Tablegen-LSP] Add support for a basic TableGen language server

This follows the same general structure of the MLIR and PDLL language
servers. This commits adds the basic functionality for setting up the server,
and initially only supports providing diagnostics. Followon commits will
build out more comprehensive behavior.

Realistically this should eventually live in llvm/, but building in MLIR is an easier
initial step given that:
* All of the necessary LSP functionality is already here
* It allows for proving out useful language features (e.g. compilation databases)
  without affecting wider scale tablegen users
* MLIR has a vscode extension that can immediately take advantage of it

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

Author: River707

mlir/docs/Tools/MLIRLSP.md

@@ -176,6 +176,8 @@ extra setup steps are required:
 
 *   Copy `mlir/utils/textmate/mlir.json` to the extension directory and rename
     to `grammar.json`.
+*   Copy `llvm/utils/textmate/tablegen.json` to the extension directory and rename
+    to `tablegen-grammar.json`.
 *   Copy
     `https://mlir.llvm.org//LogoAssets/logo/PNG/full_color/mlir-identity-03.png`
     to the extension directory and rename to `icon.png`.

mlir/include/mlir/Tools/tblgen-lsp-server/TableGenLspServerMain.h

@@ -0,0 +1,24 @@
+//===- TableGenLSPServerMain.h - TableGen Language Server main --*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Main entry function for tblgen-lsp-server when built as standalone binary.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TOOLS_TBLGENLSPSERVER_TABLEGENLSPSERVERMAIN_H
+#define MLIR_TOOLS_TBLGENLSPSERVER_TABLEGENLSPSERVERMAIN_H
+
+namespace mlir {
+struct LogicalResult;
+
+/// Implementation for tools like `tblgen-lsp-server`.
+LogicalResult TableGenLspServerMain(int argc, char **argv);
+
+} // namespace mlir
+
+#endif // MLIR_TOOLS_TBLGENLSPSERVER_TABLEGENLSPSERVERMAIN_H

mlir/lib/Tools/CMakeLists.txt

@@ -5,3 +5,4 @@ add_subdirectory(mlir-pdll-lsp-server)
 add_subdirectory(mlir-reduce)
 add_subdirectory(mlir-translate)
 add_subdirectory(PDLL)
+add_subdirectory(tblgen-lsp-server)

mlir/lib/Tools/lsp-server-support/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_mlir_library(MLIRLspServerSupportLib
   Logging.cpp
   Protocol.cpp
+  SourceMgrUtils.cpp
   Transport.cpp
 
   ADDITIONAL_HEADER_DIRS

mlir/lib/Tools/lsp-server-support/SourceMgrUtils.cpp

@@ -0,0 +1,61 @@
+//===--- SourceMgrUtils.cpp - SourceMgr LSP Utils -------------------------===//
+//
+// 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 "SourceMgrUtils.h"
+
+using namespace mlir;
+using namespace mlir::lsp;
+
+/// Find the end of a string whose contents start at the given `curPtr`. Returns
+/// the position at the end of the string, after a terminal or invalid character
+/// (e.g. `"` or `\0`).
+static const char *lexLocStringTok(const char *curPtr) {
+  while (char c = *curPtr++) {
+    // Check for various terminal characters.
+    if (StringRef("\"\n\v\f").contains(c))
+      return curPtr;
+
+    // Check for escape sequences.
+    if (c == '\\') {
+      // Check a few known escapes and \xx hex digits.
+      if (*curPtr == '"' || *curPtr == '\\' || *curPtr == 'n' || *curPtr == 't')
+        ++curPtr;
+      else if (llvm::isHexDigit(*curPtr) && llvm::isHexDigit(curPtr[1]))
+        curPtr += 2;
+      else
+        return curPtr;
+    }
+  }
+
+  // If we hit this point, we've reached the end of the buffer. Update the end
+  // pointer to not point past the buffer.
+  return curPtr - 1;
+}
+
+SMRange lsp::convertTokenLocToRange(SMLoc loc) {
+  if (!loc.isValid())
+    return SMRange();
+  const char *curPtr = loc.getPointer();
+
+  // Check if this is a string token.
+  if (*curPtr == '"') {
+    curPtr = lexLocStringTok(curPtr + 1);
+
+    // Otherwise, default to handling an identifier.
+  } else {
+    // Return if the given character is a valid identifier character.
+    auto isIdentifierChar = [](char c) {
+      return isalnum(c) || c == '$' || c == '.' || c == '_' || c == '-';
+    };
+
+    while (*curPtr && isIdentifierChar(*(++curPtr)))
+      continue;
+  }
+
+  return SMRange(loc, SMLoc::getFromPointer(curPtr));
+}

mlir/lib/Tools/lsp-server-support/SourceMgrUtils.h

@@ -0,0 +1,31 @@
+//===--- SourceMgrUtils.h - SourceMgr LSP Utils -----------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains an array of generally useful SourceMgr utilities for
+// interacting with LSP components.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LIB_MLIR_TOOLS_LSPSERVERSUPPORT_TRANSPORT_H_
+#define LIB_MLIR_TOOLS_LSPSERVERSUPPORT_TRANSPORT_H_
+
+#include "Protocol.h"
+#include "llvm/Support/SourceMgr.h"
+
+namespace mlir {
+namespace lsp {
+
+/// Returns the range of a lexical token given a SMLoc corresponding to the
+/// start of an token location. The range is computed heuristically, and
+/// supports identifier-like tokens, strings, etc.
+SMRange convertTokenLocToRange(SMLoc loc);
+
+} // namespace lsp
+} // namespace mlir
+
+#endif

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

@@ -9,6 +9,7 @@
 #include "MLIRServer.h"
 #include "../lsp-server-support/Logging.h"
 #include "../lsp-server-support/Protocol.h"
+#include "../lsp-server-support/SourceMgrUtils.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Parser/AsmParserState.h"
@@ -55,7 +56,7 @@ getLocationFromLoc(llvm::SourceMgr &sourceMgr, Location loc,
       // Use range of potential identifier starting at location, else length 1
       // range.
       location->range.end.character += 1;
-      if (Optional<SMRange> range = AsmParserState::convertIdLocToRange(loc)) {
+      if (Optional<SMRange> range = lsp::convertTokenLocToRange(loc)) {
         auto lineCol = sourceMgr.getLineAndColumn(range->End);
         location->range.end.character =
             std::max(fileLoc.getColumn() + 1, lineCol.second - 1);

mlir/lib/Tools/tblgen-lsp-server/CMakeLists.txt

@@ -0,0 +1,18 @@
+set(LLVM_LINK_COMPONENTS
+  Demangle
+  Support
+  TableGen
+)
+
+llvm_add_library(TableGenLspServerLib
+  LSPServer.cpp
+  TableGenServer.cpp
+  TableGenLspServerMain.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Tools/tblgen-lsp-server
+
+  LINK_LIBS PUBLIC
+  MLIRLspServerSupportLib
+  MLIRSupport
+  )

mlir/lib/Tools/tblgen-lsp-server/LSPServer.cpp

@@ -0,0 +1,163 @@
+//===- LSPServer.cpp - TableGen Language Server ---------------------------===//
+//
+// 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 "LSPServer.h"
+
+#include "../lsp-server-support/Logging.h"
+#include "../lsp-server-support/Protocol.h"
+#include "../lsp-server-support/Transport.h"
+#include "TableGenServer.h"
+#include "llvm/ADT/FunctionExtras.h"
+#include "llvm/ADT/StringMap.h"
+
+using namespace mlir;
+using namespace mlir::lsp;
+
+//===----------------------------------------------------------------------===//
+// LSPServer
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct LSPServer {
+  LSPServer(TableGenServer &server, JSONTransport &transport)
+      : server(server), transport(transport) {}
+
+  //===--------------------------------------------------------------------===//
+  // Initialization
+
+  void onInitialize(const InitializeParams &params,
+                    Callback<llvm::json::Value> reply);
+  void onInitialized(const InitializedParams &params);
+  void onShutdown(const NoParams &params, Callback<std::nullptr_t> reply);
+
+  //===--------------------------------------------------------------------===//
+  // Document Change
+
+  void onDocumentDidOpen(const DidOpenTextDocumentParams &params);
+  void onDocumentDidClose(const DidCloseTextDocumentParams &params);
+  void onDocumentDidChange(const DidChangeTextDocumentParams &params);
+
+  //===--------------------------------------------------------------------===//
+  // Fields
+  //===--------------------------------------------------------------------===//
+
+  TableGenServer &server;
+  JSONTransport &transport;
+
+  /// An outgoing notification used to send diagnostics to the client when they
+  /// are ready to be processed.
+  OutgoingNotification<PublishDiagnosticsParams> publishDiagnostics;
+
+  /// Used to indicate that the 'shutdown' request was received from the
+  /// Language Server client.
+  bool shutdownRequestReceived = false;
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Initialization
+
+void LSPServer::onInitialize(const InitializeParams &params,
+                             Callback<llvm::json::Value> reply) {
+  // Send a response with the capabilities of this server.
+  llvm::json::Object serverCaps{
+      {"textDocumentSync",
+       llvm::json::Object{
+           {"openClose", true},
+           {"change", (int)TextDocumentSyncKind::Full},
+           {"save", true},
+       }},
+  };
+
+  llvm::json::Object result{
+      {{"serverInfo", llvm::json::Object{{"name", "tblgen-lsp-server"},
+                                         {"version", "0.0.1"}}},
+       {"capabilities", std::move(serverCaps)}}};
+  reply(std::move(result));
+}
+void LSPServer::onInitialized(const InitializedParams &) {}
+void LSPServer::onShutdown(const NoParams &, Callback<std::nullptr_t> reply) {
+  shutdownRequestReceived = true;
+  reply(nullptr);
+}
+
+//===----------------------------------------------------------------------===//
+// Document Change
+
+void LSPServer::onDocumentDidOpen(const DidOpenTextDocumentParams &params) {
+  PublishDiagnosticsParams diagParams(params.textDocument.uri,
+                                      params.textDocument.version);
+  server.addOrUpdateDocument(params.textDocument.uri, params.textDocument.text,
+                             params.textDocument.version,
+                             diagParams.diagnostics);
+
+  // Publish any recorded diagnostics.
+  publishDiagnostics(diagParams);
+}
+void LSPServer::onDocumentDidClose(const DidCloseTextDocumentParams &params) {
+  Optional<int64_t> version = server.removeDocument(params.textDocument.uri);
+  if (!version)
+    return;
+
+  // Empty out the diagnostics shown for this document. This will clear out
+  // anything currently displayed by the client for this document (e.g. in the
+  // "Problems" pane of VSCode).
+  publishDiagnostics(
+      PublishDiagnosticsParams(params.textDocument.uri, *version));
+}
+void LSPServer::onDocumentDidChange(const DidChangeTextDocumentParams &params) {
+  // TODO: We currently only support full document updates, we should refactor
+  // to avoid this.
+  if (params.contentChanges.size() != 1)
+    return;
+  PublishDiagnosticsParams diagParams(params.textDocument.uri,
+                                      params.textDocument.version);
+  server.addOrUpdateDocument(
+      params.textDocument.uri, params.contentChanges.front().text,
+      params.textDocument.version, diagParams.diagnostics);
+
+  // Publish any recorded diagnostics.
+  publishDiagnostics(diagParams);
+}
+
+//===----------------------------------------------------------------------===//
+// Entry Point
+//===----------------------------------------------------------------------===//
+
+LogicalResult mlir::lsp::runTableGenLSPServer(TableGenServer &server,
+                                              JSONTransport &transport) {
+  LSPServer lspServer(server, transport);
+  MessageHandler messageHandler(transport);
+
+  // Initialization
+  messageHandler.method("initialize", &lspServer, &LSPServer::onInitialize);
+  messageHandler.notification("initialized", &lspServer,
+                              &LSPServer::onInitialized);
+  messageHandler.method("shutdown", &lspServer, &LSPServer::onShutdown);
+
+  // Document Changes
+  messageHandler.notification("textDocument/didOpen", &lspServer,
+                              &LSPServer::onDocumentDidOpen);
+  messageHandler.notification("textDocument/didClose", &lspServer,
+                              &LSPServer::onDocumentDidClose);
+  messageHandler.notification("textDocument/didChange", &lspServer,
+                              &LSPServer::onDocumentDidChange);
+
+  // Diagnostics
+  lspServer.publishDiagnostics =
+      messageHandler.outgoingNotification<PublishDiagnosticsParams>(
+          "textDocument/publishDiagnostics");
+
+  // Run the main loop of the transport.
+  if (llvm::Error error = transport.run(messageHandler)) {
+    Logger::error("Transport error: {0}", error);
+    llvm::consumeError(std::move(error));
+    return failure();
+  }
+  return success(lspServer.shutdownRequestReceived);
+}

mlir/lib/Tools/tblgen-lsp-server/LSPServer.h

@@ -0,0 +1,29 @@
+//===- LSPServer.h - TableGen LSP Server ------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LIB_MLIR_TOOLS_TBLGENLSPSERVER_LSPSERVER_H
+#define LIB_MLIR_TOOLS_TBLGENLSPSERVER_LSPSERVER_H
+
+#include <memory>
+
+namespace mlir {
+struct LogicalResult;
+
+namespace lsp {
+class JSONTransport;
+class TableGenServer;
+
+/// Run the main loop of the LSP server using the given TableGen server and
+/// transport.
+LogicalResult runTableGenLSPServer(TableGenServer &server,
+                                   JSONTransport &transport);
+
+} // namespace lsp
+} // namespace mlir
+
+#endif // LIB_MLIR_TOOLS_TBLGENLSPSERVER_LSPSERVER_H