[mlir][Tablegen-LSP] Add support for include file link and hover

This allows for following links to include files. This support is effectively
identical to the logic in the PDLL language server, and code is shared as
much as possible.

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

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

@@ -7,10 +7,15 @@
 //===----------------------------------------------------------------------===//
 
 #include "SourceMgrUtils.h"
+#include "llvm/Support/Path.h"
 
 using namespace mlir;
 using namespace mlir::lsp;
 
+//===----------------------------------------------------------------------===//
+// Utils
+//===----------------------------------------------------------------------===//
+
 /// 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`).
@@ -59,3 +64,46 @@ SMRange lsp::convertTokenLocToRange(SMLoc loc) {
 
   return SMRange(loc, SMLoc::getFromPointer(curPtr));
 }
+
+//===----------------------------------------------------------------------===//
+// SourceMgrInclude
+//===----------------------------------------------------------------------===//
+
+Hover SourceMgrInclude::buildHover() const {
+  Hover hover(range);
+  {
+    llvm::raw_string_ostream hoverOS(hover.contents.value);
+    hoverOS << "`" << llvm::sys::path::filename(uri.file()) << "`\n***\n"
+            << uri.file();
+  }
+  return hover;
+}
+
+void lsp::gatherIncludeFiles(llvm::SourceMgr &sourceMgr,
+                             SmallVectorImpl<SourceMgrInclude> &includes) {
+  for (unsigned i = 1, e = sourceMgr.getNumBuffers(); i < e; ++i) {
+    // Check to see if this file was included by the main file.
+    SMLoc includeLoc = sourceMgr.getBufferInfo(i + 1).IncludeLoc;
+    if (!includeLoc.isValid() || sourceMgr.FindBufferContainingLoc(
+                                     includeLoc) != sourceMgr.getMainFileID())
+      continue;
+
+    // Try to build a URI for this file path.
+    auto *buffer = sourceMgr.getMemoryBuffer(i + 1);
+    llvm::SmallString<256> path(buffer->getBufferIdentifier());
+    llvm::sys::path::remove_dots(path, /*remove_dot_dot=*/true);
+
+    llvm::Expected<URIForFile> includedFileURI = URIForFile::fromFile(path);
+    if (!includedFileURI)
+      continue;
+
+    // Find the end of the include token.
+    const char *includeStart = includeLoc.getPointer() - 2;
+    while (*(--includeStart) != '\"')
+      continue;
+
+    // Push this include.
+    SMRange includeRange(SMLoc::getFromPointer(includeStart), includeLoc);
+    includes.emplace_back(*includedFileURI, Range(sourceMgr, includeRange));
+  }
+}

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

@@ -19,12 +19,39 @@
 
 namespace mlir {
 namespace lsp {
+//===----------------------------------------------------------------------===//
+// Utils
+//===----------------------------------------------------------------------===//
 
 /// 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);
 
+//===----------------------------------------------------------------------===//
+// SourceMgrInclude
+//===----------------------------------------------------------------------===//
+
+/// This class represents a single include within a root file.
+struct SourceMgrInclude {
+  SourceMgrInclude(const lsp::URIForFile &uri, const lsp::Range &range)
+      : uri(uri), range(range) {}
+
+  /// Build a hover for the current include file.
+  Hover buildHover() const;
+
+  /// The URI of the file that is included.
+  lsp::URIForFile uri;
+
+  /// The range of the include directive.
+  lsp::Range range;
+};
+
+/// Given a source manager, gather all of the processed include files. These are
+/// assumed to be all of the files other than the main root file.
+void gatherIncludeFiles(llvm::SourceMgr &sourceMgr,
+                        SmallVectorImpl<SourceMgrInclude> &includes);
+
 } // namespace lsp
 } // namespace mlir
 

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp

@@ -11,6 +11,7 @@
 #include "../lsp-server-support/CompilationDatabase.h"
 #include "../lsp-server-support/Logging.h"
 #include "../lsp-server-support/Protocol.h"
+#include "../lsp-server-support/SourceMgrUtils.h"
 #include "mlir/Tools/PDLL/AST/Context.h"
 #include "mlir/Tools/PDLL/AST/Nodes.h"
 #include "mlir/Tools/PDLL/AST/Types.h"
@@ -106,24 +107,6 @@ getLspDiagnoticFromDiag(llvm::SourceMgr &sourceMgr, const ast::Diagnostic &diag,
   return lspDiag;
 }
 
