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

[rriddle]

Summary

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

[jpienaar]

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

504

OOC does this enable linkage? Or would one need to write this in markdown formatting?

mlir/test/mlir-pdll-lsp-server/include/included.td

1 ↗

(On Diff #423849)

Are the empty new lines at top on purpose? File heading missing?

[rriddle]

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

504

The documentLink automatically has linkage, that is what gets you the ctrl+click behavior. The hover here is just additional information about the include. I just modeled the info in the hover after what clangd provides.

mlir/test/mlir-pdll-lsp-server/include/included.td

1 ↗

(On Diff #423849)

Dropped the newlines at the top. I didn't add a file heading because no other tests I could find have them, and these files are not tested directly (thus no RUN lines), they are only used when testing other files.

llvm/include/llvm/Support/SourceMgr.h

unsigned AddNewSourceBuffer(std::unique_ptr<MemoryBuffer> F,
     SrcBuffer NB;
     NB.Buffer = std::move(F);
     NB.IncludeLoc = IncludeLoc;
     Buffers.push_back(std::move(NB));
     return Buffers.size();
   }
 
   /// 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
   /// one of the IncludeDirs.
   ///
   /// If no file is found, this returns 0, otherwise it returns the buffer ID
   /// of the stacked file. The full path to the included file can be found in
   /// \p IncludedFile.

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

public:
   llvm::SourceMgr &getSourceMgr() { return srcMgr; }
 
   /// Return a reference to the diagnostic engine used by the lexer.
   ast::DiagnosticEngine &getDiagEngine() { return diagEngine; }
 
   /// 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();
 
   /// Change the position of the lexer cursor. The next token we lex will start
   /// at the designated point in the input.
   void resetPointer(const char *newPointer) { curPtr = newPointer; }
 

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

if (!diagEngine.getHandlerFn()) {
     addedHandlerToDiagEngine = true;
   }
 }
 
 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(
-      filename.str(), SMLoc::getFromPointer(curPtr), includedFile);
-  if (!bufferID) return failure();
+  int bufferID =
+      srcMgr.AddIncludeFile(filename.str(), includeLoc.End, includedFile);
+  if (!bufferID)
+    return failure();
 
   curBufferID = bufferID;
   curBuffer = srcMgr.getMemoryBuffer(curBufferID)->getBuffer();
   curPtr = curBuffer.begin();
   return success();
 }
 
 Token Lexer::emitError(SMRange loc, const Twine &msg) {

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

LogicalResult Parser::parseInclude(SmallVectorImpl<ast::Decl *> &decls) {
   SMRange fileLoc = curToken.getLoc();
   std::string filenameStr = curToken.getStringValue();
   StringRef filename = filenameStr;
   consumeToken();
 
   // 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
-    // finish parsing the nested file.
-    Token oldToken = curToken;
+    // Make sure to update to the next token after we finish parsing the nested
+    // file.
     curToken = lexer.lexToken();
     LogicalResult result = parseModuleBody(decls);
-    curToken = oldToken;
+    curToken = lexer.lexToken();
     return result;
   }
 
   // Otherwise, this must be a `.td` include.
   if (filename.endswith(".td"))
     return parseTdInclude(filename, fileLoc, decls);
 
   return emitError(fileLoc,
if (!includeBuffer)
     return emitError(fileLoc, "unable to open include file `" + filename + "`");
 
   auto processFn = [&](llvm::RecordKeeper &records) {
     processTdIncludeRecords(records, decls);
 
     // 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),
                               parserSrcMgr.getIncludeDirs(), processFn))
     return failure();
 
   return success();
 }

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

struct SignatureHelp {
 
   /// The active parameter of the active signature.
   int activeParameter = 0;
 };
 
 /// 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
 
 namespace llvm {
 template <>
 struct format_provider<mlir::lsp::Position> {
   static void format(const mlir::lsp::Position &pos, raw_ostream &os,
                      StringRef style) {
     assert(style.empty() && "style modifiers for this type are not supported");
     os << pos;
   }
 };
 } // namespace llvm
 
 #endif

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

