[mlir-lsp-server] Add support for processing split files

[rriddle]

Summary

MLIR tools very commonly use // ----- to split a file into distinct sub documents, that are processed separately. This revision adds support to mlir-lsp-server for splitting MLIR files based on this sigil, and processing them separately.

Depends On D102644

[jpienaar]

Nice, thanks!

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

576

Should this be a reference or pointer? Seems like an expensive struct to accidentally copy.

630

Perhaps add a TODO that we make this some constant to ensure these remain in sync with mlir-opt?

mlir/test/mlir-lsp-server/definition-split-file.test

2

Could you add a comment explaining the test?

[rriddle]

update

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

576

Just deleted the copy constructor on MLIRDocument.

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

 //===----------------------------------------------------------------------===//
 // MLIRDocument
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// 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 &registry,
                std::vector<lsp::Diagnostic> &diagnostics);
+  MLIRDocument(const MLIRDocument &) = delete;
+  MLIRDocument &operator=(const MLIRDocument &) = delete;
 
   //===--------------------------------------------------------------------===//
   // Definitions and References
   //===--------------------------------------------------------------------===//
 
   void getLocationsOf(const lsp::URIForFile &uri, const lsp::Position &defPos,
                       std::vector<lsp::Location> &locations);
   void findReferencesOf(const lsp::URIForFile &uri, const lsp::Position &pos,
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);
 
-  /// The version of this document.
-  int64_t version;
-
   /// 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.
   Block parsedIR;
 
   /// The source manager containing the contents of the input file.
   llvm::SourceMgr sourceMgr;
 };
 } // namespace
 
 MLIRDocument::MLIRDocument(const lsp::URIForFile &uri, StringRef contents,
-                           int64_t version, DialectRegistry &registry,
+                           DialectRegistry &registry,
                            std::vector<lsp::Diagnostic> &diagnostics)
-    : version(version), context(registry) {
+    : context(registry) {
   context.allowUnregisteredDialects();
   ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diag) {
     diagnostics.push_back(getLspDiagnoticFromDiag(diag, uri));
   });
 
   // Try to parsed the given IR string.
   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
   if (!memBuffer) {
lsp::Hover MLIRDocument::buildHoverForBlockArgument(
   printDefBlockName(os, block);
   os << "\n\nArgument #" << arg.getArgNumber() << "\n\n"
      << "Type: `" << arg.getType() << "`\n\n";
 
   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 &registry,
+                    std::vector<lsp::Diagnostic> &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 &registry,
+               std::vector<lsp::Diagnostic> &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<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);
+
+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<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) {
+  // 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));
+
+  uint64_t lineOffset = subContents.front().count('\n');
+  for (StringRef docContents : llvm::drop_begin(subContents)) {
+    unsigned currentNumDiags = diagnostics.size();
+    auto chunk = std::make_unique<MLIRTextFileChunk>(
+        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<lsp::Location> &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<lsp::Location> &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<lsp::Hover> MLIRTextFile::findHover(const lsp::URIForFile &uri,
+                                             lsp::Position hoverPos) {
+  MLIRTextFileChunk &chunk = getChunkFor(hoverPos);
+  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;
+}
+
+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<uint64_t>(pos.line) < chunk->lineOffset;
+      });
+  MLIRTextFileChunk &chunk = it == chunks.end() ? *chunks.back() : **(--it);
+  pos.line -= chunk.lineOffset;
+  return chunk;
+}
+
+//===----------------------------------------------------------------------===//
 // MLIRServer::Impl
 //===----------------------------------------------------------------------===//
 
 struct lsp::MLIRServer::Impl {
   Impl(DialectRegistry &registry) : registry(registry) {}
 
   /// The registry containing dialects that can be recognized in parsed .mlir
   /// files.
   DialectRegistry &registry;
 
-  /// The documents held by the server, mapped by their URI file name.
-  llvm::StringMap<std::unique_ptr<MLIRDocument>> documents;
+  /// The files held by the server, mapped by their URI file name.
+  llvm::StringMap<std::unique_ptr<MLIRTextFile>> files;
 };
 
 //===----------------------------------------------------------------------===//
 // MLIRServer
 //===----------------------------------------------------------------------===//
 
 lsp::MLIRServer::MLIRServer(DialectRegistry &registry)
     : impl(std::make_unique<Impl>(registry)) {}
 lsp::MLIRServer::~MLIRServer() {}
 
 void lsp::MLIRServer::addOrUpdateDocument(
     const URIForFile &uri, StringRef contents, int64_t version,
     std::vector<Diagnostic> &diagnostics) {
-  impl->documents[uri.file()] = std::make_unique<MLIRDocument>(
+  impl->files[uri.file()] = std::make_unique<MLIRTextFile>(
       uri, contents, version, impl->registry, diagnostics);
 }
 
 Optional<int64_t> 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<Location> &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<Location> &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::Hover> 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;
 }

mlir/test/mlir-lsp-server/definition-split-file.test

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