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

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

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

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);
 
+  //===--------------------------------------------------------------------===//
+  // Hover
+
+  void onHover(const TextDocumentPositionParams &params,
+               Callback<Optional<Hover>> reply);
+
   //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
@@ -84,6 +90,7 @@ void LSPServer::onInitialize(const InitializeParams &params,
        }},
       {"definitionProvider", true},
       {"referencesProvider", true},
+      {"hoverProvider", true},
   };
 
   llvm::json::Object result{
@@ -154,6 +161,14 @@ void LSPServer::onReference(const ReferenceParams &params,
   reply(std::move(locations));
 }
 
+//===----------------------------------------------------------------------===//
+// Hover
+
+void LSPServer::onHover(const TextDocumentPositionParams &params,
+                        Callback<Optional<Hover>> reply) {
+  reply(server.findHover(params.textDocument.uri, params.position));
+}
+
 //===----------------------------------------------------------------------===//
 // Entry Point
 //===----------------------------------------------------------------------===//
@@ -183,6 +198,9 @@ LogicalResult mlir::lsp::runPdllLSPServer(PDLLServer &server,
   messageHandler.method("textDocument/references", &lspServer,
                         &LSPServer::onReference);
 
+  // Hover
+  messageHandler.method("textDocument/hover", &lspServer, &LSPServer::onHover);
+
   // Diagnostics
   lspServer.publishDiagnostics =
       messageHandler.outgoingNotification<PublishDiagnosticsParams>(

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

@@ -244,6 +244,27 @@ struct PDLDocument {
   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 SMRange &hoverRange);
+  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);
+  template <typename T>
+  lsp::Hover buildHoverForUserConstraintOrRewrite(StringRef typeName,
+                                                  const T *decl,
+                                                  const SMRange &hoverRange);
+
   //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
@@ -320,6 +341,157 @@ void PDLDocument::findReferencesOf(const lsp::URIForFile &uri,
     references.push_back(getLocationFromLoc(sourceMgr, refLoc, uri));
 }
 
+//===----------------------------------------------------------------------===//
+// PDLDocument: Hover
+//===----------------------------------------------------------------------===//
+
+Optional<lsp::Hover> PDLDocument::findHover(const lsp::URIForFile &uri,
+                                            const lsp::Position &hoverPos) {
+  SMLoc posLoc = hoverPos.getAsSMLoc(sourceMgr);
+  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 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 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();
+  }
+  return hover;
+}
+
+lsp::Hover PDLDocument::buildHoverForVariable(const ast::VariableDecl *varDecl,
+                                              const SMRange &hoverRange) {
+  lsp::Hover hover(lsp::Range(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 SMRange &hoverRange) {
+  lsp::Hover hover(lsp::Range(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 SMRange &hoverRange) {
+  lsp::Hover hover(lsp::Range(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 SMRange &hoverRange) {
+  lsp::Hover hover(lsp::Range(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
 //===----------------------------------------------------------------------===//
@@ -371,6 +543,8 @@ public:
                       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
@@ -458,6 +632,17 @@ void PDLTextFile::findReferencesOf(const lsp::URIForFile &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();
@@ -521,3 +706,11 @@ void lsp::PDLLServer::findReferencesOf(const URIForFile &uri,
   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/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h

@@ -15,6 +15,7 @@
 namespace mlir {
 namespace lsp {
 struct Diagnostic;
+struct Hover;
 struct Location;
 struct Position;
 class URIForFile;
@@ -47,6 +48,10 @@ public:
   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;
 

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

@@ -0,0 +1,133 @@
+// RUN: mlir-pdll-lsp-server -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}"
+}}}
+// -----
+// 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

@@ -6,6 +6,7 @@
 // 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,