From dcb202d32bc3d93084d37be5269d166a9c17926c Mon Sep 17 00:00:00 2001 From: Kolya Panchenko Date: Fri, 28 Feb 2025 15:34:56 -0500 Subject: [PATCH 1/2] [mlir-lsp] Abstract input and output of the `JSONTransport` The patch abstracts sending and receiving json messages of `JSONTransport` to allow custom implementation of them. For example, one concrete implementation can use pipes without a need to convert file descriptor to a `FILE` object. --- .../mlir/Tools/lsp-server-support/Transport.h | 99 +++++++++++++++---- .../Tools/lsp-server-support/Transport.cpp | 22 +++-- 2 files changed, 90 insertions(+), 31 deletions(-) diff --git a/mlir/include/mlir/Tools/lsp-server-support/Transport.h b/mlir/include/mlir/Tools/lsp-server-support/Transport.h index 6843bc76ab9dc..78f9d48e4153a 100644 --- a/mlir/include/mlir/Tools/lsp-server-support/Transport.h +++ b/mlir/include/mlir/Tools/lsp-server-support/Transport.h @@ -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; + virtual bool isEndOfInput() const = 0; + + /// Read in a message from the input stream. + virtual LogicalResult readMessage(std::string &json) { + 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 { +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 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 in, + std::unique_ptr 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(in, style)), + out(std::make_unique(out, + prettyOutput)) {} /// The following methods are used to send a message to the LSP client. void notify(StringRef method, llvm::json::Value params); @@ -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 in; - /// An output buffer used when building output messages. - SmallVector 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 out; }; //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Tools/lsp-server-support/Transport.cpp b/mlir/lib/Tools/lsp-server-support/Transport.cpp index ad8308f69aead..d6806bf5cfe13 100644 --- a/mlir/lib/Tools/lsp-server-support/Transport.cpp +++ b/mlir/lib/Tools/lsp-server-support/Transport.cpp @@ -175,7 +175,7 @@ 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)}, @@ -183,7 +183,7 @@ void JSONTransport::notify(StringRef method, llvm::json::Value 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}, @@ -193,14 +193,14 @@ void JSONTransport::call(StringRef method, llvm::json::Value params, void JSONTransport::reply(llvm::json::Value id, llvm::Expected 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())}, @@ -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 doc = llvm::json::parse(json)) { if (!handleMessage(std::move(*doc), handler)) return llvm::Error::success(); @@ -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); @@ -303,7 +303,8 @@ LogicalResult readLine(std::FILE *in, SmallVectorImpl &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; @@ -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))) { From b62319d7200931a25f1c75aa4bcf076e8b733d1a Mon Sep 17 00:00:00 2001 From: Kolya Panchenko Date: Mon, 3 Mar 2025 10:02:18 -0800 Subject: [PATCH 2/2] Addressed comments --- .../mlir/Tools/lsp-server-support/Transport.h | 54 ++++++------------- .../Tools/lsp-server-support/Transport.cpp | 14 ++--- 2 files changed, 23 insertions(+), 45 deletions(-) diff --git a/mlir/include/mlir/Tools/lsp-server-support/Transport.h b/mlir/include/mlir/Tools/lsp-server-support/Transport.h index 78f9d48e4153a..0010a475fedd2 100644 --- a/mlir/include/mlir/Tools/lsp-server-support/Transport.h +++ b/mlir/include/mlir/Tools/lsp-server-support/Transport.h @@ -50,11 +50,11 @@ class JSONTransportInput { : style(style) {} virtual ~JSONTransportInput() = default; - virtual bool getError() const = 0; + virtual bool hasError() const = 0; virtual bool isEndOfInput() const = 0; /// Read in a message from the input stream. - virtual LogicalResult readMessage(std::string &json) { + LogicalResult readMessage(std::string &json) { return style == JSONStreamStyle::Delimited ? readDelimitedMessage(json) : readStandardMessage(json); } @@ -66,15 +66,6 @@ class JSONTransportInput { 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: @@ -82,7 +73,7 @@ class JSONTransportInputOverFile : public JSONTransportInput { std::FILE *in, JSONStreamStyle style = JSONStreamStyle::Standard) : JSONTransportInput(style), in(in) {} - bool getError() const final { return ferror(in); } + bool hasError() const final { return ferror(in); } bool isEndOfInput() const final { return feof(in); } LogicalResult readDelimitedMessage(std::string &json) final; @@ -92,37 +83,19 @@ class JSONTransportInputOverFile : public JSONTransportInput { std::FILE *in; }; -/// Concrete implementation of the JSONTransportOutput that writes to a stream. -class JSONTransportOutputOverStream : public JSONTransportOutput { -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 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 in, - std::unique_ptr out) - : in(std::move(in)), out(std::move(out)) {} + JSONTransport(std::unique_ptr in, raw_ostream &out, + bool prettyOutput = false) + : in(std::move(in)), out(out), prettyOutput(prettyOutput) {} JSONTransport(std::FILE *in, raw_ostream &out, JSONStreamStyle style = JSONStreamStyle::Standard, bool prettyOutput = false) - : in(std::make_unique(in, style)), - out(std::make_unique(out, - prettyOutput)) {} + : in(std::make_unique(in, style)), out(out), + prettyOutput(prettyOutput) {} /// The following methods are used to send a message to the LSP client. void notify(StringRef method, llvm::json::Value params); @@ -132,15 +105,20 @@ 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); private: /// The input to read a message from. std::unique_ptr in; - - /// The output to send a messages to. - std::unique_ptr out; + SmallVector outputBuffer; + /// The output file stream. + raw_ostream &out; + /// If the output JSON should be formatted for easier readability. + bool prettyOutput; }; //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Tools/lsp-server-support/Transport.cpp b/mlir/lib/Tools/lsp-server-support/Transport.cpp index d6806bf5cfe13..d0863ba0ae087 100644 --- a/mlir/lib/Tools/lsp-server-support/Transport.cpp +++ b/mlir/lib/Tools/lsp-server-support/Transport.cpp @@ -175,7 +175,7 @@ llvm::Error decodeError(const llvm::json::Object &o) { } void JSONTransport::notify(StringRef method, llvm::json::Value params) { - out->sendMessage(llvm::json::Object{ + sendMessage(llvm::json::Object{ {"jsonrpc", "2.0"}, {"method", method}, {"params", std::move(params)}, @@ -183,7 +183,7 @@ void JSONTransport::notify(StringRef method, llvm::json::Value params) { } void JSONTransport::call(StringRef method, llvm::json::Value params, llvm::json::Value id) { - out->sendMessage(llvm::json::Object{ + sendMessage(llvm::json::Object{ {"jsonrpc", "2.0"}, {"id", std::move(id)}, {"method", method}, @@ -193,14 +193,14 @@ void JSONTransport::call(StringRef method, llvm::json::Value params, void JSONTransport::reply(llvm::json::Value id, llvm::Expected result) { if (result) { - return out->sendMessage(llvm::json::Object{ + return sendMessage(llvm::json::Object{ {"jsonrpc", "2.0"}, {"id", std::move(id)}, {"result", std::move(*result)}, }); } - out->sendMessage(llvm::json::Object{ + sendMessage(llvm::json::Object{ {"jsonrpc", "2.0"}, {"id", std::move(id)}, {"error", encodeError(result.takeError())}, @@ -210,7 +210,7 @@ void JSONTransport::reply(llvm::json::Value id, llvm::Error JSONTransport::run(MessageHandler &handler) { std::string json; while (!in->isEndOfInput()) { - if (in->getError()) { + if (in->hasError()) { return llvm::errorCodeToError( std::error_code(errno, std::system_category())); } @@ -227,7 +227,7 @@ llvm::Error JSONTransport::run(MessageHandler &handler) { return llvm::errorCodeToError(std::make_error_code(std::errc::io_error)); } -void JSONTransportOutputOverStream::sendMessage(const llvm::json::Value &msg) { +void JSONTransport::sendMessage(llvm::json::Value msg) { outputBuffer.clear(); llvm::raw_svector_ostream os(outputBuffer); os << llvm::formatv(prettyOutput ? "{0:2}\n" : "{0}", msg); @@ -310,7 +310,7 @@ JSONTransportInputOverFile::readStandardMessage(std::string &json) { unsigned long long contentLength = 0; llvm::SmallString<128> line; while (true) { - if (feof(in) || ferror(in) || failed(readLine(in, line))) + if (feof(in) || hasError() || failed(readLine(in, line))) return failure(); // Content-Length is a mandatory header, and the only one we handle.