-//===----------------------------------------------------------------------===//
-// PDLLInclude
-//===----------------------------------------------------------------------===//
-
-namespace {
-/// This class represents a single include within a root file.
-struct PDLLInclude {
-  PDLLInclude(const lsp::URIForFile &uri, const lsp::Range &range)
-      : uri(uri), range(range) {}
-
-  /// The URI of the file that is included.
-  lsp::URIForFile uri;
-
-  /// The range of the include directive.
-  lsp::Range range;
-};
-} // namespace
-
 //===----------------------------------------------------------------------===//
 // PDLIndex
 //===----------------------------------------------------------------------===//
@@ -288,7 +271,6 @@ struct PDLDocument {
                                  const lsp::Position &hoverPos);
   Optional<lsp::Hover> findHover(const ast::Decl *decl,
                                  const SMRange &hoverRange);
-  lsp::Hover buildHoverForInclude(const PDLLInclude &include);
   lsp::Hover buildHoverForOpName(const ods::Operation *op,
                                  const SMRange &hoverRange);
   lsp::Hover buildHoverForVariable(const ast::VariableDecl *varDecl,
@@ -343,7 +325,7 @@ struct PDLDocument {
   PDLIndex index;
 
   /// The set of includes of the parsed module.
-  std::vector<PDLLInclude> parsedIncludes;
+  SmallVector<lsp::SourceMgrInclude> parsedIncludes;
 };
 } // namespace
 
