Skip to content

[mlir-lsp] Abstract input and output of the JSONTransport #129320

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
npanchen merged 3 commits into from
Mar 14, 2025
Merged

[mlir-lsp] Abstract input and output of the JSONTransport #129320

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
dcb202d
[mlir-lsp] Abstract input and output of the `JSONTransport`
npanchen Feb 28, 2025
b62319d
Addressed comments
npanchen Mar 3, 2025
dd3bb29
Merge branch 'main' into npanchen/nfc_json_transport_inout_abstraction
npanchen Mar 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 78 additions & 21 deletions mlir/include/mlir/Tools/lsp-server-support/Transport.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,14 +43,86 @@ enum JSONStreamStyle {
Delimited
};

/// An abstract class used by the JSONTransport to read JSON message.
class JSONTransportInput {
public:
explicit JSONTransportInput(JSONStreamStyle style = JSONStreamStyle::Standard)
: style(style) {}
virtual ~JSONTransportInput() = default;

virtual bool getError() const = 0;
Copy link
Contributor

@River707 River707 Feb 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hasError? getError is kind of confusing given the return type.

Copy link
Contributor Author

@npanchen npanchen Feb 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you're right, I wanted to return error first, but it won't be easy if error type is different for custom implementation.

virtual bool isEndOfInput() const = 0;

/// Read in a message from the input stream.
virtual LogicalResult readMessage(std::string &json) {
Copy link
Contributor

@River707 River707 Feb 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why virtual here? If we already accept and dispatch on the stream style, I don't see this should be overridable.

Copy link
Contributor Author

@npanchen npanchen Feb 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

wanted to give way more freedom to whoever is going to inherit from that class with a possibility to use different stream style.
But for now I agree, that can stay as non-virtual

return style == JSONStreamStyle::Delimited ? readDelimitedMessage(json)
: readStandardMessage(json);
}
virtual LogicalResult readDelimitedMessage(std::string &json) = 0;
virtual LogicalResult readStandardMessage(std::string &json) = 0;

private:
/// The JSON stream style to use.
JSONStreamStyle style;
};

/// An abstract class used by the JSONTransport to write JSON messages.
class JSONTransportOutput {
public:
explicit JSONTransportOutput() = default;
virtual ~JSONTransportOutput() = default;

virtual void sendMessage(const llvm::json::Value &msg) = 0;
};

/// Concrete implementation of the JSONTransportInput that reads from a file.
class JSONTransportInputOverFile : public JSONTransportInput {
public:
explicit JSONTransportInputOverFile(
std::FILE *in, JSONStreamStyle style = JSONStreamStyle::Standard)
: JSONTransportInput(style), in(in) {}

bool getError() const final { return ferror(in); }
bool isEndOfInput() const final { return feof(in); }

LogicalResult readDelimitedMessage(std::string &json) final;
LogicalResult readStandardMessage(std::string &json) final;

private:
std::FILE *in;
};

