From 8dd4272ca2f5ce7c793f3c6f9e468913d768700d Mon Sep 17 00:00:00 2001 From: River Riddle Date: Fri, 11 Mar 2022 00:23:39 -0800 Subject: [PATCH] [mlir][PDLL] Add symbol support to the PDLL language server This adds support for documenting the top-level "symbols", e.g. patterns, constraints, rewrites, etc., within a PDLL file. Differential Revision: https://reviews.llvm.org/D121543 --- .../Tools/mlir-pdll-lsp-server/LSPServer.cpp | 21 ++++ .../Tools/mlir-pdll-lsp-server/PDLLServer.cpp | 101 ++++++++++++++++++ .../Tools/mlir-pdll-lsp-server/PDLLServer.h | 5 + .../document-symbols.test | 93 ++++++++++++++++ .../initialize-params.test | 1 + 5 files changed, 221 insertions(+) create mode 100644 mlir/test/mlir-pdll-lsp-server/document-symbols.test diff --git a/mlir/lib/Tools/mlir-pdll-lsp-server/LSPServer.cpp b/mlir/lib/Tools/mlir-pdll-lsp-server/LSPServer.cpp index c1e198f015ea8..214e346ed0128 100644 --- a/mlir/lib/Tools/mlir-pdll-lsp-server/LSPServer.cpp +++ b/mlir/lib/Tools/mlir-pdll-lsp-server/LSPServer.cpp @@ -58,6 +58,12 @@ struct LSPServer { void onHover(const TextDocumentPositionParams ¶ms, Callback> reply); + //===--------------------------------------------------------------------===// + // Document Symbols + + void onDocumentSymbol(const DocumentSymbolParams ¶ms, + Callback> reply); + //===--------------------------------------------------------------------===// // Fields //===--------------------------------------------------------------------===// @@ -91,6 +97,7 @@ void LSPServer::onInitialize(const InitializeParams ¶ms, {"definitionProvider", true}, {"referencesProvider", true}, {"hoverProvider", true}, + {"documentSymbolProvider", true}, }; llvm::json::Object result{ @@ -169,6 +176,16 @@ void LSPServer::onHover(const TextDocumentPositionParams ¶ms, reply(server.findHover(params.textDocument.uri, params.position)); } +//===----------------------------------------------------------------------===// +// Document Symbols + +void LSPServer::onDocumentSymbol(const DocumentSymbolParams ¶ms, + Callback> reply) { + std::vector symbols; + server.findDocumentSymbols(params.textDocument.uri, symbols); + reply(std::move(symbols)); +} + //===----------------------------------------------------------------------===// // Entry Point //===----------------------------------------------------------------------===// @@ -201,6 +218,10 @@ LogicalResult mlir::lsp::runPdllLSPServer(PDLLServer &server, // Hover messageHandler.method("textDocument/hover", &lspServer, &LSPServer::onHover); + // Document Symbols + messageHandler.method("textDocument/documentSymbol", &lspServer, + &LSPServer::onDocumentSymbol); + // Diagnostics lspServer.publishDiagnostics = messageHandler.outgoingNotification( diff --git a/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp b/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp index d0960f9c1f2f9..cc8c22f49ab23 100644 --- a/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp +++ b/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp @@ -42,6 +42,12 @@ static lsp::URIForFile getURIFromLoc(llvm::SourceMgr &mgr, SMRange loc, return mainFileURI; } +/// Returns true if the given location is in the main file of the source +/// manager. +static bool isMainFileLoc(llvm::SourceMgr &mgr, SMRange loc) { + return mgr.FindBufferContainingLoc(loc.Start) == mgr.getMainFileID(); +} + /// Returns a language server location from the given source range. static lsp::Location getLocationFromLoc(llvm::SourceMgr &mgr, SMRange range, const lsp::URIForFile &uri) { @@ -265,6 +271,12 @@ struct PDLDocument { const T *decl, const SMRange &hoverRange); + //===--------------------------------------------------------------------===// + // Document Symbols + //===--------------------------------------------------------------------===// + + void findDocumentSymbols(std::vector &symbols); + //===--------------------------------------------------------------------===// // Fields //===--------------------------------------------------------------------===// @@ -492,6 +504,48 @@ lsp::Hover PDLDocument::buildHoverForUserConstraintOrRewrite( return hover; } +//===----------------------------------------------------------------------===// +// PDLDocument: Document Symbols +//===----------------------------------------------------------------------===// + +void PDLDocument::findDocumentSymbols( + std::vector &symbols) { + if (failed(astModule)) + return; + + for (const ast::Decl *decl : (*astModule)->getChildren()) { + if (!isMainFileLoc(sourceMgr, decl->getLoc())) + continue; + + if (const auto *patternDecl = dyn_cast(decl)) { + const ast::Name *name = patternDecl->getName(); + + SMRange nameLoc = name ? name->getLoc() : patternDecl->getLoc(); + SMRange bodyLoc(nameLoc.Start, patternDecl->getBody()->getLoc().End); + + symbols.emplace_back( + name ? name->getName() : "", lsp::SymbolKind::Class, + lsp::Range(sourceMgr, bodyLoc), lsp::Range(sourceMgr, nameLoc)); + } else if (const auto *cDecl = dyn_cast(decl)) { + // TODO: Add source information for the code block body. + SMRange nameLoc = cDecl->getName().getLoc(); + SMRange bodyLoc = nameLoc; + + symbols.emplace_back( + cDecl->getName().getName(), lsp::SymbolKind::Function, + lsp::Range(sourceMgr, bodyLoc), lsp::Range(sourceMgr, nameLoc)); + } else if (const auto *cDecl = dyn_cast(decl)) { + // TODO: Add source information for the code block body. + SMRange nameLoc = cDecl->getName().getLoc(); + SMRange bodyLoc = nameLoc; + + symbols.emplace_back( + cDecl->getName().getName(), lsp::SymbolKind::Function, + lsp::Range(sourceMgr, bodyLoc), lsp::Range(sourceMgr, nameLoc)); + } + } +} + //===----------------------------------------------------------------------===// // PDLTextFileChunk //===----------------------------------------------------------------------===// @@ -545,6 +599,7 @@ class PDLTextFile { std::vector &references); Optional findHover(const lsp::URIForFile &uri, lsp::Position hoverPos); + void findDocumentSymbols(std::vector &symbols); private: /// Find the PDL document that contains the given position, and update the @@ -643,6 +698,45 @@ Optional PDLTextFile::findHover(const lsp::URIForFile &uri, return hoverInfo; } +void PDLTextFile::findDocumentSymbols( + std::vector &symbols) { + if (chunks.size() == 1) + return chunks.front()->document.findDocumentSymbols(symbols); + + // If there are multiple chunks in this file, we create top-level symbols for + // each chunk. + for (unsigned i = 0, e = chunks.size(); i < e; ++i) { + PDLTextFileChunk &chunk = *chunks[i]; + lsp::Position startPos(chunk.lineOffset); + lsp::Position endPos((i == e - 1) ? totalNumLines - 1 + : chunks[i + 1]->lineOffset); + lsp::DocumentSymbol symbol("", + lsp::SymbolKind::Namespace, + /*range=*/lsp::Range(startPos, endPos), + /*selectionRange=*/lsp::Range(startPos)); + chunk.document.findDocumentSymbols(symbol.children); + + // Fixup the locations of document symbols within this chunk. + if (i != 0) { + SmallVector symbolsToFix; + for (lsp::DocumentSymbol &childSymbol : symbol.children) + symbolsToFix.push_back(&childSymbol); + + while (!symbolsToFix.empty()) { + lsp::DocumentSymbol *symbol = symbolsToFix.pop_back_val(); + chunk.adjustLocForChunkOffset(symbol->range); + chunk.adjustLocForChunkOffset(symbol->selectionRange); + + for (lsp::DocumentSymbol &childSymbol : symbol->children) + symbolsToFix.push_back(&childSymbol); + } + } + + // Push the symbol for this chunk. + symbols.emplace_back(std::move(symbol)); + } +} + PDLTextFileChunk &PDLTextFile::getChunkFor(lsp::Position &pos) { if (chunks.size() == 1) return *chunks.front(); @@ -714,3 +808,10 @@ Optional lsp::PDLLServer::findHover(const URIForFile &uri, return fileIt->second->findHover(uri, hoverPos); return llvm::None; } + +void lsp::PDLLServer::findDocumentSymbols( + const URIForFile &uri, std::vector &symbols) { + auto fileIt = impl->files.find(uri.file()); + if (fileIt != impl->files.end()) + fileIt->second->findDocumentSymbols(symbols); +} diff --git a/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h b/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h index 9ecc5ecda8fa1..1a647f18db125 100644 --- a/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h +++ b/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h @@ -15,6 +15,7 @@ namespace mlir { namespace lsp { struct Diagnostic; +struct DocumentSymbol; struct Hover; struct Location; struct Position; @@ -52,6 +53,10 @@ class PDLLServer { /// couldn't be found. Optional findHover(const URIForFile &uri, const Position &hoverPos); + /// Find all of the document symbols within the given file. + void findDocumentSymbols(const URIForFile &uri, + std::vector &symbols); + private: struct Impl; diff --git a/mlir/test/mlir-pdll-lsp-server/document-symbols.test b/mlir/test/mlir-pdll-lsp-server/document-symbols.test new file mode 100644 index 0000000000000..a62d5c8428685 --- /dev/null +++ b/mlir/test/mlir-pdll-lsp-server/document-symbols.test @@ -0,0 +1,93 @@ +// RUN: mlir-pdll-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s +{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootUri":"test:///workspace","capabilities":{"textDocument":{"documentSymbol":{"hierarchicalDocumentSymbolSupport":true}}},"trace":"off"}} +// ----- +{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{ + "uri":"test:///foo.pdll", + "languageId":"pdll", + "version":1, + "text":"Pattern Foo {\nerase op;\n}\nConstraint Cst() -> Op{\nreturn op;\n}\n\nRewrite SomeRewrite() -> Op {\nreturn op: Op;\n}" +}}} +// ----- +{"jsonrpc":"2.0","id":1,"method":"textDocument/documentSymbol","params":{ + "textDocument":{"uri":"test:///foo.pdll"} +}} +// CHECK: "id": 1 +// CHECK-NEXT: "jsonrpc": "2.0", +// CHECK-NEXT: "result": [ +// CHECK-NEXT: { +// CHECK-NEXT: "kind": 5, +// CHECK-NEXT: "name": "Foo", +// CHECK-NEXT: "range": { +// CHECK-NEXT: "end": { +// CHECK-NEXT: "character": 1, +// CHECK-NEXT: "line": 2 +// CHECK-NEXT: }, +// CHECK-NEXT: "start": { +// CHECK-NEXT: "character": 8, +// CHECK-NEXT: "line": 0 +// CHECK-NEXT: } +// CHECK-NEXT: }, +// CHECK-NEXT: "selectionRange": { +// CHECK-NEXT: "end": { +// CHECK-NEXT: "character": 11, +// CHECK-NEXT: "line": 0 +// CHECK-NEXT: }, +// CHECK-NEXT: "start": { +// CHECK-NEXT: "character": 8, +// CHECK-NEXT: "line": 0 +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: }, +// CHECK-NEXT: { +// CHECK-NEXT: "kind": 12, +// CHECK-NEXT: "name": "Cst", +// CHECK-NEXT: "range": { +// CHECK-NEXT: "end": { +// CHECK-NEXT: "character": 14, +// CHECK-NEXT: "line": 3 +// CHECK-NEXT: }, +// CHECK-NEXT: "start": { +// CHECK-NEXT: "character": 11, +// CHECK-NEXT: "line": 3 +// CHECK-NEXT: } +// CHECK-NEXT: }, +// CHECK-NEXT: "selectionRange": { +// CHECK-NEXT: "end": { +// CHECK-NEXT: "character": 14, +// CHECK-NEXT: "line": 3 +// CHECK-NEXT: }, +// CHECK-NEXT: "start": { +// CHECK-NEXT: "character": 11, +// CHECK-NEXT: "line": 3 +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: }, +// CHECK-NEXT: { +// CHECK-NEXT: "kind": 12, +// CHECK-NEXT: "name": "SomeRewrite", +// CHECK-NEXT: "range": { +// CHECK-NEXT: "end": { +// CHECK-NEXT: "character": 19, +// CHECK-NEXT: "line": 7 +// CHECK-NEXT: }, +// CHECK-NEXT: "start": { +// CHECK-NEXT: "character": 8, +// CHECK-NEXT: "line": 7 +// CHECK-NEXT: } +// CHECK-NEXT: }, +// CHECK-NEXT: "selectionRange": { +// CHECK-NEXT: "end": { +// CHECK-NEXT: "character": 19, +// CHECK-NEXT: "line": 7 +// CHECK-NEXT: }, +// CHECK-NEXT: "start": { +// CHECK-NEXT: "character": 8, +// CHECK-NEXT: "line": 7 +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: ] +// ----- +{"jsonrpc":"2.0","id":3,"method":"shutdown"} +// ----- +{"jsonrpc":"2.0","method":"exit"} diff --git a/mlir/test/mlir-pdll-lsp-server/initialize-params.test b/mlir/test/mlir-pdll-lsp-server/initialize-params.test index e889750591dda..d2af6f514fe54 100644 --- a/mlir/test/mlir-pdll-lsp-server/initialize-params.test +++ b/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: "documentSymbolProvider": true, // CHECK-NEXT: "hoverProvider": true, // CHECK-NEXT: "referencesProvider": true, // CHECK-NEXT: "textDocumentSync": {