[mlir][PDLL] Add hover support to the PDLL language server

[rriddle]

Summary

This adds support for providing information when hovering over
operation names, variables, patters, constraints, and rewrites.

[jpienaar]

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

428

Variable (to be consistent)

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

struct LSPServer::Impl {
   // Definitions and References
 
   void onGoToDefinition(const TextDocumentPositionParams &params,
                         Callback<std::vector<Location>> reply);
   void onReference(const ReferenceParams &params,
                    Callback<std::vector<Location>> reply);
 
   //===--------------------------------------------------------------------===//
+  // Hover
+
+  void onHover(const TextDocumentPositionParams &params,
+               Callback<Optional<Hover>> reply);
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   PDLLServer &server;
   JSONTransport &transport;
 
   /// An outgoing notification used to send diagnostics to the client when they
   /// are ready to be processed.
llvm::json::Object serverCaps{
       {"textDocumentSync",
        llvm::json::Object{
            {"openClose", true},
            {"change", (int)TextDocumentSyncKind::Full},
            {"save", true},
        }},
       {"definitionProvider", true},
       {"referencesProvider", true},
+      {"hoverProvider", true},
   };
 
   llvm::json::Object result{
       {{"serverInfo", llvm::json::Object{{"name", "mlir-pdll-lsp-server"},
                                          {"version", "0.0.0"}}},
        {"capabilities", std::move(serverCaps)}}};
   reply(std::move(result));
 }

 void LSPServer::Impl::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));
 }
 
 //===----------------------------------------------------------------------===//
+// Hover
+
+void LSPServer::Impl::onHover(const TextDocumentPositionParams &params,
+                              Callback<Optional<Hover>> reply) {
+  reply(server.findHover(params.textDocument.uri, params.position));
+}
+
+//===----------------------------------------------------------------------===//
 // LSPServer
 //===----------------------------------------------------------------------===//
 
 LSPServer::LSPServer(PDLLServer &server, JSONTransport &transport)
     : impl(std::make_unique<Impl>(server, transport)) {}
 LSPServer::~LSPServer() = default;
 
 LogicalResult LSPServer::run() {
messageHandler.notification("textDocument/didChange", impl.get(),
                               &Impl::onDocumentDidChange);
 
   // Definitions and References
   messageHandler.method("textDocument/definition", impl.get(),
                         &Impl::onGoToDefinition);
   messageHandler.method("textDocument/references", impl.get(),
                         &Impl::onReference);
 
+  // Hover
+  messageHandler.method("textDocument/hover", impl.get(), &Impl::onHover);
+
   // Diagnostics
   impl->publishDiagnostics =
       messageHandler.outgoingNotification<PublishDiagnosticsParams>(
           "textDocument/publishDiagnostics");
 
   // Run the main loop of the transport.
   if (llvm::Error error = impl->transport.run(messageHandler)) {
     Logger::error("Transport error: {0}", error);
     llvm::consumeError(std::move(error));
     return failure();
   }
   return success(impl->shutdownRequestReceived);
 }

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

 #define LIB_MLIR_TOOLS_MLIRPDLLSPSERVER_SERVER_H_
 
 #include "mlir/Support/LLVM.h"
 #include <memory>
 
 namespace mlir {
 namespace lsp {
 struct Diagnostic;
+struct Hover;
 struct Location;
 struct Position;
 class URIForFile;
 
 /// This class implements all of the PDLL related functionality necessary for a
 /// language server. This class allows for keeping the PDLL specific logic
 /// separate from the logic that involves LSP server/client communication.
 class PDLLServer {
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);
 
+  /// 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;
 };
 
 } // namespace lsp
 } // namespace mlir
 
 #endif // LIB_MLIR_TOOLS_MLIRPDLLSPSERVER_SERVER_H_

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

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);
 
   //===--------------------------------------------------------------------===//
