[mlir-lsp-server] Add support for textDocument/documentSymbols

[rriddle]

Summary

This allows for building an outline of the symbols and symbol tables within the IR. This allows for easy navigations to functions/modules and other symbol/symbol table operations within the IR.

Depends On D103728

[rriddle]

update

[jpienaar]

Looks good overall thanks

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

635

Nit: I'd have this inside the else (the indenting here looks weird else as makes it feel like it is attached to block above).

mlir/lib/Tools/mlir-lsp-server/lsp/Protocol.h

476

OOC are these just default LSP ones?

[rriddle]

mlir/lib/Tools/mlir-lsp-server/lsp/Protocol.h

476

Yeah, https://microsoft.github.io/language-server-protocol/specifications/specification-3-17/#symbolKind

mlir/include/mlir/Parser/AsmParserState.h

public:
   struct OperationDefinition {
     struct ResultGroupDefinition {
       /// The result number that starts this group.
       unsigned startIndex;
       /// The source definition of the result group.
       SMDefinition definition;
     };
 
-    OperationDefinition(Operation *op, llvm::SMRange loc) : op(op), loc(loc) {}
+    OperationDefinition(Operation *op, llvm::SMRange loc, llvm::SMLoc endLoc)
+        : op(op), loc(loc), scopeLoc(loc.Start, endLoc) {}
 
     /// The operation representing this definition.
     Operation *op;
 
     /// The source location for the operation, i.e. the location of its name.
     llvm::SMRange loc;
 
+    /// The full source range of the operation definition, i.e. a range
+    /// encompassing the start and end of the full operation definition.
+    llvm::SMRange scopeLoc;
+
     /// Source definitions for any result groups of this operation.
     SmallVector<std::pair<unsigned, SMDefinition>> resultGroups;
 
     /// If this operation is a symbol operation, this vector contains symbol
     /// uses of this operation.
     SmallVector<llvm::SMRange> symbolUses;
   };
 
public:
   /// Return the definition for the given block, or nullptr if the given
   /// block does not have a definition.
   const BlockDefinition *getBlockDef(Block *block) const;
 
   /// Return a range of the OperationDefinitions held by the current parser
   /// state.
   iterator_range<OperationDefIterator> getOpDefs() const;
 
+  /// Return the definition for the given operation, or nullptr if the given
+  /// operation does not have a definition.
+  const OperationDefinition *getOpDef(Operation *op) const;
+
   /// Returns (heuristically) the range of an identifier given a SMLoc
   /// corresponding to the start of an identifier location.
   static llvm::SMRange convertIdLocToRange(llvm::SMLoc loc);
 
   //===--------------------------------------------------------------------===//
   // Populate State
   //===--------------------------------------------------------------------===//
 
   /// Initialize the state in preparation for populating more parser state under
   /// the given top-level operation.
   void initialize(Operation *topLevelOp);
 
   /// Finalize any in-progress parser state under the given top-level operation.
   void finalize(Operation *topLevelOp);
 
   /// Start a definition for an operation with the given name.
   void startOperationDefinition(const OperationName &opName);
 
   /// Finalize the most recently started operation definition.
   void finalizeOperationDefinition(
-      Operation *op, llvm::SMRange nameLoc,
+      Operation *op, llvm::SMRange nameLoc, llvm::SMLoc endLoc,
       ArrayRef<std::pair<unsigned, llvm::SMLoc>> resultGroups = llvm::None);
 
   /// Start a definition for a region nested under the current operation.
   void startRegionDefinition();
 
   /// Finalize the most recently started region definition.
   void finalizeRegionDefinition();
 

mlir/lib/Parser/AsmParserState.cpp

struct PartialOpDef {
     /// Return if this operation is a symbol table.
     bool isSymbolTable() const { return symbolTable.get(); }
 
     /// If this operation is a symbol table, the following contains symbol uses
     /// within this operation.
     std::unique_ptr<SymbolUseMap> symbolTable;
   };
 
-  /// Resolve any symbol table uses under the given partial operation.
-  void resolveSymbolUses(Operation *op, PartialOpDef &opDef);
+  /// Resolve any symbol table uses in the IR.
+  void resolveSymbolUses();
 
   /// A mapping from operations in the input source file to their parser state.
   SmallVector<std::unique_ptr<OperationDefinition>> operations;
   DenseMap<Operation *, unsigned> operationToIdx;
 
   /// A mapping from blocks in the input source file to their parser state.
   SmallVector<std::unique_ptr<BlockDefinition>> blocks;
   DenseMap<Block *, unsigned> blocksToIdx;
 
   /// A set of value definitions that are placeholders for forward references.
   /// This map should be empty if the parser finishes successfully.
   DenseMap<Value, SmallVector<llvm::SMLoc>> placeholderValueUses;
 
