[mlir:PDLL] Add support for inlay hints

These allow for displaying additional inline information,
such as the types of variables, names operands/results,
constraint/rewrite arguments, etc. This requires a bump in the
vscode extension to a newer version, as inlay hints are a new LSP feature.

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

mlir/docs/Tools/MLIRLSP.md

@@ -219,6 +219,18 @@ necessarily provided by all IDE clients.
 
 ![IMG](/mlir-pdll-lsp-server/view_output.gif)
 
+#### Inlay hints
+
+The language server provides additional information inline with the source code.
+Editors usually render this using read-only virtual text snippets interspersed
+with code. Hints may be shown for:
+
+* types of local variables
+* names of operand and result groups
+* constraint and rewrite arguments
+
+![IMG](/mlir-pdll-lsp-server/inlay_hints.png)
+
 ## TableGen LSP Language Server : `tblgen-lsp-server`
 
 MLIR provides an implementation of an LSP language server for `.td` text files

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

@@ -850,3 +850,45 @@ llvm::json::Value mlir::lsp::toJSON(const DocumentLink &value) {
       {"target", value.target},
   };
 }
+
+//===----------------------------------------------------------------------===//
+// InlayHintsParams
+//===----------------------------------------------------------------------===//
+
+bool mlir::lsp::fromJSON(const llvm::json::Value &value,
+                         InlayHintsParams &result, llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  return o && o.map("textDocument", result.textDocument) &&
+         o.map("range", result.range);
+}
+
+//===----------------------------------------------------------------------===//
+// InlayHint
+//===----------------------------------------------------------------------===//
+
+llvm::json::Value mlir::lsp::toJSON(const InlayHint &value) {
+  return llvm::json::Object{{"position", value.position},
+                            {"kind", (int)value.kind},
+                            {"label", value.label},
+                            {"paddingLeft", value.paddingLeft},
+                            {"paddingRight", value.paddingRight}};
+}
+bool mlir::lsp::operator==(const InlayHint &lhs, const InlayHint &rhs) {
+  return std::tie(lhs.position, lhs.kind, lhs.label) ==
+         std::tie(rhs.position, rhs.kind, rhs.label);
+}
+bool mlir::lsp::operator<(const InlayHint &lhs, const InlayHint &rhs) {
+  return std::tie(lhs.position, lhs.kind, lhs.label) <
+         std::tie(rhs.position, rhs.kind, rhs.label);
+}
+
+llvm::raw_ostream &mlir::lsp::operator<<(llvm::raw_ostream &os,
+                                         InlayHintKind value) {
+  switch (value) {
+  case InlayHintKind::Parameter:
+    return os << "parameter";
+  case InlayHintKind::Type:
+    return os << "type";
+  }
+  llvm_unreachable("Unknown InlayHintKind");
+}

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

@@ -1000,6 +1000,86 @@ struct DocumentLink {
 /// Add support for JSON serialization.
 llvm::json::Value toJSON(const DocumentLink &value);
 
+//===----------------------------------------------------------------------===//
+// InlayHintsParams
+//===----------------------------------------------------------------------===//
+
+/// A parameter literal used in inlay hint requests.
+struct InlayHintsParams {
+  /// The text document.
+  TextDocumentIdentifier textDocument;
+
+  /// The visible document range for which inlay hints should be computed.
+  Range range;
+};
+
+/// Add support for JSON serialization.
+bool fromJSON(const llvm::json::Value &value, InlayHintsParams &result,
+              llvm::json::Path path);
+
+//===----------------------------------------------------------------------===//
+// InlayHintKind
+//===----------------------------------------------------------------------===//
+
+/// Inlay hint kinds.
+enum class InlayHintKind {
+  /// An inlay hint that for a type annotation.
+  ///
+  /// An example of a type hint is a hint in this position:
+  ///    auto var ^ = expr;
+  /// which shows the deduced type of the variable.
+  Type = 1,
+
+  /// An inlay hint that is for a parameter.
+  ///
+  /// An example of a parameter hint is a hint in this position:
+  ///    func(^arg);
+  /// which shows the name of the corresponding parameter.
+  Parameter = 2,
+};
+
+//===----------------------------------------------------------------------===//
+// InlayHint
+//===----------------------------------------------------------------------===//
+
+/// Inlay hint information.
+struct InlayHint {
+  InlayHint(InlayHintKind kind, Position pos) : position(pos), kind(kind) {}
+
+  /// The position of this hint.
+  Position position;
+
+  /// The label of this hint. A human readable string or an array of
+  /// InlayHintLabelPart label parts.
+  ///
+  /// *Note* that neither the string nor the label part can be empty.
+  std::string label;
+
+  /// The kind of this hint. Can be omitted in which case the client should fall
+  /// back to a reasonable default.
+  InlayHintKind kind;
+
+  /// Render padding before the hint.
+  ///
+  /// Note: Padding should use the editor's background color, not the
+  /// background color of the hint itself. That means padding can be used
+  /// to visually align/separate an inlay hint.
+  bool paddingLeft = false;
+
+  /// Render padding after the hint.
+  ///
+  /// Note: Padding should use the editor's background color, not the
+  /// background color of the hint itself. That means padding can be used
+  /// to visually align/separate an inlay hint.
+  bool paddingRight = false;
+};
+
+/// Add support for JSON serialization.
+llvm::json::Value toJSON(const InlayHint &);
+bool operator==(const InlayHint &lhs, const InlayHint &rhs);
+bool operator<(const InlayHint &lhs, const InlayHint &rhs);
+llvm::raw_ostream &operator<<(llvm::raw_ostream &os, InlayHintKind value);
+
 } // namespace lsp
 } // namespace mlir
 

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

@@ -82,6 +82,12 @@ struct LSPServer {
   void onSignatureHelp(const TextDocumentPositionParams &params,
                        Callback<SignatureHelp> reply);
 
+  //===--------------------------------------------------------------------===//
+  // Inlay Hints
+
+  void onInlayHint(const InlayHintsParams &params,
+                   Callback<std::vector<InlayHint>> reply);
+
   //===--------------------------------------------------------------------===//
   // PDLL View Output
 
@@ -140,6 +146,7 @@ void LSPServer::onInitialize(const InitializeParams &params,
        }},
       {"hoverProvider", true},
       {"documentSymbolProvider", true},
+      {"inlayHintProvider", true},
   };
 
   llvm::json::Object result{
@@ -248,6 +255,16 @@ void LSPServer::onSignatureHelp(const TextDocumentPositionParams &params,
   reply(server.getSignatureHelp(params.textDocument.uri, params.position));
 }
 
+//===----------------------------------------------------------------------===//
+// Inlay Hints
+
+void LSPServer::onInlayHint(const InlayHintsParams &params,
+                            Callback<std::vector<InlayHint>> reply) {
+  std::vector<InlayHint> hints;
+  server.getInlayHints(params.textDocument.uri, params.range, hints);
+  reply(std::move(hints));
+}
+
 //===----------------------------------------------------------------------===//
 // PDLL ViewOutput
 
@@ -305,6 +322,10 @@ LogicalResult mlir::lsp::runPdllLSPServer(PDLLServer &server,
   messageHandler.method("textDocument/signatureHelp", &lspServer,
                         &LSPServer::onSignatureHelp);
 
+  // Inlay Hints
+  messageHandler.method("textDocument/inlayHint", &lspServer,
+                        &LSPServer::onInlayHint);
+
   // PDLL ViewOutput
   messageHandler.method("pdll/viewOutput", &lspServer,
                         &LSPServer::onPDLLViewOutput);