+  // Hover
+  //===--------------------------------------------------------------------===//
+
+  Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
+                                 const lsp::Position &hoverPos);
+  Optional<lsp::Hover> findHover(const ast::Decl *decl,
+                                 const llvm::SMRange &hoverRange);
+  lsp::Hover buildHoverForOpName(const ods::Operation *op,
+                                 const llvm::SMRange &hoverRange);
+  lsp::Hover buildHoverForVariable(const ast::VariableDecl *varDecl,
+                                   const llvm::SMRange &hoverRange);
+  lsp::Hover buildHoverForPattern(const ast::PatternDecl *patternDecl,
+                                  const llvm::SMRange &hoverRange);
+  lsp::Hover buildHoverForCoreConstraint(const ast::CoreConstraintDecl *decl,
+                                         const llvm::SMRange &hoverRange);
+  template <typename T>
+  lsp::Hover
+  buildHoverForUserConstraintOrRewrite(StringRef typeName, const T *decl,
+                                       const llvm::SMRange &hoverRange);
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   /// The include directories for this file.
   std::vector<std::string> includeDirs;
 
   /// The source manager containing the contents of the input file.
   llvm::SourceMgr sourceMgr;
if (!symbol)
     return;
 
   references.push_back(getLocationFromLoc(sourceMgr, symbol->getDefLoc(), uri));
   for (llvm::SMRange refLoc : symbol->references)
     references.push_back(getLocationFromLoc(sourceMgr, refLoc, uri));
 }
 
 //===----------------------------------------------------------------------===//