/// Concrete implementation of the JSONTransportOutput that writes to a stream.
class JSONTransportOutputOverStream : public JSONTransportOutput {
Copy link
Contributor

@River707 River707 Feb 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a bit confused why this is being abstracted out. Why can't your custom implementation use raw_ostream?

Copy link
Contributor Author

@npanchen npanchen Feb 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the main reason was to keep that output class simple and match to what json transport is expecting. User may decide to either implement other functionality such as colors, seek etc or not to implement them.

Copy link
Contributor

@River707 River707 Feb 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With these kinds of things I'd really prefer to avoid abstracting until we need to. Most of the time it ends up not being necessary. For example, why would you need/want colors for a transport class? You generally aren't observing the data being transferred outside of logging/testing (or shouldn't be at least).

Copy link
Contributor Author

@npanchen npanchen Feb 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I do agree that raw_ostream seems sufficient even for extreme case, such as network communication, but it does provide unnecessary functionality that isn't needed. That is, it's like to use a cannon to kill a fly.

Shouldn't output be abstracted with minimal methods ? For example, I see that MCStreamer is quite similar to our discussion (with exception of many methods inside), which is not inherited from raw_ostream, MCAsmStreamer, a concrete implementation, does use raw_ostream to output.

That said, I do understand your point about "not being necessary", at the same time I think it's better to hide footgun (unneeded and untested methods of raw_ostream) away. API-wise, user still can construct JSONTransport using raw_ostream for convenience.

Copy link
Contributor

@River707 River707 Feb 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's not the goal of the output stream abstraction, implementations of raw_ostream describe the capabilities of the stream, it's up to the user to decide which capabilities it wants to take advantage of. That is an example of a good abstraction, because it's reusable across different use-cases and domains. Consider, for example, if we needed to provide an output abstraction for file output for everything that may write to a file. Consider the ease of use today for raw_ostream, if an implementation already exists for your output type (which is the case most often) you have nothing else to do to take advantage of the API.

Copy link
Contributor Author

@npanchen npanchen Mar 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok, I will keep it raw_ostream, will look for some other approach if that won't work for me

public:
explicit JSONTransportOutputOverStream(raw_ostream &out,
bool prettyOutput = false)
: JSONTransportOutput(), out(out), prettyOutput(prettyOutput) {}

/// Writes the given message to the output stream.
void sendMessage(const llvm::json::Value &msg) final;

private:
SmallVector<char, 0> outputBuffer;
raw_ostream &out;
/// If the output JSON should be formatted for easier readability.
bool prettyOutput;
};

/// A transport class that performs the JSON-RPC communication with the LSP
/// client.
class JSONTransport {
public:
JSONTransport(std::unique_ptr<JSONTransportInput> in,
std::unique_ptr<JSONTransportOutput> out)
: in(std::move(in)), out(std::move(out)) {}

JSONTransport(std::FILE *in, raw_ostream &out,
JSONStreamStyle style = JSONStreamStyle::Standard,
bool prettyOutput = false)
: in(in), out(out), style(style), prettyOutput(prettyOutput) {}
: in(std::make_unique<JSONTransportInputOverFile>(in, style)),
out(std::make_unique<JSONTransportOutputOverStream>(out,
prettyOutput)) {}

/// The following methods are used to send a message to the LSP client.
void notify(StringRef method, llvm::json::Value params);
Expand All @@ -60,30 +132,15 @@ class JSONTransport {
/// Start executing the JSON-RPC transport.
llvm::Error run(MessageHandler &handler);

private:
/// Dispatches the given incoming json message to the message handler.
bool handleMessage(llvm::json::Value msg, MessageHandler &handler);
/// Writes the given message to the output stream.
void sendMessage(llvm::json::Value msg);

/// Read in a message from the input stream.
LogicalResult readMessage(std::string &json) {
return style == JSONStreamStyle::Delimited ? readDelimitedMessage(json)
: readStandardMessage(json);
}
LogicalResult readDelimitedMessage(std::string &json);
LogicalResult readStandardMessage(std::string &json);
private:
/// The input to read a message from.
std::unique_ptr<JSONTransportInput> in;

/// An output buffer used when building output messages.
SmallVector<char, 0> outputBuffer;
/// The input file stream.
std::FILE *in;
/// The output file stream.
raw_ostream &out;
/// The JSON stream style to use.
JSONStreamStyle style;
/// If the output JSON should be formatted for easier readability.
bool prettyOutput;
/// The output to send a messages to.
std::unique_ptr<JSONTransportOutput> out;
};

//===----------------------------------------------------------------------===//
Expand Down
22 changes: 12 additions & 10 deletions mlir/lib/Tools/lsp-server-support/Transport.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,15 +175,15 @@ llvm::Error decodeError(const llvm::json::Object &o) {
}

void JSONTransport::notify(StringRef method, llvm::json::Value params) {
sendMessage(llvm::json::Object{
out->sendMessage(llvm::json::Object{
{"jsonrpc", "2.0"},
{"method", method},
{"params", std::move(params)},
});
}
void JSONTransport::call(StringRef method, llvm::json::Value params,
llvm::json::Value id) {
sendMessage(llvm::json::Object{
out->sendMessage(llvm::json::Object{
{"jsonrpc", "2.0"},
{"id", std::move(id)},
{"method", method},
Expand All @@ -193,14 +193,14 @@ void JSONTransport::call(StringRef method, llvm::json::Value params,
void JSONTransport::reply(llvm::json::Value id,
llvm::Expected<llvm::json::Value> result) {
if (result) {
return sendMessage(llvm::json::Object{
return out->sendMessage(llvm::json::Object{
{"jsonrpc", "2.0"},
{"id", std::move(id)},
{"result", std::move(*result)},
});
}

sendMessage(llvm::json::Object{
out->sendMessage(llvm::json::Object{
{"jsonrpc", "2.0"},
{"id", std::move(id)},
{"error", encodeError(result.takeError())},
Expand All @@ -209,13 +209,13 @@ void JSONTransport::reply(llvm::json::Value id,

llvm::Error JSONTransport::run(MessageHandler &handler) {
std::string json;
while (!feof(in)) {
if (ferror(in)) {
while (!in->isEndOfInput()) {
if (in->getError()) {
return llvm::errorCodeToError(
std::error_code(errno, std::system_category()));
}

if (succeeded(readMessage(json))) {
if (succeeded(in->readMessage(json))) {
if (llvm::Expected<llvm::json::Value> doc = llvm::json::parse(json)) {
if (!handleMessage(std::move(*doc), handler))
return llvm::Error::success();
Expand All @@ -227,7 +227,7 @@ llvm::Error JSONTransport::run(MessageHandler &handler) {
return llvm::errorCodeToError(std::make_error_code(std::errc::io_error));
}

void JSONTransport::sendMessage(llvm::json::Value msg) {
void JSONTransportOutputOverStream::sendMessage(const llvm::json::Value &msg) {
outputBuffer.clear();
llvm::raw_svector_ostream os(outputBuffer);
os << llvm::formatv(prettyOutput ? "{0:2}\n" : "{0}", msg);
Expand Down Expand Up @@ -303,7 +303,8 @@ LogicalResult readLine(std::FILE *in, SmallVectorImpl<char> &out) {
// Returns std::nullopt when:
// - ferror(), feof(), or shutdownRequested() are set.
// - Content-Length is missing or empty (protocol error)
LogicalResult JSONTransport::readStandardMessage(std::string &json) {
LogicalResult
JSONTransportInputOverFile::readStandardMessage(std::string &json) {
// A Language Server Protocol message starts with a set of HTTP headers,
// delimited by \r\n, and terminated by an empty line (\r\n).
unsigned long long contentLength = 0;
Expand Down Expand Up @@ -349,7 +350,8 @@ LogicalResult JSONTransport::readStandardMessage(std::string &json) {
/// This is a testing path, so favor simplicity over performance here.
/// When returning failure: feof(), ferror(), or shutdownRequested() will be
/// set.
LogicalResult JSONTransport::readDelimitedMessage(std::string &json) {
LogicalResult
JSONTransportInputOverFile::readDelimitedMessage(std::string &json) {
json.clear();
llvm::SmallString<128> line;
while (succeeded(readLine(in, line))) {
Expand Down