[mlir-lsp-server] Add support for hover on region operations

[rriddle]

Summary

This revision adds support for hover on region operations, by temporarily removing the regions during printing. This revision also tweaks the hover format for operations to include symbol information, now that FuncOp can be shown in the hover.

[jpienaar]

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

117

Would have been nice if we could have { ... } here instead to be able to differentiate from an empty body or not. But that would require more invasive changes. ... Mmm, well we could add a op with a custom printer (which doesn't parse back :)) and then this region is just the "ElidedRegion" with the elided op.

[rriddle]

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

117

Yeah, I think if we end up with a desire to support this in general we should add a printer flag that elides the region.

mlir/include/mlir/IR/SymbolTable.h

private:
 
   /// This is a mapping from a name to the symbol with that name.
   llvm::StringMap<Operation *> symbolTable;
 
   /// This is used when name conflicts are detected.
   unsigned uniquingCounter = 0;
 };
 
+raw_ostream &operator<<(raw_ostream &os, SymbolTable::Visibility visibility);
+
 //===----------------------------------------------------------------------===//
 // SymbolTableCollection
 //===----------------------------------------------------------------------===//
 
 /// This class represents a collection of `SymbolTable`s. This simplifies
 /// certain algorithms that run recursively on nested symbol tables. Symbol
 /// tables are constructed lazily to reduce the upfront cost of constructing
 /// unnecessary tables.

mlir/lib/IR/SymbolTable.cpp

Operation *SymbolTable::lookupNearestSymbolFrom(Operation *from,
   return symbolTableOp ? lookupSymbolIn(symbolTableOp, symbol) : nullptr;
 }
 Operation *SymbolTable::lookupNearestSymbolFrom(Operation *from,
                                                 SymbolRefAttr symbol) {
   Operation *symbolTableOp = getNearestSymbolTable(from);
   return symbolTableOp ? lookupSymbolIn(symbolTableOp, symbol) : nullptr;
 }
 
+raw_ostream &mlir::operator<<(raw_ostream &os,
+                              SymbolTable::Visibility visibility) {
+  switch (visibility) {
+  case SymbolTable::Visibility::Public:
+    return os << "public";
+  case SymbolTable::Visibility::Private:
+    return os << "private";
+  case SymbolTable::Visibility::Nested:
+    return os << "nested";
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // SymbolTable Trait Types
 //===----------------------------------------------------------------------===//
 
 LogicalResult detail::verifySymbolTable(Operation *op) {
   if (op->getNumRegions() != 1)
     return op->emitOpError()
            << "Operations with a 'SymbolTable' must have exactly one region";

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

for (const auto &arg : llvm::enumerate(block.arguments)) {
           hoverRange, block.block->getArgument(arg.index()), block);
     }
   }
   return llvm::None;
 }
 
 Optional<lsp::Hover> MLIRDocument::buildHoverForOperation(
     const AsmParserState::OperationDefinition &op) {
-  // Don't show hovers for operations with regions to avoid huge hover  blocks.
-  // TODO: Should we add support for printing an op without its regions?
-  if (llvm::any_of(op.op->getRegions(),
-                   [](Region &region) { return !region.empty(); }))
-    return llvm::None;
-
   lsp::Hover hover(getRangeFromLoc(sourceMgr, op.loc));
   llvm::raw_string_ostream os(hover.contents.value);
 
-  // For hovers on an operation, show the generic form.
-  os << "```mlir\n";
+  // Add the operation name to the hover.
+  os << "\"" << op.op->getName() << "\"";
+  if (SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(op.op))
+    os << " : " << symbol.getVisibility() << " @" << symbol.getName() << "";
+  os << "\n\n";
+
+  os << "Generic Form:\n\n```mlir\n";
+
+  // Temporary drop the regions of this operation so that they don't get
+  // printed in the output. This helps keeps the size of the output hover
+  // small.
+  SmallVector<std::unique_ptr<Region>> regions;
+  for (Region &region : op.op->getRegions()) {
+    regions.emplace_back(std::make_unique<Region>());
+    regions.back()->takeBody(region);
+  }
+
   op.op->print(
       os, OpPrintingFlags().printGenericOpForm().elideLargeElementsAttrs());
   os << "\n```\n";
 
+  // Move the regions back to the current operation.
+  for (Region &region : op.op->getRegions())
+    region.takeBody(*regions.back());
+
   return hover;
 }
 
 lsp::Hover MLIRDocument::buildHoverForOperationResult(llvm::SMRange hoverRange,
                                                       Operation *op,
                                                       unsigned resultStart,
                                                       unsigned resultEnd,
                                                       llvm::SMLoc posLoc) {

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

// Hover on an operation.
   "textDocument":{"uri":"test:///foo.mlir"},
   "position":{"line":1,"character":12}
 }}
 //      CHECK:  "id": 1,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
-// CHECK-NEXT:      "value": "```mlir\n%true = \"std.constant\"() {value = true} : () -> i1\n```\n"
+// CHECK-NEXT:      "value": "\"std.constant\"\n\nGeneric Form:\n\n```mlir\n%true = \"std.constant\"() {value = true} : () -> i1\n```\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 17,
 // CHECK-NEXT:        "line": 1
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 9,
 // CHECK-NEXT:        "line": 1
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on an operation result.
-{"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
+{"jsonrpc":"2.0","id":2,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
   "position":{"line":1,"character":2}
 }}
-//      CHECK:  "id": 1,
+//      CHECK:  "id": 2,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "Operation: \"std.constant\"\n\nResult #0\n\nType: `i1`\n\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 6,
 // CHECK-NEXT:        "line": 1
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 0,
 // CHECK-NEXT:        "line": 1
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on a Block.
-{"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
+{"jsonrpc":"2.0","id":3,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
   "position":{"line":3,"character":2}
 }}
-//      CHECK:  "id": 1,
+//      CHECK:  "id": 3,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "Operation: \"func\"\n\nBlock #1\n\nPredecessors: <Block #0>\n\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 4,
 // CHECK-NEXT:        "line": 3
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 0,
 // CHECK-NEXT:        "line": 3
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on a Block argument.
-{"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
+{"jsonrpc":"2.0","id":4,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
   "position":{"line":0,"character":12}
 }}
-//      CHECK:  "id": 1,
+//      CHECK:  "id": 4,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "Operation: \"func\"\n\nBlock: <Block #0>\n\nArgument #0\n\nType: `i1`\n\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 14,
 // CHECK-NEXT:        "line": 0
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 10,
 // CHECK-NEXT:        "line": 0
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
-{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+// Hover on a region operation.
+{"jsonrpc":"2.0","id":5,"method":"textDocument/hover","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":0,"character":1}
+}}
+//      CHECK:  "id": 5,
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "contents": {
+// CHECK-NEXT:      "kind": "markdown",
+// CHECK-NEXT:      "value": "\"func\" : public @foo\n\nGeneric Form:\n\n```mlir\n\"func\"() ( {\n}) {sym_name = \"foo\", type = (i1) -> ()} : () -> ()\n```\n"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    "range": {
+// CHECK-NEXT:      "end": {
+// CHECK-NEXT:        "character": 4,
+// CHECK-NEXT:        "line": 0
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "start": {
+// CHECK-NEXT:        "character": 0,
+// CHECK-NEXT:        "line": 0
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+{"jsonrpc":"2.0","id":6,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}