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

[npanchen]

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.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Nikolay Panchenko (npanchen)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/129320.diff

2 Files Affected:

• (modified) mlir/include/mlir/Tools/lsp-server-support/Transport.h (+78-22)
• (modified) mlir/lib/Tools/lsp-server-support/Transport.cpp (+11-10)

diff --git a/mlir/include/mlir/Tools/lsp-server-support/Transport.h b/mlir/include/mlir/Tools/lsp-server-support/Transport.h
index 6843bc76ab9dc..4b01b545fc6b2 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<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);
@@ -60,30 +132,14 @@ 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;
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Tools/lsp-server-support/Transport.cpp b/mlir/lib/Tools/lsp-server-support/Transport.cpp
index ad8308f69aead..fbb9a7acc16d3 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<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())},
@@ -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();
@@ -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<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;
@@ -349,7 +350,7 @@ 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))) {

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[River707]

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

[npanchen]

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.

[River707]

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).

[npanchen]

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.

[River707]

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.

[npanchen]

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

[River707]

Abstracting over the input somewhat makes sense to me, given that we don't really have a decent input stream abstraction AFAIK. For the output stream though, I'd strongly prefer to not abstract that and just use raw_ostream (unless there is a very good reason not to).

[River707]

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

[npanchen]

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

[River707]

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

[npanchen]

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

[npanchen]

@River707 ping

[npanchen]

@River707 ping