[mlir][PDLL] Add document link and hover support to mlir-pdll-lsp-server

This allows for navigating to included files on click, and also provides hover
information about the include file (similarly to clangd).

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

llvm/include/llvm/Support/SourceMgr.h

@@ -151,11 +151,19 @@ public:
   }
 
   /// Takes the source buffers from the given source manager and append them to
+  /// the current manager. `MainBufferIncludeLoc` is an optional include
+  /// location to attach to the main buffer of `SrcMgr` after it gets moved to
   /// the current manager.
-  void takeSourceBuffersFrom(SourceMgr &SrcMgr) {
+  void takeSourceBuffersFrom(SourceMgr &SrcMgr,
+                             SMLoc MainBufferIncludeLoc = SMLoc()) {
+    if (SrcMgr.Buffers.empty())
+      return;
+
+    size_t OldNumBuffers = getNumBuffers();
     std::move(SrcMgr.Buffers.begin(), SrcMgr.Buffers.end(),
               std::back_inserter(Buffers));
     SrcMgr.Buffers.clear();
+    Buffers[OldNumBuffers].IncludeLoc = MainBufferIncludeLoc;
   }
 
   /// Search for a file with the specified name in the current directory or in

mlir/lib/Tools/PDLL/Parser/Lexer.cpp

@@ -100,11 +100,12 @@ Lexer::~Lexer() {
   if (addedHandlerToDiagEngine) diagEngine.setHandlerFn(nullptr);
 }
 
-LogicalResult Lexer::pushInclude(StringRef filename) {
+LogicalResult Lexer::pushInclude(StringRef filename, SMRange includeLoc) {
   std::string includedFile;
-  int bufferID = srcMgr.AddIncludeFile(
+  int bufferID =
-      filename.str(), SMLoc::getFromPointer(curPtr), includedFile);
+      srcMgr.AddIncludeFile(filename.str(), includeLoc.End, includedFile);
-  if (!bufferID) return failure();
+  if (!bufferID)
+    return failure();
 
   curBufferID = bufferID;
   curBuffer = srcMgr.getMemoryBuffer(curBufferID)->getBuffer();

mlir/lib/Tools/PDLL/Parser/Lexer.h

@@ -178,7 +178,7 @@ public:
   /// Push an include of the given file. This will cause the lexer to start
   /// processing the provided file. Returns failure if the file could not be
   /// opened, success otherwise.
-  LogicalResult pushInclude(StringRef filename);
+  LogicalResult pushInclude(StringRef filename, SMRange includeLoc);
 
   /// Lex the next token and return it.
   Token lexToken();

mlir/lib/Tools/PDLL/Parser/Parser.cpp

@@ -692,17 +692,16 @@ LogicalResult Parser::parseInclude(SmallVectorImpl<ast::Decl *> &decls) {
   // Check the type of include. If ending with `.pdll`, this is another pdl file
   // to be parsed along with the current module.
   if (filename.endswith(".pdll")) {
-    if (failed(lexer.pushInclude(filename)))
+    if (failed(lexer.pushInclude(filename, fileLoc)))
       return emitError(fileLoc,
                        "unable to open include file `" + filename + "`");
 
     // If we added the include successfully, parse it into the current module.
-    // Make sure to save the current token so that we can restore it when we
+    // Make sure to update to the next token after we finish parsing the nested
-    // finish parsing the nested file.
+    // file.
-    Token oldToken = curToken;
     curToken = lexer.lexToken();
     LogicalResult result = parseModuleBody(decls);
-    curToken = oldToken;
+    curToken = lexer.lexToken();
     return result;
   }
 
@@ -750,7 +749,7 @@ LogicalResult Parser::parseTdInclude(StringRef filename, llvm::SMRange fileLoc,
     // After we are done processing, move all of the tablegen source buffers to
     // the main parser source mgr. This allows for directly using source
     // locations from the .td files without needing to remap them.
-    parserSrcMgr.takeSourceBuffersFrom(llvm::SrcMgr);
+    parserSrcMgr.takeSourceBuffersFrom(llvm::SrcMgr, fileLoc.End);
     return false;
   };
   if (llvm::TableGenParseFile(std::move(*includeBuffer),

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

@@ -796,3 +796,24 @@ llvm::json::Value mlir::lsp::toJSON(const SignatureHelp &value) {
       {"signatures", llvm::json::Array(value.signatures)},
   };
 }
+
+//===----------------------------------------------------------------------===//
+// DocumentLinkParams
+//===----------------------------------------------------------------------===//
+
+bool mlir::lsp::fromJSON(const llvm::json::Value &value,
+                         DocumentLinkParams &result, llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  return o && o.map("textDocument", result.textDocument);
+}
+
+//===----------------------------------------------------------------------===//
+// DocumentLink
+//===----------------------------------------------------------------------===//
+
+llvm::json::Value mlir::lsp::toJSON(const DocumentLink &value) {
+  return llvm::json::Object{
+      {"range", value.range},
+      {"target", value.target},
+  };
+}

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

@@ -930,6 +930,56 @@ struct SignatureHelp {
 /// Add support for JSON serialization.
 llvm::json::Value toJSON(const SignatureHelp &value);
 
+//===----------------------------------------------------------------------===//
+// DocumentLinkParams
+//===----------------------------------------------------------------------===//
+
+/// Parameters for the document link request.
+struct DocumentLinkParams {
+  /// The document to provide document links for.
+  TextDocumentIdentifier textDocument;
+};
+
+/// Add support for JSON serialization.
+bool fromJSON(const llvm::json::Value &value, DocumentLinkParams &result,
+              llvm::json::Path path);
+
+//===----------------------------------------------------------------------===//
+// DocumentLink
+//===----------------------------------------------------------------------===//
+
+/// A range in a text document that links to an internal or external resource,
+/// like another text document or a web site.
+struct DocumentLink {
+  DocumentLink() = default;
+  DocumentLink(Range range, URIForFile target)
+      : range(range), target(std::move(target)) {}
+
+  /// The range this link applies to.
+  Range range;
+
+  /// The uri this link points to. If missing a resolve request is sent later.
+  URIForFile target;
+
+  // TODO: The following optional fields defined by the language server protocol
+  // are unsupported:
+  //
+  // data?: any - A data entry field that is preserved on a document link
+  //              between a DocumentLinkRequest and a
+  //              DocumentLinkResolveRequest.
+
+  friend bool operator==(const DocumentLink &lhs, const DocumentLink &rhs) {
+    return lhs.range == rhs.range && lhs.target == rhs.target;
+  }
+
+  friend bool operator!=(const DocumentLink &lhs, const DocumentLink &rhs) {
+    return !(lhs == rhs);
+  }
+};
+
+/// Add support for JSON serialization.
+llvm::json::Value toJSON(const DocumentLink &value);
+
 } // namespace lsp
 } // namespace mlir
 

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

@@ -52,6 +52,12 @@ struct LSPServer {
   void onReference(const ReferenceParams &params,
                    Callback<std::vector<Location>> reply);
 
+  //===----------------------------------------------------------------------===//
+  // DocumentLink
+
+  void onDocumentLink(const DocumentLinkParams &params,
+                      Callback<std::vector<DocumentLink>> reply);
+
   //===--------------------------------------------------------------------===//
   // Hover
 
@@ -121,6 +127,10 @@ void LSPServer::onInitialize(const InitializeParams &params,
        }},
       {"definitionProvider", true},
       {"referencesProvider", true},
+      {"documentLinkProvider",
+       llvm::json::Object{
+           {"resolveProvider", false},
+       }},
       {"hoverProvider", true},
       {"documentSymbolProvider", true},
   };
@@ -193,6 +203,16 @@ void LSPServer::onReference(const ReferenceParams &params,
   reply(std::move(locations));
 }
 
+//===----------------------------------------------------------------------===//
+// 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
 
@@ -256,6 +276,10 @@ LogicalResult mlir::lsp::runPdllLSPServer(PDLLServer &server,
   messageHandler.method("textDocument/references", &lspServer,
                         &LSPServer::onReference);
 
+  // Document Link
+  messageHandler.method("textDocument/documentLink", &lspServer,
+                        &LSPServer::onDocumentLink);
+
   // Hover
   messageHandler.method("textDocument/hover", &lspServer, &LSPServer::onHover);
 

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

@@ -104,6 +104,24 @@ 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
 //===----------------------------------------------------------------------===//
@@ -253,6 +271,13 @@ struct PDLDocument {
   void findReferencesOf(const lsp::URIForFile &uri, const lsp::Position &pos,
                         std::vector<lsp::Location> &references);
 
+  //===--------------------------------------------------------------------===//
+  // Document Links
+  //===--------------------------------------------------------------------===//
+
+  void getDocumentLinks(const lsp::URIForFile &uri,
+                        std::vector<lsp::DocumentLink> &links);
+
   //===--------------------------------------------------------------------===//
   // Hover
   //===--------------------------------------------------------------------===//
@@ -261,6 +286,7 @@ 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,
@@ -313,6 +339,9 @@ struct PDLDocument {
 
   /// The index of the parsed module.
   PDLIndex index;
+
+  /// The set of includes of the parsed module.
+  std::vector<PDLLInclude> parsedIncludes;
 };
 } // namespace
 
@@ -340,7 +369,41 @@ PDLDocument::PDLDocument(const lsp::URIForFile &uri, StringRef contents,
       diagnostics.push_back(std::move(*lspDiag));
   });
   astModule = parsePDLAST(astContext, sourceMgr);
-  if (succeeded(astModule))
+
+  // 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));
+  }
+
+  // If we failed to parse the module, there is nothing left to initialize.
+  if (failed(astModule))
+    return;
+
+  // Prepare the AST index with the parsed module.
   index.initialize(**astModule, odsContext);
 }
 
@@ -372,6 +435,16 @@ void PDLDocument::findReferencesOf(const lsp::URIForFile &uri,
     references.push_back(getLocationFromLoc(sourceMgr, refLoc, uri));
 }
 
+//===--------------------------------------------------------------------===//
+// PDLDocument: Document Links
+//===--------------------------------------------------------------------===//
+
+void PDLDocument::getDocumentLinks(const lsp::URIForFile &uri,
+                                   std::vector<lsp::DocumentLink> &links) {
+  for (const PDLLInclude &include : parsedIncludes)
+    links.emplace_back(include.range, include.uri);
+}
+
 //===----------------------------------------------------------------------===//
 // PDLDocument: Hover
 //===----------------------------------------------------------------------===//
@@ -379,6 +452,14 @@ void PDLDocument::findReferencesOf(const lsp::URIForFile &uri,
 Optional<lsp::Hover> PDLDocument::findHover(const lsp::URIForFile &uri,
                                             const lsp::Position &hoverPos) {
   SMLoc posLoc = hoverPos.getAsSMLoc(sourceMgr);
+
+  // Check for a reference to an include.
+  for (const PDLLInclude &include : parsedIncludes) {
+    if (include.range.contains(hoverPos))
+      return buildHoverForInclude(include);
+  }
+
+  // Find the symbol at the given location.
   SMRange hoverRange;
   const PDLIndexSymbol *symbol = index.lookup(posLoc, &hoverRange);
   if (!symbol)
@@ -416,6 +497,17 @@ 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));
@@ -1040,6 +1132,8 @@ public:
                       std::vector<lsp::Location> &locations);
   void findReferencesOf(const lsp::URIForFile &uri, lsp::Position pos,
                         std::vector<lsp::Location> &references);
