[Mlir-commits] [mlir] 8dd4272 - [mlir][PDLL] Add symbol support to the PDLL language server

River Riddle llvmlistbot at llvm.org
Sat Mar 19 13:29:32 PDT 2022


Author: River Riddle
Date: 2022-03-19T13:28:23-07:00
New Revision: 8dd4272ca2f5ce7c793f3c6f9e468913d768700d

URL: https://github.com/llvm/llvm-project/commit/8dd4272ca2f5ce7c793f3c6f9e468913d768700d
DIFF: https://github.com/llvm/llvm-project/commit/8dd4272ca2f5ce7c793f3c6f9e468913d768700d.diff

LOG: [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

Added: 
    mlir/test/mlir-pdll-lsp-server/document-symbols.test

Modified: 
    mlir/lib/Tools/mlir-pdll-lsp-server/LSPServer.cpp
    mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp
    mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h
    mlir/test/mlir-pdll-lsp-server/initialize-params.test

Removed: 
    


################################################################################
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 &params,
                Callback<Optional<Hover>> reply);
 
+  //===--------------------------------------------------------------------===//
+  // Document Symbols
+
+  void onDocumentSymbol(const DocumentSymbolParams &params,
+                        Callback<std::vector<DocumentSymbol>> reply);
+
   //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
@@ -91,6 +97,7 @@ void LSPServer::onInitialize(const InitializeParams &params,
       {"definitionProvider", true},
       {"referencesProvider", true},
       {"hoverProvider", true},
+      {"documentSymbolProvider", true},
   };
 
   llvm::json::Object result{
@@ -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
 //===----------------------------------------------------------------------===//
@@ -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>(

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<lsp::DocumentSymbol> &symbols);
+
   //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
@@ -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
 //===----------------------------------------------------------------------===//
@@ -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
@@ -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();
@@ -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);
+}

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<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;
 

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<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"}

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": {


        


More information about the Mlir-commits mailing list