Skip to content

Commit

Permalink
[mlir][PDLL] Add symbol support to the PDLL language server
Browse files Browse the repository at this point in the history 
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
  • Loading branch information
@River707
River707 committed Mar 19, 2022
1 parent 41ae211 commit 8dd4272
Show file tree
Hide file tree
Showing 5 changed files with 221 additions and 0 deletions.
21 changes: 21 additions & 0 deletions mlir/lib/Tools/mlir-pdll-lsp-server/LSPServer.cpp
View file
Expand Up @@ -58,6 +58,12 @@ struct LSPServer {
void onHover(const TextDocumentPositionParams &params,
Callback<Optional<Hover>> reply);

//===--------------------------------------------------------------------===//
// Document Symbols

void onDocumentSymbol(const DocumentSymbolParams &params,
Callback<std::vector<DocumentSymbol>> reply);

//===--------------------------------------------------------------------===//
// Fields
//===--------------------------------------------------------------------===//
Expand Down Expand Up @@ -91,6 +97,7 @@ void LSPServer::onInitialize(const InitializeParams &params,
{"definitionProvider", true},
{"referencesProvider", true},
{"hoverProvider", true},
{"documentSymbolProvider", true},
};

llvm::json::Object result{
Expand Down Expand Up @@ -169,6 +176,16 @@ void LSPServer::onHover(const TextDocumentPositionParams &params,
reply(server.findHover(params.textDocument.uri, params.position));
}

//===----------------------------------------------------------------------===//
// Document Symbols

void LSPServer::onDocumentSymbol(const DocumentSymbolParams &params,
Callback<std::vector<DocumentSymbol>> reply) {
std::vector<DocumentSymbol> symbols;
server.findDocumentSymbols(params.textDocument.uri, symbols);
reply(std::move(symbols));
}

//===----------------------------------------------------------------------===//
// Entry Point
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -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<PublishDiagnosticsParams>(
Expand Down
101 changes: 101 additions & 0 deletions mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp
View file
Expand Up @@ -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) {
Expand Down Expand Up @@ -265,6 +271,12 @@ struct PDLDocument {
const T *decl,
const SMRange &hoverRange);

//===--------------------------------------------------------------------===//
// Document Symbols
//===--------------------------------------------------------------------===//

void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);

//===--------------------------------------------------------------------===//
// Fields
//===--------------------------------------------------------------------===//
Expand Down Expand Up @@ -492,6 +504,48 @@ lsp::Hover PDLDocument::buildHoverForUserConstraintOrRewrite(
return hover;
}

//===----------------------------------------------------------------------===//
// PDLDocument: Document Symbols
//===----------------------------------------------------------------------===//

void PDLDocument::findDocumentSymbols(
std::vector<lsp::DocumentSymbol> &symbols) {
if (failed(astModule))
return;

for (const ast::Decl *decl : (*astModule)->getChildren()) {
if (!isMainFileLoc(sourceMgr, decl->getLoc()))
continue;

if (const auto *patternDecl = dyn_cast<ast::PatternDecl>(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() : "<pattern>", lsp::SymbolKind::Class,
lsp::Range(sourceMgr, bodyLoc), lsp::Range(sourceMgr, nameLoc));
} else if (const auto *cDecl = dyn_cast<ast::UserConstraintDecl>(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<ast::UserRewriteDecl>(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
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -545,6 +599,7 @@ class PDLTextFile {
std::vector<lsp::Location> &references);
Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
lsp::Position hoverPos);
void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);

private:
/// Find the PDL document that contains the given position, and update the
Expand Down Expand Up @@ -643,6 +698,45 @@ Optional<lsp::Hover> PDLTextFile::findHover(const lsp::URIForFile &uri,
return hoverInfo;
}

void PDLTextFile::findDocumentSymbols(
std::vector<lsp::DocumentSymbol> &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("<file-split-" + Twine(i) + ">",
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<lsp::DocumentSymbol *> 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();
Expand Down Expand Up @@ -714,3 +808,10 @@ Optional<lsp::Hover> lsp::PDLLServer::findHover(const URIForFile &uri,
return fileIt->second->findHover(uri, hoverPos);
return llvm::None;
}

void lsp::PDLLServer::findDocumentSymbols(
const URIForFile &uri, std::vector<DocumentSymbol> &symbols) {
auto fileIt = impl->files.find(uri.file());
if (fileIt != impl->files.end())
fileIt->second->findDocumentSymbols(symbols);
}
5 changes: 5 additions & 0 deletions mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h
View file
Expand Up @@ -15,6 +15,7 @@
namespace mlir {
namespace lsp {
struct Diagnostic;
struct DocumentSymbol;
struct Hover;
struct Location;
struct Position;
Expand Down Expand Up @@ -52,6 +53,10 @@ class PDLLServer {
/// couldn't be found.
Optional<Hover> findHover(const URIForFile &uri, const Position &hoverPos);

/// Find all of the document symbols within the given file.
void findDocumentSymbols(const URIForFile &uri,
std::vector<DocumentSymbol> &symbols);

private:
struct Impl;

Expand Down
93 changes: 93 additions & 0 deletions mlir/test/mlir-pdll-lsp-server/document-symbols.test
View file
@@ -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<foo.op>;\n}\nConstraint Cst() -> Op{\nreturn op<toy.test>;\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"}
1 change: 1 addition & 0 deletions mlir/test/mlir-pdll-lsp-server/initialize-params.test
View file
Expand Up @@ -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": {
Expand Down

0 comments on commit 8dd4272

Please sign in to comment.