+  void getDocumentLinks(const lsp::URIForFile &uri,
+                        std::vector<lsp::DocumentLink> &links);
   Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
                                  lsp::Position hoverPos);
   void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);
@@ -1135,6 +1229,20 @@ void PDLTextFile::findReferencesOf(const lsp::URIForFile &uri,
       chunk.adjustLocForChunkOffset(loc.range);
 }
 
+void PDLTextFile::getDocumentLinks(const lsp::URIForFile &uri,
+                                   std::vector<lsp::DocumentLink> &links) {
+  chunks.front()->document.getDocumentLinks(uri, links);
+  for (const auto &it : llvm::drop_begin(chunks)) {
+    size_t currentNumLinks = links.size();
+    it->document.getDocumentLinks(uri, links);
+
+    // Adjust any links within this file to account for the offset of this
+    // chunk.
+    for (auto &link : llvm::drop_begin(links, currentNumLinks))
+      it->adjustLocForChunkOffset(link.range);
+  }
+}
+
 Optional<lsp::Hover> PDLTextFile::findHover(const lsp::URIForFile &uri,
                                             lsp::Position hoverPos) {
   PDLTextFileChunk &chunk = getChunkFor(hoverPos);
@@ -1285,6 +1393,13 @@ void lsp::PDLLServer::findReferencesOf(const URIForFile &uri,
     fileIt->second->findReferencesOf(uri, pos, references);
 }
 
+void lsp::PDLLServer::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::PDLLServer::findHover(const URIForFile &uri,
                                                 const Position &hoverPos) {
   auto fileIt = impl->files.find(uri.file());

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

@@ -19,6 +19,7 @@ namespace lsp {
 struct Diagnostic;
 class CompilationDatabase;
 struct CompletionList;
+struct DocumentLink;
 struct DocumentSymbol;
 struct Hover;
 struct Location;
@@ -67,6 +68,10 @@ public:
   void findReferencesOf(const URIForFile &uri, const Position &pos,
                         std::vector<Location> &references);
 
+  /// 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);

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

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

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

@@ -1,11 +1,11 @@
-// RUN: mlir-pdll-lsp-server -lit-test < %s | FileCheck %s
+// RUN: mlir-pdll-lsp-server -pdll-extra-dir %S -pdll-extra-dir %S/../../include -lit-test < %s | FileCheck %s
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"pdll","capabilities":{},"trace":"off"}}
 // -----
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
   "uri":"test:///foo.pdll",
   "languageId":"pdll",
   "version":1,
-  "text":"Constraint FooCst();\nRewrite FooRewrite(op: Op) -> Op;\nPattern Foo {\nlet root: Op;\nerase root;\n}"
+  "text":"Constraint FooCst();\nRewrite FooRewrite(op: Op) -> Op;\nPattern Foo {\nlet root: Op;\nerase root;\n}\n#include \"include/included.td\"\n#include \"include/included.pdll\""
 }}}
 // -----
 // Hover on a variable.