@@ -373,33 +355,7 @@ PDLDocument::PDLDocument(const lsp::URIForFile &uri, StringRef contents,
   astModule = parsePDLAST(astContext, sourceMgr);
 
   // Initialize the set of parsed includes.
-  for (unsigned i = 1, e = sourceMgr.getNumBuffers(); i < e; ++i) {
-    // Check to see if this file was included by the main file.
-    SMLoc includeLoc = sourceMgr.getBufferInfo(i + 1).IncludeLoc;
-    if (!includeLoc.isValid() || sourceMgr.FindBufferContainingLoc(
-                                     includeLoc) != sourceMgr.getMainFileID())
-      continue;
-
-    // Try to build a URI for this file path.
-    auto *buffer = sourceMgr.getMemoryBuffer(i + 1);
-    llvm::SmallString<256> path(buffer->getBufferIdentifier());
-    llvm::sys::path::remove_dots(path, /*remove_dot_dot=*/true);
-
-    llvm::Expected<lsp::URIForFile> includedFileURI =
-        lsp::URIForFile::fromFile(path);
-    if (!includedFileURI)
-      continue;
-
-    // Find the end of the include token.
-    const char *includeStart = includeLoc.getPointer() - 2;
-    while (*(--includeStart) != '\"')
-      continue;
-
-    // Push this include.
-    SMRange includeRange(SMLoc::getFromPointer(includeStart), includeLoc);
-    parsedIncludes.emplace_back(*includedFileURI,
-                                lsp::Range(sourceMgr, includeRange));
-  }
+  lsp::gatherIncludeFiles(sourceMgr, parsedIncludes);
 
   // If we failed to parse the module, there is nothing left to initialize.
   if (failed(astModule))
@@ -443,7 +399,7 @@ void PDLDocument::findReferencesOf(const lsp::URIForFile &uri,
 
 void PDLDocument::getDocumentLinks(const lsp::URIForFile &uri,
                                    std::vector<lsp::DocumentLink> &links) {
-  for (const PDLLInclude &include : parsedIncludes)
+  for (const lsp::SourceMgrInclude &include : parsedIncludes)
     links.emplace_back(include.range, include.uri);
 }
 
@@ -456,10 +412,9 @@ Optional<lsp::Hover> PDLDocument::findHover(const lsp::URIForFile &uri,
   SMLoc posLoc = hoverPos.getAsSMLoc(sourceMgr);
 
   // Check for a reference to an include.
-  for (const PDLLInclude &include : parsedIncludes) {
+  for (const lsp::SourceMgrInclude &include : parsedIncludes)
     if (include.range.contains(hoverPos))
-      return buildHoverForInclude(include);
-  }
+      return include.buildHover();
 
   // Find the symbol at the given location.
   SMRange hoverRange;
@@ -499,17 +454,6 @@ Optional<lsp::Hover> PDLDocument::findHover(const ast::Decl *decl,
   return llvm::None;
 }
 
-lsp::Hover PDLDocument::buildHoverForInclude(const PDLLInclude &include) {
-  lsp::Hover hover(include.range);
-  {
-    llvm::raw_string_ostream hoverOS(hover.contents.value);
-    hoverOS << "`" << llvm::sys::path::filename(include.uri.file())
-            << "`\n***\n"
-            << include.uri.file();
-  }
-  return hover;
-}
-
 lsp::Hover PDLDocument::buildHoverForOpName(const ods::Operation *op,
                                             const SMRange &hoverRange) {
   lsp::Hover hover(lsp::Range(sourceMgr, hoverRange));

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

@@ -42,6 +42,18 @@ struct LSPServer {
   void onDocumentDidClose(const DidCloseTextDocumentParams &params);
   void onDocumentDidChange(const DidChangeTextDocumentParams &params);
 
+  //===----------------------------------------------------------------------===//
+  // DocumentLink
+
+  void onDocumentLink(const DocumentLinkParams &params,
+                      Callback<std::vector<DocumentLink>> reply);
+
+  //===--------------------------------------------------------------------===//
+  // Hover
+
+  void onHover(const TextDocumentPositionParams &params,
+               Callback<Optional<Hover>> reply);
+
   //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
@@ -72,6 +84,11 @@ void LSPServer::onInitialize(const InitializeParams &params,
            {"change", (int)TextDocumentSyncKind::Full},
            {"save", true},
        }},
+      {"documentLinkProvider",
+       llvm::json::Object{
+           {"resolveProvider", false},
+       }},
+      {"hoverProvider", true},
   };
 
   llvm::json::Object result{
@@ -125,6 +142,24 @@ void LSPServer::onDocumentDidChange(const DidChangeTextDocumentParams &params) {
   publishDiagnostics(diagParams);
 }
 
+//===----------------------------------------------------------------------===//
+// DocumentLink
+
+void LSPServer::onDocumentLink(const DocumentLinkParams &params,
+                               Callback<std::vector<DocumentLink>> reply) {
+  std::vector<DocumentLink> links;
+  server.getDocumentLinks(params.textDocument.uri, links);
+  reply(std::move(links));
+}
+
+//===----------------------------------------------------------------------===//
+// Hover
+
+void LSPServer::onHover(const TextDocumentPositionParams &params,
+                        Callback<Optional<Hover>> reply) {
+  reply(server.findHover(params.textDocument.uri, params.position));
+}
+
 //===----------------------------------------------------------------------===//
 // Entry Point
 //===----------------------------------------------------------------------===//
@@ -148,6 +183,13 @@ LogicalResult mlir::lsp::runTableGenLSPServer(TableGenServer &server,
   messageHandler.notification("textDocument/didChange", &lspServer,
                               &LSPServer::onDocumentDidChange);
 
+  // Document Link
+  messageHandler.method("textDocument/documentLink", &lspServer,
+                        &LSPServer::onDocumentLink);
+
+  // Hover
+  messageHandler.method("textDocument/hover", &lspServer, &LSPServer::onHover);
+
   // Diagnostics
   lspServer.publishDiagnostics =
       messageHandler.outgoingNotification<PublishDiagnosticsParams>(

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

@@ -103,6 +103,20 @@ public:
   /// Return the current version of this text file.
   int64_t getVersion() const { return version; }
 
+  //===--------------------------------------------------------------------===//
+  // Document Links
+  //===--------------------------------------------------------------------===//
+
+  void getDocumentLinks(const lsp::URIForFile &uri,
+                        std::vector<lsp::DocumentLink> &links);
+
+  //===--------------------------------------------------------------------===//
+  // Hover
+  //===--------------------------------------------------------------------===//
+
+  Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
+                                 const lsp::Position &hoverPos);
+
 private:
   /// The full string contents of the file.
   std::string contents;
@@ -118,6 +132,9 @@ private:
 
   /// The record keeper containing the parsed tablegen constructs.
   llvm::RecordKeeper recordKeeper;
+
+  /// The set of includes of the parsed file.
+  SmallVector<lsp::SourceMgrInclude> parsedIncludes;
 };
 } // namespace
 
@@ -157,10 +174,38 @@ TableGenTextFile::TableGenTextFile(
           ctx->diagnostics.push_back(*lspDiag);
       },
       &handlerContext);
-  if (llvm::TableGenParseFile(sourceMgr, recordKeeper))
+  bool failedToParse = llvm::TableGenParseFile(sourceMgr, recordKeeper);
+
+  // Process all of the include files.
+  lsp::gatherIncludeFiles(sourceMgr, parsedIncludes);
+  if (failedToParse)
     return;
 }
 
+//===--------------------------------------------------------------------===//
+// TableGenTextFile: Document Links
+//===--------------------------------------------------------------------===//
+
+void TableGenTextFile::getDocumentLinks(const lsp::URIForFile &uri,
+                                        std::vector<lsp::DocumentLink> &links) {
+  for (const lsp::SourceMgrInclude &include : parsedIncludes)
+    links.emplace_back(include.range, include.uri);
+}
+
+//===----------------------------------------------------------------------===//
+// TableGenTextFile: Hover
+//===----------------------------------------------------------------------===//
+
+Optional<lsp::Hover>
+TableGenTextFile::findHover(const lsp::URIForFile &uri,
+                            const lsp::Position &hoverPos) {
+  // Check for a reference to an include.
+  for (const lsp::SourceMgrInclude &include : parsedIncludes)
+    if (include.range.contains(hoverPos))
+      return include.buildHover();
+  return llvm::None;
+}
+
 //===----------------------------------------------------------------------===//
 // TableGenServer::Impl
 //===----------------------------------------------------------------------===//
@@ -209,3 +254,18 @@ Optional<int64_t> lsp::TableGenServer::removeDocument(const URIForFile &uri) {
   impl->files.erase(it);
   return version;
 }
+
+void lsp::TableGenServer::getDocumentLinks(
+    const URIForFile &uri, std::vector<DocumentLink> &documentLinks) {
+  auto fileIt = impl->files.find(uri.file());
+  if (fileIt != impl->files.end())
+    return fileIt->second->getDocumentLinks(uri, documentLinks);
+}
+
+Optional<lsp::Hover> lsp::TableGenServer::findHover(const URIForFile &uri,
+                                                    const Position &hoverPos) {
+  auto fileIt = impl->files.find(uri.file());
+  if (fileIt != impl->files.end())
+    return fileIt->second->findHover(uri, hoverPos);
+  return llvm::None;
+}

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

@@ -17,6 +17,9 @@
 namespace mlir {
 namespace lsp {
 struct Diagnostic;
+struct DocumentLink;
+struct Hover;
+struct Position;
 class URIForFile;
 
 /// This class implements all of the TableGen related functionality necessary
@@ -52,6 +55,14 @@ public:
   /// the server.
   Optional<int64_t> removeDocument(const URIForFile &uri);
 
+  /// Return the document links referenced by the given file.
+  void getDocumentLinks(const URIForFile &uri,
+                        std::vector<DocumentLink> &documentLinks);
+
+  /// Find a hover description for the given hover position, or None if one
+  /// couldn't be found.
+  Optional<Hover> findHover(const URIForFile &uri, const Position &hoverPos);
+
 private:
   struct Impl;
   std::unique_ptr<Impl> impl;

mlir/test/tblgen-lsp-server/document-links.test

@@ -0,0 +1,34 @@
+// RUN: tblgen-lsp-server -tablegen-extra-dir %S -tablegen-extra-dir %S/../../include -lit-test < %s | FileCheck %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"tablegen","capabilities":{},"trace":"off"}}
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///foo.td",
+  "languageId":"tablegen",
+  "version":1,
+  "text":"include \"include/included.td\""
+}}}
+// -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/documentLink","params":{
+  "textDocument":{"uri":"test:///foo.td"}
+}}
+//      CHECK:  "id": 1,
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 29,
+// CHECK-NEXT:          "line": 0
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 8,
+// CHECK-NEXT:          "line": 0
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "target": "file:{{.*}}included.td"
+// CHECK-NEXT:    }
+// CHECK-NEXT:  ]
+// -----
+{"jsonrpc":"2.0","id":7,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}