+// PDLDocument: Hover
+//===----------------------------------------------------------------------===//
+
+Optional<lsp::Hover> PDLDocument::findHover(const lsp::URIForFile &uri,
+                                            const lsp::Position &hoverPos) {
+  llvm::SMLoc posLoc = getPosFromLoc(sourceMgr, hoverPos);
+  llvm::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);
+  const auto *decl = symbol->definition.get<const ast::Decl *>();
+  return findHover(decl, hoverRange);
+}
+
+Optional<lsp::Hover> PDLDocument::findHover(const ast::Decl *decl,
+                                            const llvm::SMRange &hoverRange) {
+  // Add hover for variables.
+  if (const auto *varDecl = dyn_cast<ast::VariableDecl>(decl))
+    return buildHoverForVariable(varDecl, hoverRange);
+
+  // Add hover for patterns.
+  if (const auto *patternDecl = dyn_cast<ast::PatternDecl>(decl))
+    return buildHoverForPattern(patternDecl, hoverRange);
+
+  // Add hover for core constraints.
+  if (const auto *cst = dyn_cast<ast::CoreConstraintDecl>(decl))
+    return buildHoverForCoreConstraint(cst, hoverRange);
+
+  // Add hover for user constraints.
+  if (const auto *cst = dyn_cast<ast::UserConstraintDecl>(decl))
+    return buildHoverForUserConstraintOrRewrite("Constraint", cst, hoverRange);
+
+  // 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::buildHoverForOpName(const ods::Operation *op,
+                                            const llvm::SMRange &hoverRange) {
+  lsp::Hover hover(getRangeFromLoc(sourceMgr, hoverRange));
+  {
+    llvm::raw_string_ostream hoverOS(hover.contents.value);
+    hoverOS << "**OpName**: `" << op->getName() << "`\n***\n"
+            << op->getSummary() << "\n***\n"
+            << op->getDescription();
+  }
+  return hover;
+}
+
+lsp::Hover PDLDocument::buildHoverForVariable(const ast::VariableDecl *varDecl,
+                                              const llvm::SMRange &hoverRange) {
+  lsp::Hover hover(getRangeFromLoc(sourceMgr, hoverRange));
+  {
+    llvm::raw_string_ostream hoverOS(hover.contents.value);
+    hoverOS << "**variable**: `" << varDecl->getName().getName() << "`\n***\n"
+            << "Type: `" << varDecl->getType() << "`\n";
+  }
+  return hover;
+}
+
+lsp::Hover
+PDLDocument::buildHoverForPattern(const ast::PatternDecl *patternDecl,
+                                  const llvm::SMRange &hoverRange) {
+  lsp::Hover hover(getRangeFromLoc(sourceMgr, hoverRange));
+  {
+    llvm::raw_string_ostream hoverOS(hover.contents.value);
+    hoverOS << "**Pattern**";
+    if (const ast::Name *name = patternDecl->getName())
+      hoverOS << ": `" << name->getName() << "`";
+    hoverOS << "\n***\n";
+    if (Optional<uint16_t> benefit = patternDecl->getBenefit())
+      hoverOS << "Benefit: " << *benefit << "\n";
+    if (patternDecl->hasBoundedRewriteRecursion())
+      hoverOS << "HasBoundedRewriteRecursion\n";
+    hoverOS << "RootOp: `"
+            << patternDecl->getRootRewriteStmt()->getRootOpExpr()->getType()
+            << "`\n";
+  }
+  return hover;
+}
+
+lsp::Hover
+PDLDocument::buildHoverForCoreConstraint(const ast::CoreConstraintDecl *decl,
+                                         const llvm::SMRange &hoverRange) {
+  lsp::Hover hover(getRangeFromLoc(sourceMgr, hoverRange));
+  {
+    llvm::raw_string_ostream hoverOS(hover.contents.value);
+    hoverOS << "**Constraint**: `";
+    TypeSwitch<const ast::Decl *>(decl)
+        .Case([&](const ast::AttrConstraintDecl *) { hoverOS << "Attr"; })
+        .Case([&](const ast::OpConstraintDecl *opCst) {
+          hoverOS << "Op";
+          if (Optional<StringRef> name = opCst->getName())
+            hoverOS << "<" << name << ">";
+        })
+        .Case([&](const ast::TypeConstraintDecl *) { hoverOS << "Type"; })
+        .Case([&](const ast::TypeRangeConstraintDecl *) {
+          hoverOS << "TypeRange";
+        })
+        .Case([&](const ast::ValueConstraintDecl *) { hoverOS << "Value"; })
+        .Case([&](const ast::ValueRangeConstraintDecl *) {
+          hoverOS << "ValueRange";
+        });
+    hoverOS << "`\n";
+  }
+  return hover;
+}
+
+template <typename T>
+lsp::Hover PDLDocument::buildHoverForUserConstraintOrRewrite(
+    StringRef typeName, const T *decl, const llvm::SMRange &hoverRange) {
+  lsp::Hover hover(getRangeFromLoc(sourceMgr, hoverRange));
+  {
+    llvm::raw_string_ostream hoverOS(hover.contents.value);
+    hoverOS << "**" << typeName << "**: `" << decl->getName().getName()
+            << "`\n***\n";
+    ArrayRef<ast::VariableDecl *> inputs = decl->getInputs();
+    if (!inputs.empty()) {
+      hoverOS << "Parameters:\n";
+      for (const ast::VariableDecl *input : inputs)
+        hoverOS << "* " << input->getName().getName() << ": `"
+                << input->getType() << "`\n";
+      hoverOS << "***\n";
+    }
+    ast::Type resultType = decl->getResultType();
+    if (auto resultTupleTy = resultType.dyn_cast<ast::TupleType>()) {
+      if (resultTupleTy.empty())
+        return hover;
+
+      hoverOS << "Results:\n";
+      for (auto it : llvm::zip(resultTupleTy.getElementNames(),
+                               resultTupleTy.getElementTypes())) {
+        StringRef name = std::get<0>(it);
+        hoverOS << "* " << (name.empty() ? "" : (name + ": ")) << "`"
+                << std::get<1>(it) << "`\n";
+      }
+    } else {
+      hoverOS << "Results:\n* `" << resultType << "`\n";
+    }
+    hoverOS << "***\n";
+  }
+  return hover;
+}
+
+//===----------------------------------------------------------------------===//
 // PDLTextFileChunk
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class represents a single chunk of an PDL text file.
 struct PDLTextFileChunk {
   PDLTextFileChunk(uint64_t lineOffset, const lsp::URIForFile &uri,
                    StringRef contents,
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);
+  Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
+                                 lsp::Position hoverPos);
 
 private:
   /// Find the PDL document that contains the given position, and update the
   /// position to be anchored at the start of the found chunk instead of the
   /// beginning of the file.
   PDLTextFileChunk &getChunkFor(lsp::Position &pos);
 
   /// The full string contents of the file.
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);
 }
 
