[mlir][lsp] Use rootUri and rootPath for source files (#185479)

When looking up source files, consider the rootUri and rootPath (if set) and filename is non local path.
This commit is contained in:
Jacques Pienaar 2026-04-07 16:13:47 +02:00 committed by GitHub
parent 6689389174
commit 0b070ab623
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 130 additions and 31 deletions

View file

@ -215,6 +215,13 @@ struct InitializeParams {
/// The initial trace setting. If omitted trace is disabled ('off').
std::optional<TraceLevel> trace;
/// The root URI of the workspace. Is null if no folder is open.
std::optional<std::string> rootUri;
/// The root path of the workspace. Is null if no folder is open.
/// This is deprecated, use rootUri instead, but kept for more compatibility.
std::optional<std::string> rootPath;
};
/// Add support for JSON serialization.

View file

@ -335,6 +335,8 @@ bool llvm::lsp::fromJSON(const llvm::json::Value &Value,
// We deliberately don't fail if we can't parse individual fields.
O.map("capabilities", Result.capabilities);
O.map("trace", Result.trace);
O.map("rootUri", Result.rootUri);
O.map("rootPath", Result.rootPath);
mapOptOrNull(Value, "clientInfo", Result.clientInfo, Path);
return true;

View file

@ -130,6 +130,17 @@ struct LSPServer {
void LSPServer::onInitialize(const InitializeParams &params,
Callback<llvm::json::Value> reply) {
// Configure the workspace root if it was provided.
if (params.rootUri) {
llvm::Expected<URIForFile> rootURI = URIForFile::fromURI(*params.rootUri);
if (rootURI)
server.setWorkspaceRoot(rootURI->file());
else
consumeError(rootURI.takeError());
} else if (params.rootPath) {
server.setWorkspaceRoot(*params.rootPath);
}
// Send a response with the capabilities of this server.
llvm::json::Object serverCaps{
{"textDocumentSync",

View file

@ -19,7 +19,9 @@
#include "mlir/Tools/lsp-server-support/SourceMgrUtils.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/Support/Base64.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/LSP/Logging.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/SourceMgr.h"
#include <optional>
@ -34,14 +36,28 @@ static SMRange convertTokenLocToRange(SMLoc loc) {
/// Returns a language server location from the given MLIR file location.
/// `uriScheme` is the scheme to use when building new uris.
static std::optional<lsp::Location> getLocationFromLoc(StringRef uriScheme,
FileLineColLoc loc) {
static std::optional<lsp::Location>
getLocationFromLoc(StringRef uriScheme, FileLineColLoc loc,
StringRef workspaceRoot) {
StringRef filename = loc.getFilename();
SmallString<128> absPath;
// Always make the path absolute. Skip paths that start with a separator:
// prevents incorrect resolution of virtual paths used in tests on Windows.
if (!llvm::sys::path::is_absolute(filename) && !filename.starts_with("/") &&
!filename.starts_with("\\")) {
if (!workspaceRoot.empty())
llvm::sys::path::append(absPath, workspaceRoot, filename);
else
absPath = filename;
llvm::sys::fs::make_absolute(absPath);
filename = absPath;
}
llvm::Expected<lsp::URIForFile> sourceURI =
lsp::URIForFile::fromFile(loc.getFilename(), uriScheme);
lsp::URIForFile::fromFile(filename, uriScheme);
if (!sourceURI) {
llvm::lsp::Logger::error("Failed to create URI for file `{0}`: {1}",
loc.getFilename(),
llvm::toString(sourceURI.takeError()));
filename, llvm::toString(sourceURI.takeError()));
return std::nullopt;
}
@ -57,15 +73,16 @@ static std::optional<lsp::Location> getLocationFromLoc(StringRef uriScheme,
/// present, is used to filter sub locations that do not share the same uri.
static std::optional<lsp::Location>
getLocationFromLoc(llvm::SourceMgr &sourceMgr, Location loc,
StringRef uriScheme, const lsp::URIForFile *uri = nullptr) {
StringRef uriScheme, StringRef workspaceRoot,
const lsp::URIForFile *uri = nullptr) {
std::optional<lsp::Location> location;
loc->walk([&](Location nestedLoc) {
FileLineColLoc fileLoc = dyn_cast<FileLineColLoc>(nestedLoc);
auto fileLoc = dyn_cast<FileLineColLoc>(nestedLoc);
if (!fileLoc)
return WalkResult::advance();
std::optional<lsp::Location> sourceLoc =
getLocationFromLoc(uriScheme, fileLoc);
getLocationFromLoc(uriScheme, fileLoc, workspaceRoot);
if (sourceLoc && (!uri || sourceLoc->uri == *uri)) {
location = *sourceLoc;
SMLoc loc = sourceMgr.FindLocForLineAndColumn(
@ -90,7 +107,8 @@ getLocationFromLoc(llvm::SourceMgr &sourceMgr, Location loc,
/// contained within the given URI.
static void collectLocationsFromLoc(Location loc,
std::vector<lsp::Location> &locations,
const lsp::URIForFile &uri) {
const lsp::URIForFile &uri,
StringRef workspaceRoot) {
SetVector<Location> visitedLocs;
loc->walk([&](Location nestedLoc) {
FileLineColLoc fileLoc = dyn_cast<FileLineColLoc>(nestedLoc);
@ -98,7 +116,7 @@ static void collectLocationsFromLoc(Location loc,
return WalkResult::advance();
std::optional<lsp::Location> sourceLoc =
getLocationFromLoc(uri.scheme(), fileLoc);
getLocationFromLoc(uri.scheme(), fileLoc, workspaceRoot);
if (sourceLoc && sourceLoc->uri != uri)
locations.push_back(*sourceLoc);
return WalkResult::advance();
@ -193,7 +211,8 @@ static void printDefBlockName(raw_ostream &os,
/// Convert the given MLIR diagnostic to the LSP form.
static lsp::Diagnostic getLspDiagnoticFromDiag(llvm::SourceMgr &sourceMgr,
Diagnostic &diag,
const lsp::URIForFile &uri) {
const lsp::URIForFile &uri,
StringRef workspaceRoot) {
lsp::Diagnostic lspDiag;
lspDiag.source = "mlir";
@ -205,8 +224,8 @@ static lsp::Diagnostic getLspDiagnoticFromDiag(llvm::SourceMgr &sourceMgr,
// TODO: For simplicity, we just grab the first one. It may be likely that we
// will need a more interesting heuristic here.'
StringRef uriScheme = uri.scheme();
std::optional<lsp::Location> lspLocation =
getLocationFromLoc(sourceMgr, diag.getLocation(), uriScheme, &uri);
std::optional<lsp::Location> lspLocation = getLocationFromLoc(
sourceMgr, diag.getLocation(), uriScheme, workspaceRoot, &uri);
if (lspLocation)
lspDiag.range = lspLocation->range;
@ -230,8 +249,8 @@ static lsp::Diagnostic getLspDiagnoticFromDiag(llvm::SourceMgr &sourceMgr,
std::vector<llvm::lsp::DiagnosticRelatedInformation> relatedDiags;
for (Diagnostic &note : diag.getNotes()) {
lsp::Location noteLoc;
if (std::optional<lsp::Location> loc =
getLocationFromLoc(sourceMgr, note.getLocation(), uriScheme))
if (std::optional<lsp::Location> loc = getLocationFromLoc(
sourceMgr, note.getLocation(), uriScheme, workspaceRoot))
noteLoc = *loc;
else
noteLoc.uri = uri;
@ -252,7 +271,8 @@ namespace {
/// document.
struct MLIRDocument {
MLIRDocument(MLIRContext &context, const lsp::URIForFile &uri,
StringRef contents, std::vector<lsp::Diagnostic> &diagnostics);
StringRef contents, StringRef workspaceRoot,
std::vector<lsp::Diagnostic> &diagnostics);
MLIRDocument(const MLIRDocument &) = delete;
MLIRDocument &operator=(const MLIRDocument &) = delete;
@ -337,14 +357,19 @@ struct MLIRDocument {
/// The source manager containing the contents of the input file.
llvm::SourceMgr sourceMgr;
/// The workspace root of the server.
std::string workspaceRoot;
};
} // namespace
MLIRDocument::MLIRDocument(MLIRContext &context, const lsp::URIForFile &uri,
StringRef contents,
std::vector<lsp::Diagnostic> &diagnostics) {
StringRef contents, StringRef workspaceRoot,
std::vector<lsp::Diagnostic> &diagnostics)
: workspaceRoot(workspaceRoot.str()) {
ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diag) {
diagnostics.push_back(getLspDiagnoticFromDiag(sourceMgr, diag, uri));
diagnostics.push_back(
getLspDiagnoticFromDiag(sourceMgr, diag, uri, workspaceRoot));
});
// Try to parsed the given IR string.
@ -387,14 +412,17 @@ void MLIRDocument::getLocationsOf(const lsp::URIForFile &uri,
// Check all definitions related to operations.
for (const AsmParserState::OperationDefinition &op : asmState.getOpDefs()) {
if (contains(op.loc, posLoc))
return collectLocationsFromLoc(op.op->getLoc(), locations, uri);
return collectLocationsFromLoc(op.op->getLoc(), locations, uri,
workspaceRoot);
for (const auto &result : op.resultGroups)
if (containsPosition(result.definition))
return collectLocationsFromLoc(op.op->getLoc(), locations, uri);
return collectLocationsFromLoc(op.op->getLoc(), locations, uri,
workspaceRoot);
for (const auto &symUse : op.symbolUses) {
if (contains(symUse, posLoc)) {
locations.emplace_back(uri, sourceMgr, op.loc);
return collectLocationsFromLoc(op.op->getLoc(), locations, uri);
return collectLocationsFromLoc(op.op->getLoc(), locations, uri,
workspaceRoot);
}
}
}
@ -968,8 +996,10 @@ namespace {
struct MLIRTextFileChunk {
MLIRTextFileChunk(MLIRContext &context, uint64_t lineOffset,
const lsp::URIForFile &uri, StringRef contents,
StringRef workspaceRoot,
std::vector<lsp::Diagnostic> &diagnostics)
: lineOffset(lineOffset), document(context, uri, contents, diagnostics) {}
: lineOffset(lineOffset),
document(context, uri, contents, workspaceRoot, diagnostics) {}
/// Adjust the line number of the given range to anchor at the beginning of
/// the file, instead of the beginning of this chunk.
@ -998,6 +1028,7 @@ class MLIRTextFile {
public:
MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents,
int64_t version, lsp::DialectRegistryFn registryFn,
StringRef workspaceRoot,
std::vector<lsp::Diagnostic> &diagnostics);
/// Return the current version of this text file.
@ -1047,6 +1078,7 @@ private:
MLIRTextFile::MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents,
int64_t version, lsp::DialectRegistryFn registryFn,
StringRef workspaceRoot,
std::vector<lsp::Diagnostic> &diagnostics)
: context(registryFn(uri), MLIRContext::Threading::DISABLED),
contents(fileContents.str()), version(version) {
@ -1056,13 +1088,14 @@ MLIRTextFile::MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents,
SmallVector<StringRef, 8> subContents;
StringRef(contents).split(subContents, kDefaultSplitMarker);
chunks.emplace_back(std::make_unique<MLIRTextFileChunk>(
context, /*lineOffset=*/0, uri, subContents.front(), diagnostics));
context, /*lineOffset=*/0, uri, subContents.front(), workspaceRoot,
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>(context, lineOffset, uri,
docContents, diagnostics);
auto chunk = std::make_unique<MLIRTextFileChunk>(
context, lineOffset, uri, docContents, workspaceRoot, diagnostics);
lineOffset += docContents.count('\n');
// Adjust locations used in diagnostics to account for the offset from the
@ -1271,6 +1304,9 @@ struct lsp::MLIRServer::Impl {
/// The files held by the server, mapped by their URI file name.
llvm::StringMap<std::unique_ptr<MLIRTextFile>> files;
/// The workspace root of the server.
std::string workspaceRoot;
};
//===----------------------------------------------------------------------===//
@ -1284,8 +1320,9 @@ lsp::MLIRServer::~MLIRServer() = default;
void lsp::MLIRServer::addOrUpdateDocument(
const URIForFile &uri, StringRef contents, int64_t version,
std::vector<llvm::lsp::Diagnostic> &diagnostics) {
impl->files[uri.file()] = std::make_unique<MLIRTextFile>(
uri, contents, version, impl->registryFn, diagnostics);
impl->files[uri.file()] =
std::make_unique<MLIRTextFile>(uri, contents, version, impl->registryFn,
impl->workspaceRoot, diagnostics);
}
std::optional<int64_t> lsp::MLIRServer::removeDocument(const URIForFile &uri) {
@ -1407,3 +1444,7 @@ lsp::MLIRServer::convertToBytecode(const URIForFile &uri) {
}
return fileIt->second->convertToBytecode();
}
void lsp::MLIRServer::setWorkspaceRoot(StringRef root) {
impl->workspaceRoot = root.str();
}

View file

@ -87,6 +87,9 @@ public:
llvm::Expected<MLIRConvertBytecodeResult>
convertToBytecode(const URIForFile &uri);
/// Set the workspace root for the server.
void setWorkspaceRoot(StringRef root);
private:
struct Impl;

View file

@ -27,7 +27,7 @@
// CHECK-NEXT: "source": "mlir"
// CHECK-NEXT: }
// CHECK-NEXT: ],
// CHECK-NEXT: "uri": "test:///foo.mlir",
// CHECK-NEXT: "uri": "test:///{{([a-zA-Z]:/)?}}foo.mlir",
// CHECK-NEXT: "version": 1
// CHECK-NEXT: }
// -----
@ -57,7 +57,7 @@
// CHECK-NEXT: "source": "mlir"
// CHECK-NEXT: }
// CHECK-NEXT: ],
// CHECK-NEXT: "uri": "test:///foo.mlir",
// CHECK-NEXT: "uri": "test:///{{([a-zA-Z]:/)?}}foo.mlir",
// CHECK-NEXT: "version": 1
// CHECK-NEXT: }
// -----

View file

@ -0,0 +1,35 @@
// RUN: mlir-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootUri":"test:///root/","capabilities":{},"trace":"off"}}
// -----
{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
"uri":"test:///foo.mlir",
"languageId":"mlir",
"version":1,
"text":"func.func @foo() {\n %0 = arith.constant true loc(\"indexer.cc\":10:10)\n return\n}"
}}}
// -----
{"jsonrpc":"2.0","id":1,"method":"textDocument/definition","params":{
"textDocument":{"uri":"test:///foo.mlir"},
"position":{"line":1,"character":12}
}}
// CHECK: "id": 1
// CHECK-NEXT: "jsonrpc": "2.0",
// CHECK-NEXT: "result": [
// CHECK-NEXT: {
// CHECK-NEXT: "range": {
// CHECK-NEXT: "end": {
// CHECK-NEXT: "character": 9,
// CHECK-NEXT: "line": 9
// CHECK-NEXT: },
// CHECK-NEXT: "start": {
// CHECK-NEXT: "character": 9,
// CHECK-NEXT: "line": 9
// CHECK-NEXT: }
// CHECK-NEXT: },
// CHECK-NEXT: "uri": "test:///{{([a-zA-Z]:/)?}}root/indexer.cc"
// CHECK-NEXT: }
// CHECK-NEXT: ]
// -----
{"jsonrpc":"2.0","id":3,"method":"shutdown"}
// -----
{"jsonrpc":"2.0","method":"exit"}

View file

@ -2,7 +2,7 @@
"name": "vscode-mlir",
"displayName": "MLIR",
"description": "MLIR Language Extension",
"version": "0.0.14",
"version": "0.0.15",
"publisher": "llvm-vs-code-extensions",
"homepage": "https://mlir.llvm.org/",
"icon": "icon.png",