+  /// The symbol table operations within the IR.
+  SmallVector<std::pair<Operation *, std::unique_ptr<SymbolUseMap>>>
+      symbolTableOperations;
+
   /// A stack of partial operation definitions that have been started but not
   /// yet finalized.
   SmallVector<PartialOpDef> partialOperations;
 
   /// A stack of symbol use scopes. This is used when collecting symbol table
   /// uses during parsing.
   SmallVector<SymbolUseMap *> symbolUseScopes;
 
   /// A symbol table containing all of the symbol table operations in the IR.
   SymbolTableCollection symbolTable;
 };
 
-void AsmParserState::Impl::resolveSymbolUses(Operation *op,
-                                             PartialOpDef &opDef) {
-  assert(opDef.isSymbolTable() && "expected op to be a symbol table");
-
+void AsmParserState::Impl::resolveSymbolUses() {
   SmallVector<Operation *> symbolOps;
-  for (auto &it : *opDef.symbolTable) {
+  for (auto &opAndUseMapIt : symbolTableOperations) {
+    for (auto &it : *opAndUseMapIt.second) {
       symbolOps.clear();
-    if (failed(symbolTable.lookupSymbolIn(op, it.first.cast<SymbolRefAttr>(),
-                                          symbolOps)))
+      if (failed(symbolTable.lookupSymbolIn(
+              opAndUseMapIt.first, it.first.cast<SymbolRefAttr>(), symbolOps)))
         continue;
 
       for (ArrayRef<llvm::SMRange> useRange : it.second) {
         for (const auto &symIt : llvm::zip(symbolOps, useRange)) {
           auto opIt = operationToIdx.find(std::get<0>(symIt));
           if (opIt != operationToIdx.end())
             operations[opIt->second]->symbolUses.push_back(std::get<1>(symIt));
         }
       }
     }
   }
+}
 
 //===----------------------------------------------------------------------===//
 // AsmParserState
 //===----------------------------------------------------------------------===//
 
 AsmParserState::AsmParserState() : impl(std::make_unique<Impl>()) {}
 AsmParserState::~AsmParserState() {}
 AsmParserState &AsmParserState::operator=(AsmParserState &&other) {
auto AsmParserState::getBlockDef(Block *block) const
   auto it = impl->blocksToIdx.find(block);
   return it == impl->blocksToIdx.end() ? nullptr : &*impl->blocks[it->second];
 }
 
 auto AsmParserState::getOpDefs() const -> iterator_range<OperationDefIterator> {
   return llvm::make_pointee_range(llvm::makeArrayRef(impl->operations));
 }
 
-/// Returns (heuristically) the range of an identifier given a SMLoc
-/// corresponding to the start of an identifier location.
+auto AsmParserState::getOpDef(Operation *op) const
+    -> const OperationDefinition * {
+  auto it = impl->operationToIdx.find(op);
+  return it == impl->operationToIdx.end() ? nullptr
+                                          : &*impl->operations[it->second];
+}
+
 llvm::SMRange AsmParserState::convertIdLocToRange(llvm::SMLoc loc) {
   if (!loc.isValid())
     return llvm::SMRange();
 
   // Return if the given character is a valid identifier character.
   auto isIdentifierChar = [](char c) {
     return isalnum(c) || c == '$' || c == '.' || c == '_' || c == '-';
   };
 
   const char *curPtr = loc.getPointer();
-  while (isIdentifierChar(*(++curPtr)))
+  while (*curPtr && isIdentifierChar(*(++curPtr)))
     continue;
   return llvm::SMRange(loc, llvm::SMLoc::getFromPointer(curPtr));
 }
 
 //===----------------------------------------------------------------------===//
 // Populate State
 
 void AsmParserState::initialize(Operation *topLevelOp) {
   startOperationDefinition(topLevelOp->getName());
 
   // If the top-level operation is a symbol table, push a new symbol scope.
   Impl::PartialOpDef &partialOpDef = impl->partialOperations.back();
   if (partialOpDef.isSymbolTable())
     impl->symbolUseScopes.push_back(partialOpDef.symbolTable.get());
 }
 
 void AsmParserState::finalize(Operation *topLevelOp) {
   assert(!impl->partialOperations.empty() &&
          "expected valid partial operation definition");
   Impl::PartialOpDef partialOpDef = impl->partialOperations.pop_back_val();
 
   // If this operation is a symbol table, resolve any symbol uses.
-  if (partialOpDef.isSymbolTable())
-    impl->resolveSymbolUses(topLevelOp, partialOpDef);
+  if (partialOpDef.isSymbolTable()) {
+    impl->symbolTableOperations.emplace_back(
+        topLevelOp, std::move(partialOpDef.symbolTable));
+  }
+  impl->resolveSymbolUses();
 }
 
 void AsmParserState::startOperationDefinition(const OperationName &opName) {
   impl->partialOperations.emplace_back(opName);
 }
 
 void AsmParserState::finalizeOperationDefinition(
-    Operation *op, llvm::SMRange nameLoc,
+    Operation *op, llvm::SMRange nameLoc, llvm::SMLoc endLoc,
     ArrayRef<std::pair<unsigned, llvm::SMLoc>> resultGroups) {
   assert(!impl->partialOperations.empty() &&
          "expected valid partial operation definition");
   Impl::PartialOpDef partialOpDef = impl->partialOperations.pop_back_val();
 
   // Build the full operation definition.
   std::unique_ptr<OperationDefinition> def =
-      std::make_unique<OperationDefinition>(op, nameLoc);
+      std::make_unique<OperationDefinition>(op, nameLoc, endLoc);
   for (auto &resultGroup : resultGroups)
     def->resultGroups.emplace_back(resultGroup.first,
                                    convertIdLocToRange(resultGroup.second));
   impl->operationToIdx.try_emplace(op, impl->operations.size());
   impl->operations.emplace_back(std::move(def));
 
   // If this operation is a symbol table, resolve any symbol uses.
-  if (partialOpDef.isSymbolTable())
-    impl->resolveSymbolUses(op, partialOpDef);
+  if (partialOpDef.isSymbolTable()) {
+    impl->symbolTableOperations.emplace_back(
+        op, std::move(partialOpDef.symbolTable));
+  }
 }
 
 void AsmParserState::startRegionDefinition() {
   assert(!impl->partialOperations.empty() &&
          "expected valid partial operation definition");
 
   // If the parent operation of this region is a symbol table, we also push a
   // new symbol scope.

mlir/lib/Parser/Parser.cpp

 // OperationParser
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class provides support for parsing operations and regions of
 /// operations.
 class OperationParser : public Parser {
 public:
-  OperationParser(ParserState &state, Operation *topLevelOp);
+  OperationParser(ParserState &state, ModuleOp topLevelOp);
   ~OperationParser();
 
   /// After parsing is finished, this function must be called to see if there
   /// are any remaining issues.
-  ParseResult finalize(Operation *topLevelOp);
+  ParseResult finalize();
 
   //===--------------------------------------------------------------------===//
   // SSA Value Handling
   //===--------------------------------------------------------------------===//
 
   /// This represents a use of an SSA value in the program.  The first two
   /// entries in the tuple are the name and result number of a reference.  The
   /// third is the location of the reference, which is used in case this ends
private:
   /// The builder used when creating parsed operation instances.
   OpBuilder opBuilder;
 
   /// The top level operation that holds all of the parsed operations.
   Operation *topLevelOp;
 };
 } // end anonymous namespace
 
-OperationParser::OperationParser(ParserState &state, Operation *topLevelOp)
-    : Parser(state), opBuilder(topLevelOp->getRegion(0)),
-      topLevelOp(topLevelOp) {
+OperationParser::OperationParser(ParserState &state, ModuleOp topLevelOp)
+    : Parser(state), opBuilder(topLevelOp.getRegion()), topLevelOp(topLevelOp) {
   // The top level operation starts a new name scope.
   pushSSANameScope(/*isIsolated=*/true);
 
   // If we are populating the parser state, prepare it for parsing.
   if (state.asmState)
     state.asmState->initialize(topLevelOp);
 }
 
for (const auto &fwd : scope) {
       fwd.first->dropAllUses();
       delete fwd.first;
     }
   }
 }
 
 /// After parsing is finished, this function must be called to see if there are
 /// any remaining issues.
-ParseResult OperationParser::finalize(Operation *topLevelOp) {
+ParseResult OperationParser::finalize() {
   // Check for any forward references that are left.  If we find any, error
   // out.
   if (!forwardRefPlaceholders.empty()) {
     SmallVector<const char *, 4> errors;
     // Iteration over the map isn't deterministic, so sort by source location.
     for (auto entry : forwardRefPlaceholders)
       errors.push_back(entry.second.getPointer());
     llvm::array_pod_sort(errors.begin(), errors.end());
if (!locAttr)
              << "expected location, but found '" << attr << "'";
     auto opOrArgument = it.first;
     if (auto *op = opOrArgument.dyn_cast<Operation *>())
       op->setLoc(locAttr);
     else
       opOrArgument.get<BlockArgument>().setLoc(locAttr);
   }
 
+  // Pop the top level name scope.
+  if (failed(popSSANameScope()))
+    return failure();
+
+  // Verify that the parsed operations are valid.
+  if (failed(verify(topLevelOp)))
+    return failure();
+
   // If we are populating the parser state, finalize the top-level operation.
   if (state.asmState)
     state.asmState->finalize(topLevelOp);
-
-  // Pop the top level name scope.
-  return popSSANameScope();
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // SSA Value Handling
 //===----------------------------------------------------------------------===//
 
 void OperationParser::pushSSANameScope(bool isIsolated) {
   blocksByName.push_back(DenseMap<StringRef, BlockDefinition>());
if (!resultIDs.empty()) {
     if (state.asmState) {
       unsigned resultIt = 0;
       SmallVector<std::pair<unsigned, llvm::SMLoc>> asmResultGroups;
       asmResultGroups.reserve(resultIDs.size());
       for (ResultRecord &record : resultIDs) {
         asmResultGroups.emplace_back(resultIt, std::get<2>(record));
         resultIt += std::get<1>(record);
       }
-      state.asmState->finalizeOperationDefinition(op, nameTok.getLocRange(),
-                                                  asmResultGroups);
+      state.asmState->finalizeOperationDefinition(
+          op, nameTok.getLocRange(), /*endLoc=*/getToken().getLoc(),
     }
 
     // Add definitions for each of the result groups.
     unsigned opResI = 0;
     for (ResultRecord &resIt : resultIDs) {
       for (unsigned subRes : llvm::seq<unsigned>(0, std::get<1>(resIt))) {
         if (addDefinition({std::get<0>(resIt), subRes, std::get<2>(resIt)},
                           op->getResult(opResI++)))
           return failure();
       }
     }
 
     // Add this operation to the assembly state if it was provided to populate.
   } else if (state.asmState) {
-    state.asmState->finalizeOperationDefinition(op, nameTok.getLocRange());
+    state.asmState->finalizeOperationDefinition(op, nameTok.getLocRange(),
+                                                /*endLoc=*/getToken().getLoc());
   }
 
   return success();
 }
 
 /// Parse a single operation successor.
 ///
 ///   successor ::= block-id
Operation *OperationParser::parseGenericOperation(Block *insertBlock,
   opBuilder.setInsertionPoint(insertBlock, insertPt);
   Operation *op = parseGenericOperation();
   if (!op)
     return nullptr;
 
   // If we are populating the parser asm state, finalize this operation
   // definition.
   if (state.asmState)
-    state.asmState->finalizeOperationDefinition(op, nameToken.getLocRange());
+    state.asmState->finalizeOperationDefinition(op, nameToken.getLocRange(),
+                                                /*endLoc=*/getToken().getLoc());
   return op;
 }
 
 namespace {
 class CustomOpAsmParser : public OpAsmParser {
 public:
   CustomOpAsmParser(
       SMLoc nameLoc, ArrayRef<OperationParser::ResultRecord> resultIDs,
ParseResult OperationParser::parseRegionBody(
   // If this block is not defined in the source file, add a definition for it
   // now in the assembly state. Blocks with a name will be defined when the name
   // is parsed.
   if (state.asmState && getToken().isNot(Token::caret_identifier))
     state.asmState->addDefinition(block, startLoc);
 
   // Add arguments to the entry block.
   if (!entryArguments.empty()) {
+    // If we had named arguments, then don't allow a block name.
+    if (getToken().is(Token::caret_identifier))
+      return emitError("invalid block name in region with named arguments");
+
     for (auto &placeholderArgPair : entryArguments) {
       auto &argInfo = placeholderArgPair.first;
 
       // Ensure that the argument was not already defined.
       if (auto defLoc = getReferenceLoc(argInfo.name, argInfo.number)) {
         return emitError(argInfo.loc, "region entry argument '" + argInfo.name +
                                           "' is already in use")
                    .attachNote(getEncodedSourceLocation(*defLoc))
                << "previously referenced here";
       }
       auto loc = getEncodedSourceLocation(placeholderArgPair.first.loc);
       BlockArgument arg = block->addArgument(placeholderArgPair.second, loc);
 
       // Add a definition of this arg to the assembly state if provided.
       if (state.asmState)
         state.asmState->addDefinition(arg, argInfo.loc);
 
       // Record the definition for this argument.
       if (addDefinition(argInfo, arg))
         return failure();
     }
-
-    // If we had named arguments, then don't allow a block name.
-    if (getToken().is(Token::caret_identifier))
-      return emitError("invalid block name in region with named arguments");
   }
 
   if (parseBlock(block))
     return failure();
 
   // Verify that no other arguments were parsed.
   if (!entryArguments.empty() &&
       block->getNumArguments() > entryArguments.size()) {
ParseResult TopLevelOperationParser::parseTypeAliasDef() {
   // Register this alias with the parser state.
   state.symbols.typeAliasDefinitions.try_emplace(aliasName, aliasedType);
   return success();
 }
 
 ParseResult TopLevelOperationParser::parse(Block *topLevelBlock,
                                            Location parserLoc) {
   // Create a top-level operation to contain the parsed state.
-  OwningOpRef<Operation *> topLevelOp(ModuleOp::create(parserLoc));
+  OwningOpRef<ModuleOp> topLevelOp(ModuleOp::create(parserLoc));
   OperationParser opParser(state, topLevelOp.get());
   while (true) {
     switch (getToken().getKind()) {
     default:
       // Parse a top-level operation.
       if (opParser.parseOperation())
         return failure();
       break;
 
     // If we got to the end of the file, then we're done.
     case Token::eof: {
-      if (opParser.finalize(topLevelOp.get()))
-        return failure();
-
-      // Verify that the parsed operations are valid.
-      if (failed(verify(topLevelOp.get())))
+      if (opParser.finalize())
         return failure();
 
       // Splice the blocks of the parsed operation over to the provided
       // top-level block.
-      auto &parsedOps = (*topLevelOp)->getRegion(0).front().getOperations();
+      auto &parsedOps = topLevelOp->getBody()->getOperations();
       auto &destOps = topLevelBlock->getOperations();
       destOps.splice(destOps.empty() ? destOps.end() : std::prev(destOps.end()),
                      parsedOps, parsedOps.begin(), parsedOps.end());
       return success();
     }
 
     // If we got an error token, then the lexer already emitted an error, just
     // stop.  Someday we could introduce error recovery if there was demand

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

void onReference(const ReferenceParams &params,
                    Callback<std::vector<Location>> reply);
 
   //===--------------------------------------------------------------------===//
   // Hover
 
   void onHover(const TextDocumentPositionParams &params,
                Callback<Optional<Hover>> reply);
 
+  //===--------------------------------------------------------------------===//
+  // Document Symbols
+
+  void onDocumentSymbol(const DocumentSymbolParams &params,
+                        Callback<std::vector<DocumentSymbol>> reply);
+
+  //===--------------------------------------------------------------------===//
+  // Fields
+  //===--------------------------------------------------------------------===//
+
   MLIRServer &server;
   JSONTransport &transport;
 
   /// An outgoing notification used to send diagnostics to the client when they
   /// are ready to be processed.
   OutgoingNotification<PublishDiagnosticsParams> publishDiagnostics;
 
   /// Used to indicate that the 'shutdown' request was received from the
   /// Language Server client.
   bool shutdownRequestReceived = false;
 };
 
 //===----------------------------------------------------------------------===//
 // Initialization
 
 void LSPServer::Impl::onInitialize(const InitializeParams &params,
                                    Callback<llvm::json::Value> reply) {
+  // Send a response with the capabilities of this server.
   llvm::json::Object serverCaps{
       {"textDocumentSync",
        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-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(MLIRServer &server, JSONTransport &transport)
     : impl(std::make_unique<Impl>(server, transport)) {}
 LSPServer::~LSPServer() {}
 
 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.
   LogicalResult result = success();
   if (llvm::Error error = impl->transport.run(messageHandler)) {
     Logger::error("Transport error: {0}", error);
     llvm::consumeError(std::move(error));
     result = failure();
   } else {
     result = success(impl->shutdownRequestReceived);
   }
   return result;
 }

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

 #include "mlir/Support/LLVM.h"
 #include <memory>
 
 namespace mlir {
 class DialectRegistry;
 
 namespace lsp {
 struct Diagnostic;
+struct DocumentSymbol;
 struct Hover;
 struct Location;
 struct Position;
 class URIForFile;
 
 /// This class implements all of the MLIR related functionality necessary for a
 /// language server. This class allows for keeping the MLIR 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_MLIRLSPSERVER_SERVER_H_

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

lsp::Hover buildHoverForOperationResult(llvm::SMRange hoverRange,
                                           unsigned resultEnd,
                                           llvm::SMLoc posLoc);
   lsp::Hover buildHoverForBlock(llvm::SMRange hoverRange,
                                 const AsmParserState::BlockDefinition &block);
   lsp::Hover
   buildHoverForBlockArgument(llvm::SMRange hoverRange, BlockArgument arg,
                              const AsmParserState::BlockDefinition &block);
 
+  //===--------------------------------------------------------------------===//
+  // Document Symbols
+  //===--------------------------------------------------------------------===//
+
+  void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);
+  void findDocumentSymbols(Operation *op,
+                           std::vector<lsp::DocumentSymbol> &symbols);
+
+  //===--------------------------------------------------------------------===//
+  // Fields
+  //===--------------------------------------------------------------------===//
+
   /// The context used to hold the state contained by the parsed document.
   MLIRContext context;
 
   /// The high level parser state used to find definitions and references within
   /// the source file.
   AsmParserState asmState;
 
   /// The container for the IR parsed from the input file.
lsp::Hover MLIRDocument::buildHoverForBlockArgument(
   printDefBlockName(os, block);
   os << "\n\nArgument #" << arg.getArgNumber() << "\n\n"
      << "Type: `" << arg.getType() << "`\n\n";
 
   return hover;
 }
 
 //===----------------------------------------------------------------------===//
+// MLIRDocument: Document Symbols
+//===----------------------------------------------------------------------===//
+
+void MLIRDocument::findDocumentSymbols(
+    std::vector<lsp::DocumentSymbol> &symbols) {
+  for (Operation &op : parsedIR)
+    findDocumentSymbols(&op, symbols);
+}
+
+void MLIRDocument::findDocumentSymbols(
+    Operation *op, std::vector<lsp::DocumentSymbol> &symbols) {
+  std::vector<lsp::DocumentSymbol> *childSymbols = &symbols;
+
+  // Check for the source information of this operation.
+  if (const AsmParserState::OperationDefinition *def = asmState.getOpDef(op)) {
+    // If this operation defines a symbol, record it.
+    if (SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(op)) {
+      symbols.emplace_back(symbol.getName(),
+                           op->hasTrait<OpTrait::FunctionLike>()
+                               ? lsp::SymbolKind::Function
+                               : lsp::SymbolKind::Class,
+                           getRangeFromLoc(sourceMgr, def->scopeLoc),
+                           getRangeFromLoc(sourceMgr, def->loc));
+      childSymbols = &symbols.back().children;
+
+    } else if (op->hasTrait<OpTrait::SymbolTable>()) {
+      // Otherwise, if this is a symbol table push an anonymous document symbol.
+      symbols.emplace_back("<" + op->getName().getStringRef() + ">",
+                           lsp::SymbolKind::Namespace,
+                           getRangeFromLoc(sourceMgr, def->scopeLoc),
+                           getRangeFromLoc(sourceMgr, def->loc));
+      childSymbols = &symbols.back().children;
+    }
+  }
+
+  // Recurse into the regions of this operation.
+  if (!op->getNumRegions())
+    return;
+  for (Region &region : op->getRegions())
+    for (Operation &childOp : region.getOps())
+      findDocumentSymbols(&childOp, *childSymbols);
+}
+
+//===----------------------------------------------------------------------===//
 // MLIRTextFileChunk
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class represents a single chunk of an MLIR text file.
 struct MLIRTextFileChunk {
   MLIRTextFileChunk(uint64_t lineOffset, const lsp::URIForFile &uri,
                     StringRef contents, DialectRegistry &registry,
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 MLIR 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.
   MLIRTextFileChunk &getChunkFor(lsp::Position &pos);
 
   /// The full string contents of the file.
   std::string contents;
 
   /// The version of this file.
   int64_t version;
 
+  /// The number of lines in the file.
+  int64_t totalNumLines;
+
   /// The chunks of this file. The order of these chunks is the order in which
   /// they appear in the text file.
   std::vector<std::unique_ptr<MLIRTextFileChunk>> chunks;
 };
 } // namespace
 
 MLIRTextFile::MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents,
                            int64_t version, DialectRegistry &registry,
                            std::vector<lsp::Diagnostic> &diagnostics)
-    : contents(fileContents.str()), version(version) {
+    : contents(fileContents.str()), version(version), totalNumLines(0) {
   // Split the file into separate MLIR documents.
   // TODO: Find a way to share the split file marker with other tools. We don't
   // want to use `splitAndProcessBuffer` here, but we do want to make sure this
   // marker doesn't go out of sync.
   SmallVector<StringRef, 8> subContents;
   StringRef(contents).split(subContents, "// -----");
   chunks.emplace_back(std::make_unique<MLIRTextFileChunk>(
       /*lineOffset=*/0, uri, subContents.front(), registry, diagnostics));
for (lsp::Diagnostic &diag :
       if (!diag.relatedInformation)
         continue;
       for (auto &it : *diag.relatedInformation)
         if (it.location.uri == uri)
           chunk->adjustLocForChunkOffset(it.location.range);
     }
     chunks.emplace_back(std::move(chunk));
   }
+  totalNumLines = lineOffset;
 }
 
 void MLIRTextFile::getLocationsOf(const lsp::URIForFile &uri,
                                   lsp::Position defPos,
                                   std::vector<lsp::Location> &locations) {
   MLIRTextFileChunk &chunk = getChunkFor(defPos);
   chunk.document.getLocationsOf(uri, defPos, locations);
 
Optional<lsp::Hover> MLIRTextFile::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 MLIRTextFile::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) {
+    MLIRTextFileChunk &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));
+  }
+}
+
 MLIRTextFileChunk &MLIRTextFile::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::MLIRServer::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::MLIRServer::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/lib/Tools/mlir-lsp-server/lsp/Protocol.h

 
 /// Add support for JSON serialization.
 llvm::json::Value toJSON(const URIForFile &value);
 bool fromJSON(const llvm::json::Value &value, URIForFile &result,
               llvm::json::Path path);
 raw_ostream &operator<<(raw_ostream &os, const URIForFile &value);
 
 //===----------------------------------------------------------------------===//
+// ClientCapabilities
+//===----------------------------------------------------------------------===//
+
+struct ClientCapabilities {
+  /// Client supports hierarchical document symbols.
+  /// textDocument.documentSymbol.hierarchicalDocumentSymbolSupport
+  bool hierarchicalDocumentSymbol = false;
+};
+
+/// Add support for JSON serialization.
+bool fromJSON(const llvm::json::Value &value, ClientCapabilities &result,
+              llvm::json::Path path);
+
+//===----------------------------------------------------------------------===//
 // InitializeParams
 //===----------------------------------------------------------------------===//
 
 enum class TraceLevel {
   Off = 0,
   Messages = 1,
   Verbose = 2,
 };
 
 /// Add support for JSON serialization.
 bool fromJSON(const llvm::json::Value &value, TraceLevel &result,
               llvm::json::Path path);
 
 struct InitializeParams {
+  /// The capabilities provided by the client (editor or tool).
+  ClientCapabilities capabilities;
+
   /// The initial trace setting. If omitted trace is disabled ('off').
   Optional<TraceLevel> trace;
 };
 
 /// Add support for JSON serialization.
 bool fromJSON(const llvm::json::Value &value, InitializeParams &result,
               llvm::json::Path path);
 

 bool fromJSON(const llvm::json::Value &value,
               VersionedTextDocumentIdentifier &result, llvm::json::Path path);
 
 //===----------------------------------------------------------------------===//
 // Position
 //===----------------------------------------------------------------------===//
 
 struct Position {
+  Position(int line = 0, int character = 0)
+      : line(line), character(character) {}
+
   /// Line position in a document (zero-based).
   int line = 0;
 
   /// Character offset on a line in a document (zero-based).
   int character = 0;
 
   friend bool operator==(const Position &lhs, const Position &rhs) {
     return std::tie(lhs.line, lhs.character) ==
struct Hover {
   /// visualize a hover, e.g. by changing the background color.
   Optional<Range> range;
 };
 
 /// Add support for JSON serialization.
 llvm::json::Value toJSON(const Hover &hover);
 
 //===----------------------------------------------------------------------===//
+// SymbolKind
+//===----------------------------------------------------------------------===//
+
+enum class SymbolKind {
+  File = 1,
+  Module = 2,
+  Namespace = 3,
+  Package = 4,
+  Class = 5,
+  Method = 6,
+  Property = 7,
+  Field = 8,
+  Constructor = 9,
+  Enum = 10,
+  Interface = 11,
+  Function = 12,
+  Variable = 13,
+  Constant = 14,
+  String = 15,
+  Number = 16,
+  Boolean = 17,
+  Array = 18,
+  Object = 19,
+  Key = 20,
+  Null = 21,
+  EnumMember = 22,
+  Struct = 23,
+  Event = 24,
+  Operator = 25,
+  TypeParameter = 26
+};
+
+//===----------------------------------------------------------------------===//
+// DocumentSymbol
+//===----------------------------------------------------------------------===//
+
+/// Represents programming constructs like variables, classes, interfaces etc.
+/// that appear in a document. Document symbols can be hierarchical and they
+/// have two ranges: one that encloses its definition and one that points to its
+/// most interesting range, e.g. the range of an identifier.
+struct DocumentSymbol {
+  DocumentSymbol() = default;
+  DocumentSymbol(DocumentSymbol &&) = default;
+  DocumentSymbol(const Twine &name, SymbolKind kind, Range range,
+                 Range selectionRange)
+      : name(name.str()), kind(kind), range(range),
+        selectionRange(selectionRange) {}
+
+  /// The name of this symbol.
+  std::string name;
+
+  /// More detail for this symbol, e.g the signature of a function.
+  std::string detail;
+
+  /// The kind of this symbol.
+  SymbolKind kind;
+
+  /// The range enclosing this symbol not including leading/trailing whitespace
+  /// but everything else like comments. This information is typically used to
+  /// determine if the clients cursor is inside the symbol to reveal in the
+  /// symbol in the UI.
+  Range range;
+
+  /// The range that should be selected and revealed when this symbol is being
+  /// picked, e.g the name of a function. Must be contained by the `range`.
+  Range selectionRange;
+
+  /// Children of this symbol, e.g. properties of a class.
+  std::vector<DocumentSymbol> children;
+};
+
+/// Add support for JSON serialization.
+llvm::json::Value toJSON(const DocumentSymbol &symbol);
+
+//===----------------------------------------------------------------------===//
+// DocumentSymbolParams
+//===----------------------------------------------------------------------===//
+
+struct DocumentSymbolParams {
+  // The text document to find symbols in.
+  TextDocumentIdentifier textDocument;
+};
+
+/// Add support for JSON serialization.
+bool fromJSON(const llvm::json::Value &value, DocumentSymbolParams &result,
+              llvm::json::Path path);
+
+//===----------------------------------------------------------------------===//
 // DiagnosticRelatedInformation
 //===----------------------------------------------------------------------===//
 
 /// Represents a related message and source code location for a diagnostic.
 /// This should be used to point to code locations that cause or related to a
 /// diagnostics, e.g. when duplicating a symbol in a scope.
 struct DiagnosticRelatedInformation {
   DiagnosticRelatedInformation(Location location, std::string message)

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

llvm::json::Value mlir::lsp::toJSON(const URIForFile &value) {
   return value.uri();
 }
 
 raw_ostream &mlir::lsp::operator<<(raw_ostream &os, const URIForFile &value) {
   return os << value.uri();
 }
 
 //===----------------------------------------------------------------------===//
+// ClientCapabilities
+//===----------------------------------------------------------------------===//
+
+bool mlir::lsp::fromJSON(const llvm::json::Value &value,
+                         ClientCapabilities &result, llvm::json::Path path) {
+  const llvm::json::Object *o = value.getAsObject();
+  if (!o) {
+    path.report("expected object");
+    return false;
+  }
+  if (const llvm::json::Object *textDocument = o->getObject("textDocument")) {
+    if (const llvm::json::Object *documentSymbol =
+            textDocument->getObject("documentSymbol")) {
+      if (Optional<bool> hierarchicalSupport =
+              documentSymbol->getBoolean("hierarchicalDocumentSymbolSupport"))
+        result.hierarchicalDocumentSymbol = *hierarchicalSupport;
+    }
+  }
+  return true;
+}
+
+//===----------------------------------------------------------------------===//
 // InitializeParams
 //===----------------------------------------------------------------------===//
 
 bool mlir::lsp::fromJSON(const llvm::json::Value &value, TraceLevel &result,
                          llvm::json::Path path) {
   if (Optional<StringRef> str = value.getAsString()) {
     if (*str == "off") {
       result = TraceLevel::Off;

 }
 
 bool mlir::lsp::fromJSON(const llvm::json::Value &value,
                          InitializeParams &result, llvm::json::Path path) {
   llvm::json::ObjectMapper o(value, path);
   if (!o)
     return false;
   // We deliberately don't fail if we can't parse individual fields.
+  o.map("capabilities", result.capabilities);
   o.map("trace", result.trace);
   return true;
 }
 
 //===----------------------------------------------------------------------===//
 // TextDocumentItem
 //===----------------------------------------------------------------------===//
 

 llvm::json::Value mlir::lsp::toJSON(const Hover &hover) {
   llvm::json::Object result{{"contents", toJSON(hover.contents)}};
   if (hover.range.hasValue())
     result["range"] = toJSON(*hover.range);
   return std::move(result);
 }
 
 //===----------------------------------------------------------------------===//
+// DocumentSymbol
+//===----------------------------------------------------------------------===//
+
+llvm::json::Value mlir::lsp::toJSON(const DocumentSymbol &symbol) {
+  llvm::json::Object result{{"name", symbol.name},
+                            {"kind", static_cast<int>(symbol.kind)},
+                            {"range", symbol.range},
+                            {"selectionRange", symbol.selectionRange}};
+
+  if (!symbol.detail.empty())
+    result["detail"] = symbol.detail;
+  if (!symbol.children.empty())
+    result["children"] = symbol.children;
+  return std::move(result);
+}
+
+//===----------------------------------------------------------------------===//
+// DocumentSymbolParams
+//===----------------------------------------------------------------------===//
+
+bool mlir::lsp::fromJSON(const llvm::json::Value &value,
+                         DocumentSymbolParams &result, llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  return o && o.map("textDocument", result.textDocument);
+}
+
+//===----------------------------------------------------------------------===//
 // DiagnosticRelatedInformation
 //===----------------------------------------------------------------------===//
 
 llvm::json::Value mlir::lsp::toJSON(const DiagnosticRelatedInformation &info) {
   return llvm::json::Object{
       {"location", info.location},
       {"message", info.message},
   };

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

if (ferror(in)) {
       return llvm::errorCodeToError(
           std::error_code(errno, std::system_category()));
     }
 
     if (succeeded(readMessage(json))) {
       if (llvm::Expected<llvm::json::Value> doc = llvm::json::parse(json)) {
         if (!handleMessage(std::move(*doc), handler))
           return llvm::Error::success();
+      } else {
+        Logger::error("JSON parse error: {0}", llvm::toString(doc.takeError()));
       }
     }
   }
   return llvm::errorCodeToError(std::make_error_code(std::errc::io_error));
 }
 
 void JSONTransport::sendMessage(llvm::json::Value msg) {
   outputBuffer.clear();

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

+// RUN: mlir-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.mlir",
+  "languageId":"mlir",
+  "version":1,
+  "text":"module {\nfunc private @foo()\n}"
+}}}
+// -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/documentSymbol","params":{
+  "textDocument":{"uri":"test:///foo.mlir"}
+}}
+//      CHECK:  "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "children": [
+// CHECK-NEXT:        {
+// CHECK-NEXT:          "kind": 12,
+// CHECK-NEXT:          "name": "foo",
+// CHECK-NEXT:          "range": {
+// CHECK-NEXT:            "end": {
+// CHECK-NEXT:              "character": {{.*}},
+// CHECK-NEXT:              "line": {{.*}}
+// CHECK-NEXT:            },
+// CHECK-NEXT:            "start": {
+// CHECK-NEXT:              "character": {{.*}},
+// CHECK-NEXT:              "line": {{.*}}
+// CHECK-NEXT:            }
+// CHECK-NEXT:          },
+// CHECK-NEXT:          "selectionRange": {
+// CHECK-NEXT:            "end": {
+// CHECK-NEXT:              "character": 4,
+// CHECK-NEXT:              "line": {{.*}}
+// CHECK-NEXT:            },
+// CHECK-NEXT:            "start": {
+// CHECK-NEXT:              "character": {{.*}},
+// CHECK-NEXT:              "line": {{.*}}
+// CHECK-NEXT:            }
+// CHECK-NEXT:          }
+// CHECK-NEXT:        }
+// CHECK-NEXT:      ],
+// CHECK-NEXT:      "kind": 3,
+// CHECK-NEXT:      "name": "<module>",
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": {{.*}},
+// CHECK-NEXT:          "line": {{.*}}
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": {{.*}},
+// CHECK-NEXT:          "line": {{.*}}
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "selectionRange": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": {{.*}},
+// CHECK-NEXT:          "line": {{.*}}
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": {{.*}},
+// CHECK-NEXT:          "line": {{.*}}
+// CHECK-NEXT:        }
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  ]
+// -----
+{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}

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

 // RUN: mlir-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:    },