+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);
+  return hoverInfo;
+}
+
 PDLTextFileChunk &PDLTextFile::getChunkFor(lsp::Position &pos) {
   if (chunks.size() == 1)
     return *chunks.front();
 
   // Search for the first chunk with a greater line offset, the previous chunk
   // is the one that contains `pos`.
   auto it = llvm::upper_bound(
       chunks, pos, [](const lsp::Position &pos, const auto &chunk) {

 
 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);
 }
+
+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/hover.test

+// RUN: mlir-pdll-lsp-server -lit-test < %s | FileCheck -strict-whitespace %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}"
+}}}
+// -----
+// 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:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "contents": {
+// CHECK-NEXT:      "kind": "markdown",
+// CHECK-NEXT:      "value": "**variable**: `root`\n***\nType: `Op`\n"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    "range": {
+// CHECK-NEXT:      "end": {
+// CHECK-NEXT:        "character": 8,
+// CHECK-NEXT:        "line": 3
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "start": {
+// CHECK-NEXT:        "character": 4,
+// CHECK-NEXT:        "line": 3
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+// Hover on a pattern.
+{"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
+  "textDocument":{"uri":"test:///foo.pdll"},
+  "position":{"line":2,"character":9}
+}}
+//      CHECK:  "id": 1,
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "contents": {
+// CHECK-NEXT:      "kind": "markdown",
+// CHECK-NEXT:      "value": "**Pattern**: `Foo`\n***\nRootOp: `Op`\n"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    "range": {
+// CHECK-NEXT:      "end": {
+// CHECK-NEXT:        "character": 11,
+// CHECK-NEXT:        "line": 2
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "start": {
+// CHECK-NEXT:        "character": 8,
+// CHECK-NEXT:        "line": 2
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+// Hover on a core constraint.
+{"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
+  "textDocument":{"uri":"test:///foo.pdll"},
+  "position":{"line":3,"character":11}
+}}
+//      CHECK:  "id": 1,
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "contents": {
+// CHECK-NEXT:      "kind": "markdown",
+// CHECK-NEXT:      "value": "**Constraint**: `Op`\n"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    "range": {
+// CHECK-NEXT:      "end": {
+// CHECK-NEXT:        "character": 12,
+// CHECK-NEXT:        "line": 3
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "start": {
+// CHECK-NEXT:        "character": 10,
+// CHECK-NEXT:        "line": 3
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+// Hover on a user constraint.
+{"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
+  "textDocument":{"uri":"test:///foo.pdll"},
+  "position":{"line":0,"character":14}
+}}
+//      CHECK:  "id": 1,
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "contents": {
+// CHECK-NEXT:      "kind": "markdown",
+// CHECK-NEXT:      "value": "**Constraint**: `FooCst`\n***\n"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    "range": {
+// CHECK-NEXT:      "end": {
+// CHECK-NEXT:        "character": 17,
+// CHECK-NEXT:        "line": 0
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "start": {
+// CHECK-NEXT:        "character": 11,
+// CHECK-NEXT:        "line": 0
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+// Hover on a user rewrite.
+{"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
+  "textDocument":{"uri":"test:///foo.pdll"},
+  "position":{"line":1,"character":11}
+}}
+//      CHECK:  "id": 1,
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "contents": {
+// CHECK-NEXT:      "kind": "markdown",
+// CHECK-NEXT:      "value": "**Rewrite**: `FooRewrite`\n***\nParameters:\n* op: `Op`\n***\nResults:\n* `Op`\n***\n"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    "range": {
+// CHECK-NEXT:      "end": {
+// CHECK-NEXT:        "character": 18,
+// CHECK-NEXT:        "line": 1
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "start": {
+// CHECK-NEXT:        "character": 8,
+// CHECK-NEXT:        "line": 1
+// 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 -strict-whitespace %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:      "definitionProvider": true,
+// CHECK-NEXT:      "hoverProvider": true,
 // CHECK-NEXT:      "referencesProvider": true,
 // CHECK-NEXT:      "textDocumentSync": {
 // CHECK-NEXT:        "change": 1,
 // CHECK-NEXT:        "openClose": true,
 // CHECK-NEXT:        "save": true
 // CHECK-NEXT:      }
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "serverInfo": {