@@ -128,6 +128,54 @@
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
+// Hover on an include file.
+{"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
+  "textDocument":{"uri":"test:///foo.pdll"},
+  "position":{"line":6,"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": 30,
+// CHECK-NEXT:        "line": 6
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "start": {
+// CHECK-NEXT:        "character": 9,
+// CHECK-NEXT:        "line": 6
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+// Hover on an include file.
+{"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
+  "textDocument":{"uri":"test:///foo.pdll"},
+  "position":{"line":7,"character":15}
+}}
+//      CHECK:  "id": 1,
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "contents": {
+// CHECK-NEXT:      "kind": "markdown",
+// CHECK-NEXT:      "value": "`included.pdll`\n***\n{{.*}}included.pdll"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    "range": {
+// CHECK-NEXT:      "end": {
+// CHECK-NEXT:        "character": 32,
+// CHECK-NEXT:        "line": 7
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "start": {
+// CHECK-NEXT:        "character": 9,
+// CHECK-NEXT:        "line": 7
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
 {"jsonrpc":"2.0","id":7,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}

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

@@ -13,6 +13,9 @@
 // CHECK:             ]
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "definitionProvider": true,
+// CHECK-NEXT:      "documentLinkProvider": {
+// CHECK-NEXT:        "resolveProvider": false
+// CHECK-NEXT:      },
 // CHECK-NEXT:      "documentSymbolProvider": true,
 // CHECK-NEXT:      "hoverProvider": true,
 // CHECK-NEXT:      "referencesProvider": true,