mlir/test/tblgen-lsp-server/hover.test

@@ -0,0 +1,37 @@
+// RUN: tblgen-lsp-server -tablegen-extra-dir %S -tablegen-extra-dir %S/../../include -lit-test < %s | FileCheck %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"tablegen","capabilities":{},"trace":"off"}}
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///foo.td",
+  "languageId":"tablegen",
+  "version":1,
+  "text":"include \"include/included.td\""
+}}}
+// -----
+// Hover on an include file.
+{"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
+  "textDocument":{"uri":"test:///foo.td"},
+  "position":{"line":0,"character":15}
+}}
+//      CHECK:  "id": 1,
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "contents": {
+// CHECK-NEXT:      "kind": "markdown",
+// CHECK-NEXT:      "value": "`included.td`\n***\n{{.*}}included.td"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    "range": {
+// CHECK-NEXT:      "end": {
+// CHECK-NEXT:        "character": 29,
+// CHECK-NEXT:        "line": 0
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "start": {
+// CHECK-NEXT:        "character": 8,
+// CHECK-NEXT:        "line": 0
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+{"jsonrpc":"2.0","id":7,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}

mlir/test/tblgen-lsp-server/initialize-params.test

@@ -5,6 +5,10 @@
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "capabilities": {
+// CHECK-NEXT:    "documentLinkProvider": {
+// CHECK-NEXT:      "resolveProvider": false
+// CHECK-NEXT:    },
+// CHECK-NEXT:    "hoverProvider": true,
 // CHECK-NEXT:      "textDocumentSync": {
 // CHECK-NEXT:        "change": 1,
 // CHECK-NEXT:        "openClose": true,