[mlir][PDLL] Add symbol support to the PDLL language server

[rriddle]

Summary

This adds support for documenting the top-level "symbols",
e.g. patterns, constraints, rewrites, etc., within a PDLL file.

[jpienaar]

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

100

Why/What is the consequence of this? And if we restricted to just top-level for symbol would this be useful for where hierarchical is not supported?

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

551

Could you expand here how this relates to the below assignment? (I don't quite follow how using the nameLoc as the bodyLoc helps here or changes behavior, I'm guessing that it now points to place of template definition)

[rriddle]

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

100

It changes how we need to construct the symbol set, and I don't have access to an IDE that doesn't support hierarchical to see/decide what is actually useful there. For now though, PDLL doesn't have nested symbols so I just removed the optionality here (this was copied from the MLIR server).

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

551

Thanks for calling this out, this isn't necessary at this point. It's part of template support, which is a later addition (and would actually make sense when seeing the template code).

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

struct LSPServer::Impl {
 
   //===--------------------------------------------------------------------===//
   // Hover
 
   void onHover(const TextDocumentPositionParams &params,
                Callback<Optional<Hover>> reply);
 
   //===--------------------------------------------------------------------===//
+  // Document Symbols
+
+  void onDocumentSymbol(const DocumentSymbolParams &params,
+                        Callback<std::vector<DocumentSymbol>> reply);
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   PDLLServer &server;
   JSONTransport &transport;
 
   /// An outgoing notification used to send diagnostics to the client when they
   /// are ready to be processed.
llvm::json::Object serverCaps{
        llvm::json::Object{
            {"openClose", true},
            {"change", (int)TextDocumentSyncKind::Full},
            {"save", true},
        }},
       {"definitionProvider", true},
       {"referencesProvider", true},
       {"hoverProvider", true},
+
+      // For now we only support documenting symbols when the client supports
+      // hierarchical symbols.
+      {"documentSymbolProvider",
+       params.capabilities.hierarchicalDocumentSymbol},
   };
 
   llvm::json::Object result{
       {{"serverInfo", llvm::json::Object{{"name", "mlir-pdll-lsp-server"},
                                          {"version", "0.0.0"}}},
        {"capabilities", std::move(serverCaps)}}};
   reply(std::move(result));
 }

 // Hover
 
 void LSPServer::Impl::onHover(const TextDocumentPositionParams &params,
                               Callback<Optional<Hover>> reply) {
   reply(server.findHover(params.textDocument.uri, params.position));
 }
 
 //===----------------------------------------------------------------------===//
+// Document Symbols
+
+void LSPServer::Impl::onDocumentSymbol(
+    const DocumentSymbolParams &params,
+    Callback<std::vector<DocumentSymbol>> reply) {
+  std::vector<DocumentSymbol> symbols;
+  server.findDocumentSymbols(params.textDocument.uri, symbols);
+  reply(std::move(symbols));
+}
+
+//===----------------------------------------------------------------------===//
 // LSPServer
 //===----------------------------------------------------------------------===//
 
 LSPServer::LSPServer(PDLLServer &server, JSONTransport &transport)
     : impl(std::make_unique<Impl>(server, transport)) {}
 LSPServer::~LSPServer() = default;
 
 LogicalResult LSPServer::run() {
LogicalResult LSPServer::run() {
   messageHandler.method("textDocument/definition", impl.get(),
                         &Impl::onGoToDefinition);
   messageHandler.method("textDocument/references", impl.get(),
                         &Impl::onReference);
 
   // Hover
   messageHandler.method("textDocument/hover", impl.get(), &Impl::onHover);
 
+  // Document Symbols
+  messageHandler.method("textDocument/documentSymbol", impl.get(),
+                        &Impl::onDocumentSymbol);
+
   // Diagnostics
   impl->publishDiagnostics =
       messageHandler.outgoingNotification<PublishDiagnosticsParams>(
           "textDocument/publishDiagnostics");
 
   // Run the main loop of the transport.
   if (llvm::Error error = impl->transport.run(messageHandler)) {
     Logger::error("Transport error: {0}", error);
     llvm::consumeError(std::move(error));
     return failure();
   }
   return success(impl->shutdownRequestReceived);
 }

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h

 #define LIB_MLIR_TOOLS_MLIRPDLLSPSERVER_SERVER_H_
 
 #include "mlir/Support/LLVM.h"
 #include <memory>
 
 namespace mlir {
 namespace lsp {
 struct Diagnostic;
+struct DocumentSymbol;
 struct Hover;
 struct Location;
 struct Position;
 class URIForFile;
 
 /// This class implements all of the PDLL related functionality necessary for a
 /// language server. This class allows for keeping the PDLL specific logic
 /// separate from the logic that involves LSP server/client communication.
public:
   /// Find all references of the object pointed at by the given position.
   void findReferencesOf(const URIForFile &uri, const Position &pos,
                         std::vector<Location> &references);
 
   /// Find a hover description for the given hover position, or None if one
   /// 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;
 
   std::unique_ptr<Impl> impl;
 };
 
 } // namespace lsp
 } // namespace mlir
 
 #endif // LIB_MLIR_TOOLS_MLIRPDLLSPSERVER_SERVER_H_

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

llvm::Expected<lsp::URIForFile> fileForLoc = lsp::URIForFile::fromFile(
       mgr.getBufferInfo(bufferId).Buffer->getBufferIdentifier());
   if (fileForLoc)
     return *fileForLoc;
   lsp::Logger::error("Failed to create URI for include file: {0}",
                      llvm::toString(fileForLoc.takeError()));
   return mainFileURI;
 }
 