llvm::json::Value mlir::lsp::toJSON(const SignatureHelp &value) {
   assert(value.activeParameter >= 0 &&
          "Unexpected negative value for active parameter index");
   return llvm::json::Object{
       {"activeSignature", value.activeSignature},
       {"activeParameter", value.activeParameter},
       {"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/mlir-pdll-lsp-server/LSPServer.cpp

struct LSPServer {
   //===--------------------------------------------------------------------===//
   // Definitions and References
 
   void onGoToDefinition(const TextDocumentPositionParams &params,
                         Callback<std::vector<Location>> reply);
   void onReference(const ReferenceParams &params,
                    Callback<std::vector<Location>> reply);
 
+  //===----------------------------------------------------------------------===//
+  // DocumentLink
+
+  void onDocumentLink(const DocumentLinkParams &params,
+                      Callback<std::vector<DocumentLink>> reply);
+
   //===--------------------------------------------------------------------===//
   // Hover
 
   void onHover(const TextDocumentPositionParams &params,
                Callback<Optional<Hover>> reply);
 
   //===--------------------------------------------------------------------===//
   // Document Symbols
llvm::json::Object serverCaps{
            {"triggerCharacters", {".", ">", "(", "{", ",", "<", ":", "[", " "}},
        }},
       {"signatureHelpProvider",
        llvm::json::Object{
            {"triggerCharacters", {"(", ","}},
        }},
       {"definitionProvider", true},
       {"referencesProvider", true},
+      {"documentLinkProvider",
+       llvm::json::Object{
+           {"resolveProvider", false},
+       }},
       {"hoverProvider", true},
       {"documentSymbolProvider", true},
   };
 
   llvm::json::Object result{
       {{"serverInfo", llvm::json::Object{{"name", "mlir-pdll-lsp-server"},
                                          {"version", "0.0.1"}}},
        {"capabilities", std::move(serverCaps)}}};

 void LSPServer::onReference(const ReferenceParams &params,
                             Callback<std::vector<Location>> reply) {
   std::vector<Location> locations;
   server.findReferencesOf(params.textDocument.uri, params.position, locations);
   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
 
 void LSPServer::onHover(const TextDocumentPositionParams &params,
                         Callback<Optional<Hover>> reply) {
   reply(server.findHover(params.textDocument.uri, params.position));
 }
 
 //===----------------------------------------------------------------------===//
messageHandler.notification("textDocument/didChange", &lspServer,
                               &LSPServer::onDocumentDidChange);
 
   // Definitions and References
   messageHandler.method("textDocument/definition", &lspServer,
                         &LSPServer::onGoToDefinition);
   messageHandler.method("textDocument/references", &lspServer,
                         &LSPServer::onReference);
 
+  // Document Link
+  messageHandler.method("textDocument/documentLink", &lspServer,
+                        &LSPServer::onDocumentLink);
+
   // Hover
   messageHandler.method("textDocument/hover", &lspServer, &LSPServer::onHover);
 
   // Document Symbols
   messageHandler.method("textDocument/documentSymbol", &lspServer,
                         &LSPServer::onDocumentSymbol);
 
   // Code Completion

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

 #include <memory>
 #include <string>
 
 namespace mlir {
 namespace lsp {
 struct Diagnostic;
 class CompilationDatabase;
 struct CompletionList;
+struct DocumentLink;
 struct DocumentSymbol;
 struct Hover;
 struct Location;
 struct Position;
 struct SignatureHelp;
 class URIForFile;
 
 /// This class implements all of the PDLL related functionality necessary for a
public:
   /// Return the locations of the object pointed at by the given position.
   void getLocationsOf(const URIForFile &uri, const Position &defPos,
                       std::vector<Location> &locations);
 
   /// Find all references of the object pointed at by the given position.
   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);
 
   /// Find all of the document symbols within the given file.
   void findDocumentSymbols(const URIForFile &uri,
                            std::vector<DocumentSymbol> &symbols);
 

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

getLspDiagnoticFromDiag(llvm::SourceMgr &sourceMgr, const ast::Diagnostic &diag,
   }
   if (!relatedDiags.empty())
     lspDiag.relatedInformation = std::move(relatedDiags);
 
   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
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct PDLIndexSymbol {
   explicit PDLIndexSymbol(const ast::Decl *definition)
       : definition(definition) {}
   explicit PDLIndexSymbol(const ods::Operation *definition)
struct PDLDocument {
   //===--------------------------------------------------------------------===//
 
   void getLocationsOf(const lsp::URIForFile &uri, const lsp::Position &defPos,
                       std::vector<lsp::Location> &locations);
   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
   //===--------------------------------------------------------------------===//
 
   Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
                                  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,
                                    const SMRange &hoverRange);
   lsp::Hover buildHoverForPattern(const ast::PatternDecl *patternDecl,
                                   const SMRange &hoverRange);
   lsp::Hover buildHoverForCoreConstraint(const ast::CoreConstraintDecl *decl,
                                          const SMRange &hoverRange);
struct PDLDocument {
   ods::Context odsContext;
   ast::Context astContext;
 
   /// The parsed AST module, or failure if the file wasn't valid.
   FailureOr<ast::Module *> astModule;
 
   /// The index of the parsed module.
   PDLIndex index;
+
+  /// The set of includes of the parsed module.
+  std::vector<PDLLInclude> parsedIncludes;
 };
 } // namespace
 
 PDLDocument::PDLDocument(const lsp::URIForFile &uri, StringRef contents,
                          const std::vector<std::string> &extraDirs,
                          std::vector<lsp::Diagnostic> &diagnostics)
     : astContext(odsContext) {
   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
PDLDocument::PDLDocument(const lsp::URIForFile &uri, StringRef contents,
   sourceMgr.setIncludeDirs(includeDirs);
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
 
   astContext.getDiagEngine().setHandlerFn([&](const ast::Diagnostic &diag) {
     if (auto lspDiag = getLspDiagnoticFromDiag(sourceMgr, diag, uri))
       diagnostics.push_back(std::move(*lspDiag));
   });
   astModule = parsePDLAST(astContext, sourceMgr);
-  if (succeeded(astModule))
-    index.initialize(**astModule, odsContext);
+
+  // 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.
 }
 
 //===----------------------------------------------------------------------===//
 // PDLDocument: Definitions and References
 //===----------------------------------------------------------------------===//
 
 void PDLDocument::getLocationsOf(const lsp::URIForFile &uri,
                                  const lsp::Position &defPos,
void PDLDocument::findReferencesOf(const lsp::URIForFile &uri,
   if (!symbol)
     return;
 
   references.push_back(getLocationFromLoc(sourceMgr, symbol->getDefLoc(), uri));
   for (SMRange refLoc : symbol->references)
     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
 //===----------------------------------------------------------------------===//
 
 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)
     return llvm::None;
 
   // Add hover for operation names.
   if (const auto *op = symbol->definition.dyn_cast<const ods::Operation *>())
     return buildHoverForOpName(op, hoverRange);
Optional<lsp::Hover> PDLDocument::findHover(const ast::Decl *decl,
 
   // Add hover for user rewrites.
   if (const auto *rewrite = dyn_cast<ast::UserRewriteDecl>(decl))
     return buildHoverForUserConstraintOrRewrite("Rewrite", rewrite, hoverRange);
 
   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));
   {
     llvm::raw_string_ostream hoverOS(hover.contents.value);
     hoverOS << "**OpName**: `" << op->getName() << "`\n***\n"
             << op->getSummary() << "\n***\n"
             << op->getDescription();
public:
   //===--------------------------------------------------------------------===//
   // LSP Queries
   //===--------------------------------------------------------------------===//
 
   void getLocationsOf(const lsp::URIForFile &uri, lsp::Position defPos,
                       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);
   lsp::CompletionList getCodeCompletion(const lsp::URIForFile &uri,
                                         lsp::Position completePos);
   lsp::SignatureHelp getSignatureHelp(const lsp::URIForFile &uri,
                                       lsp::Position helpPos);
 
void PDLTextFile::findReferencesOf(const lsp::URIForFile &uri,
   // Adjust any locations within this file for the offset of this chunk.
   if (chunk.lineOffset == 0)
     return;
   for (lsp::Location &loc : references)
     if (loc.uri == 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);
   Optional<lsp::Hover> hoverInfo = chunk.document.findHover(uri, hoverPos);
 
   // Adjust any locations within this file for the offset of this chunk.
   if (chunk.lineOffset != 0 && hoverInfo && hoverInfo->range)
     chunk.adjustLocForChunkOffset(*hoverInfo->range);

 void lsp::PDLLServer::findReferencesOf(const URIForFile &uri,
                                        const Position &pos,
                                        std::vector<Location> &references) {
   auto fileIt = impl->files.find(uri.file());
   if (fileIt != impl->files.end())
     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());
   if (fileIt != impl->files.end())
     return fileIt->second->findHover(uri, hoverPos);
   return llvm::None;
 }
 

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

+// 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

-// 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.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
   "position":{"line":3,"character":6}
 }}
 //      CHECK:  "id": 1,

 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 8,
 // CHECK-NEXT:        "line": 1
 // 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":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

 // RUN: mlir-pdll-lsp-server -lit-test < %s | FileCheck %s
 // Test initialize request parameters with rootUri
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootUri":"test:///workspace","capabilities":{},"trace":"off"}}
 //      CHECK:  "id": 0,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "capabilities": {
 // CHECK-NEXT:      "completionProvider": {
 // CHECK-NEXT:        "allCommitCharacters": [
 // CHECK:             ],
 // CHECK-NEXT:        "resolveProvider": false,
 // CHECK-NEXT:        "triggerCharacters": [ 
 // 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,
 // CHECK-NEXT:      "signatureHelpProvider": {
 // CHECK-NEXT:        "triggerCharacters": [
 // CHECK-NEXT:          "(",
 // CHECK-NEXT:          ","
 // CHECK-NEXT:        ]