diff --git a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp --- a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp +++ b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp @@ -252,9 +252,11 @@ /// This class represents all of the information pertaining to a specific MLIR /// document. struct MLIRDocument { - MLIRDocument(const lsp::URIForFile &uri, StringRef contents, int64_t version, + MLIRDocument(const lsp::URIForFile &uri, StringRef contents, DialectRegistry ®istry, std::vector &diagnostics); + MLIRDocument(const MLIRDocument &) = delete; + MLIRDocument &operator=(const MLIRDocument &) = delete; //===--------------------------------------------------------------------===// // Definitions and References @@ -283,9 +285,6 @@ buildHoverForBlockArgument(llvm::SMRange hoverRange, BlockArgument arg, const AsmParserState::BlockDefinition &block); - /// The version of this document. - int64_t version; - /// The context used to hold the state contained by the parsed document. MLIRContext context; @@ -302,9 +301,9 @@ } // namespace MLIRDocument::MLIRDocument(const lsp::URIForFile &uri, StringRef contents, - int64_t version, DialectRegistry ®istry, + DialectRegistry ®istry, std::vector &diagnostics) - : version(version), context(registry) { + : context(registry) { context.allowUnregisteredDialects(); ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diag) { diagnostics.push_back(getLspDiagnoticFromDiag(diag, uri)); @@ -548,6 +547,170 @@ return hover; } +//===----------------------------------------------------------------------===// +// 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 ®istry, + std::vector &diagnostics) + : lineOffset(lineOffset), document(uri, contents, registry, diagnostics) { + } + + /// Adjust the line number of the given range to anchor at the beginning of + /// the file, instead of the beginning of this chunk. + void adjustLocForChunkOffset(lsp::Range &range) { + adjustLocForChunkOffset(range.start); + adjustLocForChunkOffset(range.end); + } + /// Adjust the line number of the given position to anchor at the beginning of + /// the file, instead of the beginning of this chunk. + void adjustLocForChunkOffset(lsp::Position &pos) { pos.line += lineOffset; } + + /// The line offset of this chunk from the beginning of the file. + uint64_t lineOffset; + /// The document referred to by this chunk. + MLIRDocument document; +}; +} // namespace + +//===----------------------------------------------------------------------===// +// MLIRTextFile +//===----------------------------------------------------------------------===// + +namespace { +/// This class represents a text file containing one or more MLIR documents. +class MLIRTextFile { +public: + MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents, + int64_t version, DialectRegistry ®istry, + std::vector &diagnostics); + + /// Return the current version of this text file. + int64_t getVersion() const { return version; } + + //===--------------------------------------------------------------------===// + // LSP Queries + //===--------------------------------------------------------------------===// + + void getLocationsOf(const lsp::URIForFile &uri, lsp::Position defPos, + std::vector &locations); + void findReferencesOf(const lsp::URIForFile &uri, lsp::Position pos, + std::vector &references); + Optional findHover(const lsp::URIForFile &uri, + lsp::Position hoverPos); + +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 chunks of this file. The order of these chunks is the order in which + /// they appear in the text file. + std::vector> chunks; +}; +} // namespace + +MLIRTextFile::MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents, + int64_t version, DialectRegistry ®istry, + std::vector &diagnostics) + : contents(fileContents.str()), version(version) { + // 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 subContents; + StringRef(contents).split(subContents, "// -----"); + chunks.emplace_back(std::make_unique( + /*lineOffset=*/0, uri, subContents.front(), registry, diagnostics)); + + uint64_t lineOffset = subContents.front().count('\n'); + for (StringRef docContents : llvm::drop_begin(subContents)) { + unsigned currentNumDiags = diagnostics.size(); + auto chunk = std::make_unique( + lineOffset, uri, docContents, registry, diagnostics); + lineOffset += docContents.count('\n'); + + // Adjust locations used in diagnostics to account for the offset from the + // beginning of the file. + for (lsp::Diagnostic &diag : + llvm::drop_begin(diagnostics, currentNumDiags)) { + chunk->adjustLocForChunkOffset(diag.range); + + 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)); + } +} + +void MLIRTextFile::getLocationsOf(const lsp::URIForFile &uri, + lsp::Position defPos, + std::vector &locations) { + MLIRTextFileChunk &chunk = getChunkFor(defPos); + chunk.document.getLocationsOf(uri, defPos, locations); + + // Adjust any locations within this file for the offset of this chunk. + if (chunk.lineOffset == 0) + return; + for (lsp::Location &loc : locations) + if (loc.uri == uri) + chunk.adjustLocForChunkOffset(loc.range); +} + +void MLIRTextFile::findReferencesOf(const lsp::URIForFile &uri, + lsp::Position pos, + std::vector &references) { + MLIRTextFileChunk &chunk = getChunkFor(pos); + chunk.document.findReferencesOf(uri, pos, references); + + // Adjust any locations within this file for the offset of this chunk. + if (chunk.lineOffset == 0) + return; + for (lsp::Location &loc : references) + if (loc.uri == uri) + chunk.adjustLocForChunkOffset(loc.range); +} + +Optional MLIRTextFile::findHover(const lsp::URIForFile &uri, + lsp::Position hoverPos) { + MLIRTextFileChunk &chunk = getChunkFor(hoverPos); + Optional 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; +} + +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) { + return static_cast(pos.line) < chunk->lineOffset; + }); + MLIRTextFileChunk &chunk = it == chunks.end() ? *chunks.back() : **(--it); + pos.line -= chunk.lineOffset; + return chunk; +} + //===----------------------------------------------------------------------===// // MLIRServer::Impl //===----------------------------------------------------------------------===// @@ -559,8 +722,8 @@ /// files. DialectRegistry ®istry; - /// The documents held by the server, mapped by their URI file name. - llvm::StringMap> documents; + /// The files held by the server, mapped by their URI file name. + llvm::StringMap> files; }; //===----------------------------------------------------------------------===// @@ -574,40 +737,40 @@ void lsp::MLIRServer::addOrUpdateDocument( const URIForFile &uri, StringRef contents, int64_t version, std::vector &diagnostics) { - impl->documents[uri.file()] = std::make_unique( + impl->files[uri.file()] = std::make_unique( uri, contents, version, impl->registry, diagnostics); } Optional lsp::MLIRServer::removeDocument(const URIForFile &uri) { - auto it = impl->documents.find(uri.file()); - if (it == impl->documents.end()) + auto it = impl->files.find(uri.file()); + if (it == impl->files.end()) return llvm::None; - int64_t version = it->second->version; - impl->documents.erase(it); + int64_t version = it->second->getVersion(); + impl->files.erase(it); return version; } void lsp::MLIRServer::getLocationsOf(const URIForFile &uri, const Position &defPos, std::vector &locations) { - auto fileIt = impl->documents.find(uri.file()); - if (fileIt != impl->documents.end()) + auto fileIt = impl->files.find(uri.file()); + if (fileIt != impl->files.end()) fileIt->second->getLocationsOf(uri, defPos, locations); } void lsp::MLIRServer::findReferencesOf(const URIForFile &uri, const Position &pos, std::vector &references) { - auto fileIt = impl->documents.find(uri.file()); - if (fileIt != impl->documents.end()) + auto fileIt = impl->files.find(uri.file()); + if (fileIt != impl->files.end()) fileIt->second->findReferencesOf(uri, pos, references); } Optional lsp::MLIRServer::findHover(const URIForFile &uri, const Position &hoverPos) { - auto fileIt = impl->documents.find(uri.file()); - if (fileIt != impl->documents.end()) + auto fileIt = impl->files.find(uri.file()); + if (fileIt != impl->files.end()) return fileIt->second->findHover(uri, hoverPos); return llvm::None; } diff --git a/mlir/test/mlir-lsp-server/definition-split-file.test b/mlir/test/mlir-lsp-server/definition-split-file.test new file mode 100644 --- /dev/null +++ b/mlir/test/mlir-lsp-server/definition-split-file.test @@ -0,0 +1,37 @@ +// RUN: mlir-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s +// This test checks support for split files by attempting to find the definition +// of a symbol in a split file. The interesting part of this test is that the +// file chunk before the one we are looking for the definition in has an error. +{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"mlir","capabilities":{},"trace":"off"}} +// ----- +{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{ + "uri":"test:///foo.mlir", + "languageId":"mlir", + "version":1, + "text":"func @foo() -> {}\n// -----\nfunc @foo() -> i1 {\n%value = constant true\nreturn %value : i1\n}" +}}} +// ----- +{"jsonrpc":"2.0","id":1,"method":"textDocument/definition","params":{ + "textDocument":{"uri":"test:///foo.mlir"}, + "position":{"line":4,"character":12} +}} +// CHECK: "id": 1 +// CHECK-NEXT: "jsonrpc": "2.0", +// CHECK-NEXT: "result": [ +// CHECK-NEXT: { +// CHECK-NEXT: "range": { +// CHECK-NEXT: "end": { +// CHECK-NEXT: "character": 6, +// CHECK-NEXT: "line": 3 +// CHECK-NEXT: }, +// CHECK-NEXT: "start": { +// CHECK-NEXT: "character": 1, +// CHECK-NEXT: "line": 3 +// CHECK-NEXT: } +// CHECK-NEXT: }, +// CHECK-NEXT: "uri": "{{.*}}/foo.mlir" +// CHECK-NEXT: } +// ----- +{"jsonrpc":"2.0","id":3,"method":"shutdown"} +// ----- +{"jsonrpc":"2.0","method":"exit"}