+/// Returns true if the given location is in the main file of the source
+/// manager.
+static bool isMainFileLoc(llvm::SourceMgr &mgr, llvm::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,
                                         llvm::SMRange range,
                                         const lsp::URIForFile &uri) {
   return lsp::Location{getURIFromLoc(mgr, range, uri),
                        getRangeFromLoc(mgr, range)};
 }
 
struct PDLDocument {
   lsp::Hover buildHoverForCoreConstraint(const ast::CoreConstraintDecl *decl,
                                          const llvm::SMRange &hoverRange);
   template <typename T>
   lsp::Hover
   buildHoverForUserConstraintOrRewrite(StringRef typeName, const T *decl,
                                        const llvm::SMRange &hoverRange);
 
   //===--------------------------------------------------------------------===//
+  // Document Symbols
+  //===--------------------------------------------------------------------===//
+
+  void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   /// The include directories for this file.
   std::vector<std::string> includeDirs;
 
   /// The source manager containing the contents of the input file.
   llvm::SourceMgr sourceMgr;
if (auto resultTupleTy = resultType.dyn_cast<ast::TupleType>()) {
       hoverOS << "Results:\n* `" << resultType << "`\n";
     }
     hoverOS << "***\n";
   }
   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();
+
+      llvm::SMRange nameLoc = name ? name->getLoc() : patternDecl->getLoc();
+      llvm::SMRange bodyLoc(nameLoc.Start,
+                            patternDecl->getBody()->getLoc().End);
+
+      // If this is a template instantiation, the location of the body may be
+      // non-contiguous from the derived pattern.
+      if (bodyLoc.End.getPointer() < bodyLoc.Start.getPointer())
+        bodyLoc = nameLoc;
+
+      symbols.emplace_back(name ? name->getName() : "<pattern>",
+                           lsp::SymbolKind::Class,
+                           getRangeFromLoc(sourceMgr, bodyLoc),
+                           getRangeFromLoc(sourceMgr, nameLoc));
+    } else if (const auto *cDecl = dyn_cast<ast::UserConstraintDecl>(decl)) {
+      // TODO: Add source information for the code block body.
+      llvm::SMRange nameLoc = cDecl->getName().getLoc();
+      llvm::SMRange bodyLoc = nameLoc;
+
+      symbols.emplace_back(cDecl->getName().getName(),
+                           lsp::SymbolKind::Function,
+                           getRangeFromLoc(sourceMgr, bodyLoc),
+                           getRangeFromLoc(sourceMgr, nameLoc));
+    } else if (const auto *cDecl = dyn_cast<ast::UserRewriteDecl>(decl)) {
+      // TODO: Add source information for the code block body.
+      llvm::SMRange nameLoc = cDecl->getName().getLoc();
+      llvm::SMRange bodyLoc = nameLoc;
+
+      symbols.emplace_back(cDecl->getName().getName(),
+                           lsp::SymbolKind::Function,
+                           getRangeFromLoc(sourceMgr, bodyLoc),
+                           getRangeFromLoc(sourceMgr, nameLoc));
+    }
+  }
+}
+
+//===----------------------------------------------------------------------===//
 // PDLTextFileChunk
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class represents a single chunk of an PDL text file.
 struct PDLTextFileChunk {
   PDLTextFileChunk(uint64_t lineOffset, const lsp::URIForFile &uri,
                    StringRef contents,
public:
   //===--------------------------------------------------------------------===//
 
   void getLocationsOf(const lsp::URIForFile &uri, lsp::Position defPos,
                       std::vector<lsp::Location> &locations);
   void findReferencesOf(const lsp::URIForFile &uri, lsp::Position pos,
                         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
   /// position to be anchored at the start of the found chunk instead of the
   /// beginning of the file.
   PDLTextFileChunk &getChunkFor(lsp::Position &pos);
 
   /// The full string contents of the file.
Optional<lsp::Hover> PDLTextFile::findHover(const lsp::URIForFile &uri,
   Optional<lsp::Hover> hoverInfo = chunk.document.findHover(uri, hoverPos);
 
   // Adjust any locations within this file for the offset of this chunk.
   if (chunk.lineOffset != 0 && hoverInfo && hoverInfo->range)
     chunk.adjustLocForChunkOffset(*hoverInfo->range);
   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();
 
   // Search for the first chunk with a greater line offset, the previous chunk
   // is the one that contains `pos`.
   auto it = llvm::upper_bound(
       chunks, pos, [](const lsp::Position &pos, const auto &chunk) {

 
 Optional<lsp::Hover> lsp::PDLLServer::findHover(const URIForFile &uri,
                                                 const Position &hoverPos) {
   auto fileIt = impl->files.find(uri.file());
   if (fileIt != impl->files.end())
     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);
+}

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

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

mlir/test/mlir-pdll-lsp-server/initialize-params.test

 // RUN: mlir-pdll-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
 // Test initialize request parameters with rootUri
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootUri":"test:///workspace","capabilities":{},"trace":"off"}}
 //      CHECK:  "id": 0,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "capabilities": {
 // CHECK-NEXT:      "definitionProvider": true,
+// CHECK-NEXT:      "documentSymbolProvider": false,
 // CHECK-NEXT:      "hoverProvider": true,
 // CHECK-NEXT:      "referencesProvider": true,
 // CHECK-NEXT:      "textDocumentSync": {
 // CHECK-NEXT:        "change": 1,
 // CHECK-NEXT:        "openClose": true,
 // CHECK-NEXT:        "save": true
 // CHECK-NEXT:      }
 // CHECK-NEXT:    },