[mlir-lsp-server] Add support for sending diagnostics to the client

[rriddle]

Summary

This allows for diagnostics emitted during parsing/verification to be surfaced to the user by the language client, as opposed to just being emitted to the logs like they are now.

[jpienaar]

Looks good in general thanks

mlir/include/mlir/Parser/AsmParserState.h

88

When would one do this vs just creating a new instance? (I'm assuming that is effectively what this is doing, keeping a reference constant but hollowing out - which a unique_ptr to parser state could also do)

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

197

Not: s/lsp/LSP/ in comment text

204

It's not just conceptually :-)

218

https://llvm.org/docs/CodingStandards.html#assert-liberally unreachable?

221

Could you add comment explaining why a warning is an error when converted?

mlir/lib/Tools/mlir-lsp-server/lsp/Protocol.h

416

Not: newline in between

471

Same nit

[rriddle]

update

mlir/include/mlir/Parser/AsmParserState.h

88

Switched to using a r-val operator=, which is what I originally wanted but couldn't get to work yesterday for some reason.

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

221

Oops, thanks for the catch.

mlir/lib/Tools/mlir-lsp-server/lsp/Protocol.h

416

This follows the style of the rest of the file, will fix the whole thing in a followup if that's okay.

471

Same as above.

mlir/include/mlir/Parser/AsmParserState.h

struct BlockDefinition {
     SMDefinition definition;
 
     /// Source definitions for any arguments of this block.
     SmallVector<SMDefinition> arguments;
   };
 
   AsmParserState();
   ~AsmParserState();
+  AsmParserState &operator=(AsmParserState &&other);
 
   //===--------------------------------------------------------------------===//
   // Access State
   //===--------------------------------------------------------------------===//
 
   using BlockDefIterator = llvm::pointee_iterator<
       ArrayRef<std::unique_ptr<BlockDefinition>>::iterator>;
   using OperationDefIterator = llvm::pointee_iterator<
       ArrayRef<std::unique_ptr<OperationDefinition>>::iterator>;
 
   /// Return a range of the BlockDefinitions held by the current parser state.

mlir/lib/Parser/AsmParserState.cpp

 };
 
 //===----------------------------------------------------------------------===//
 // AsmParserState
 //===----------------------------------------------------------------------===//
 
 AsmParserState::AsmParserState() : impl(std::make_unique<Impl>()) {}
 AsmParserState::~AsmParserState() {}
+AsmParserState &AsmParserState::operator=(AsmParserState &&other) {
+  impl = std::move(other.impl);
+  return *this;
+}
 
 //===----------------------------------------------------------------------===//
 // Access State
 
 auto AsmParserState::getBlockDefs() const -> iterator_range<BlockDefIterator> {
   return llvm::make_pointee_range(llvm::makeArrayRef(impl->blocks));
 }
 

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

struct LSPServer::Impl {
   // Hover
 
   void onHover(const TextDocumentPositionParams &params,
                Callback<Optional<Hover>> reply);
 
   MLIRServer &server;
   JSONTransport &transport;
 
+  /// An outgoing notification used to send diagnostics to the client when they
+  /// are ready to be processed.
+  OutgoingNotification<PublishDiagnosticsParams> publishDiagnostics;
+
   /// Used to indicate that the 'shutdown' request was received from the
   /// Language Server client.
   bool shutdownRequestReceived = false;
 };
 
 //===----------------------------------------------------------------------===//
 // Initialization
 
void LSPServer::Impl::onShutdown(const NoParams &,
   reply(nullptr);
 }
 
 //===----------------------------------------------------------------------===//
 // Document Change
 
 void LSPServer::Impl::onDocumentDidOpen(
     const DidOpenTextDocumentParams &params) {
-  server.addOrUpdateDocument(params.textDocument.uri, params.textDocument.text);
+  PublishDiagnosticsParams diagParams(params.textDocument.uri);
+  server.addOrUpdateDocument(params.textDocument.uri, params.textDocument.text,
+                             diagParams.diagnostics);
+
+  // Publish any recorded diagnostics.
+  publishDiagnostics(diagParams);
 }
 void LSPServer::Impl::onDocumentDidClose(
     const DidCloseTextDocumentParams &params) {
   server.removeDocument(params.textDocument.uri);
+
+  // Empty out the diagnostics shown for this document. This will clear out
+  // anything currently displayed by the client for this document (e.g. in the
+  // "Problems" pane of VSCode).
+  publishDiagnostics(PublishDiagnosticsParams(params.textDocument.uri));
 }
 void LSPServer::Impl::onDocumentDidChange(
     const DidChangeTextDocumentParams &params) {
   // TODO: We currently only support full document updates, we should refactor
   // to avoid this.
   if (params.contentChanges.size() != 1)
     return;
+  PublishDiagnosticsParams diagParams(params.textDocument.uri);
   server.addOrUpdateDocument(params.textDocument.uri,
-                             params.contentChanges.front().text);
+                             params.contentChanges.front().text,
+                             diagParams.diagnostics);
+
+  // Publish any recorded diagnostics.
+  publishDiagnostics(diagParams);
 }
 
 //===----------------------------------------------------------------------===//
 // Definitions and References
 
 void LSPServer::Impl::onGoToDefinition(const TextDocumentPositionParams &params,
                                        Callback<std::vector<Location>> reply) {
   std::vector<Location> locations;
LogicalResult LSPServer::run() {
   messageHandler.method("textDocument/definition", impl.get(),
                         &Impl::onGoToDefinition);
   messageHandler.method("textDocument/references", impl.get(),
                         &Impl::onReference);
 
   // Hover
   messageHandler.method("textDocument/hover", impl.get(), &Impl::onHover);
 
+  // Diagnostics
+  impl->publishDiagnostics =
+      messageHandler.outgoingNotification<PublishDiagnosticsParams>(
+          "textDocument/publishDiagnostics");
+
   // Run the main loop of the transport.
   LogicalResult result = success();
   if (llvm::Error error = impl->transport.run(messageHandler)) {
     Logger::error("Transport error: {0}", error);
     llvm::consumeError(std::move(error));
     result = failure();
   } else {
     result = success(impl->shutdownRequestReceived);
   }
   return result;
 }

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

 
 #include "mlir/Support/LLVM.h"
 #include <memory>
 
 namespace mlir {
 class DialectRegistry;
 
 namespace lsp {
+struct Diagnostic;
 struct Hover;
 struct Location;
 struct Position;
 class URIForFile;
 
 /// This class implements all of the MLIR related functionality necessary for a
 /// language server. This class allows for keeping the MLIR specific logic
 /// separate from the logic that involves LSP server/client communication.
 class MLIRServer {
 public:
   /// Construct a new server with the given dialect regitstry.
   MLIRServer(DialectRegistry &registry);
   ~MLIRServer();
 
-  /// Add or update the document at the given URI.
-  void addOrUpdateDocument(const URIForFile &uri, StringRef contents);
+  /// Add or update the document at the given URI. Any diagnostics emitted for
+  /// this document should be added to `diagnostics`
+  void addOrUpdateDocument(const URIForFile &uri, StringRef contents,
+                           std::vector<Diagnostic> &diagnostics);
 
   /// Remove the document with the given uri.
   void removeDocument(const URIForFile &uri);
 
   /// Return the locations of the object pointed at by the given position.
   void getLocationsOf(const URIForFile &uri, const Position &defPos,
                       std::vector<Location> &locations);
 

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

lsp::Logger::error("Failed to create URI for file `{0}`: {1}",
                        loc.getFilename(),
                        llvm::toString(sourceURI.takeError()));
     return llvm::None;
   }
 
   lsp::Position position;
   position.line = loc.getLine() - 1;
   position.character = loc.getColumn();
-  return lsp::Location{*sourceURI, lsp::Range{position, position}};
+  return lsp::Location{*sourceURI, lsp::Range(position)};
+}
+
+/// Returns a language server location from the given MLIR location, or None if
+/// one couldn't be created. `uri` is an optional additional filter that, when
+/// present, is used to filter sub locations that do not share the same uri.
+static Optional<lsp::Location>
+getLocationFromLoc(Location loc, const lsp::URIForFile *uri = nullptr) {
+  Optional<lsp::Location> location;
+  loc->walk([&](Location nestedLoc) {
+    FileLineColLoc fileLoc = nestedLoc.dyn_cast<FileLineColLoc>();
+    if (!fileLoc)
+      return WalkResult::advance();
+
+    Optional<lsp::Location> sourceLoc = getLocationFromLoc(fileLoc);
+    if (sourceLoc && (!uri || sourceLoc->uri == *uri)) {
+      location = *sourceLoc;
+      return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  });
+  return location;
 }
 
 /// Collect all of the locations from the given MLIR location that are not
 /// contained within the given URI.
 static void collectLocationsFromLoc(Location loc,
                                     std::vector<lsp::Location> &locations,
                                     const lsp::URIForFile &uri) {
   SetVector<Location> visitedLocs;
static void printDefBlockName(raw_ostream &os, Block *block,
   // Otherwise, we don't have a name so print the block number.
   os << "<Block #" << getBlockNumber(block) << ">";
 }
 static void printDefBlockName(raw_ostream &os,
                               const AsmParserState::BlockDefinition &def) {
   printDefBlockName(os, def.block, def.definition.loc);
 }
 
+/// Convert the given MLIR diagnostic to the LSP form.
+static lsp::Diagnostic getLspDiagnoticFromDiag(Diagnostic &diag,
+                                               const lsp::URIForFile &uri) {
+  lsp::Diagnostic lspDiag;
+  lspDiag.source = "mlir";
+
+  // Note: Right now all of the diagnostics are treated as parser issues, but
+  // some are parser and some are verifier.
+  lspDiag.category = "Parse Error";
+
+  // Try to grab a file location for this diagnostic.
+  // TODO: For simplicity, we just grab the first one. It may be likely that we
+  // will need a more interesting heuristic here.'
+  Optional<lsp::Location> lspLocation =
+      getLocationFromLoc(diag.getLocation(), &uri);
+  if (lspLocation)
+    lspDiag.range = lspLocation->range;
+
+  // Convert the severity for the diagnostic.
+  switch (diag.getSeverity()) {
+  case DiagnosticSeverity::Note:
+    llvm_unreachable("expected notes to be handled separately");
+    break;
+  case DiagnosticSeverity::Warning:
+    lspDiag.severity = lsp::DiagnosticSeverity::Warning;
+    break;
+  case DiagnosticSeverity::Error:
+    lspDiag.severity = lsp::DiagnosticSeverity::Error;
+    break;
+  case DiagnosticSeverity::Remark:
+    lspDiag.severity = lsp::DiagnosticSeverity::Information;
+    break;
+  }
+  lspDiag.message = diag.str();
+
+  // Attach any notes to the main diagnostic as related information.
+  std::vector<lsp::DiagnosticRelatedInformation> relatedDiags;
+  for (Diagnostic &note : diag.getNotes()) {
+    lsp::Location noteLoc;
+    if (Optional<lsp::Location> loc = getLocationFromLoc(note.getLocation()))
+      noteLoc = *loc;
+    else
+      noteLoc.uri = uri;
+    relatedDiags.emplace_back(noteLoc, note.str());
+  }
+  if (!relatedDiags.empty())
+    lspDiag.relatedInformation = std::move(relatedDiags);
+
+  return lspDiag;
+}
+
 //===----------------------------------------------------------------------===//
 // MLIRDocument
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class represents all of the information pertaining to a specific MLIR
 /// document.
 struct MLIRDocument {
   MLIRDocument(const lsp::URIForFile &uri, StringRef contents,
-               DialectRegistry &registry);
+               DialectRegistry &registry,
+               std::vector<lsp::Diagnostic> &diagnostics);
 
   //===--------------------------------------------------------------------===//
   // 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,
struct MLIRDocument {
   Block parsedIR;
 
   /// The source manager containing the contents of the input file.
   llvm::SourceMgr sourceMgr;
 };
 } // namespace
 
 MLIRDocument::MLIRDocument(const lsp::URIForFile &uri, StringRef contents,
-                           DialectRegistry &registry)
+                           DialectRegistry &registry,
+                           std::vector<lsp::Diagnostic> &diagnostics)
     : context(registry) {
   context.allowUnregisteredDialects();
   ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diag) {
-    // TODO: What should we do with these diagnostics?
-    //  * Cache and show to the user?
-    //  * Ignore?
-    lsp::Logger::error("Error when parsing MLIR document `{0}`: `{1}`",
-                       uri.file(), diag.str());
+    diagnostics.push_back(getLspDiagnoticFromDiag(diag, uri));
   });
 
   // Try to parsed the given IR string.
   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
   if (!memBuffer) {
     lsp::Logger::error("Failed to create memory buffer for file", uri.file());
     return;
   }
 
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), llvm::SMLoc());
-  if (failed(
-          parseSourceFile(sourceMgr, &parsedIR, &context, nullptr, &asmState)))
+  if (failed(parseSourceFile(sourceMgr, &parsedIR, &context, nullptr,
+                             &asmState))) {
+    // If parsing failed, clear out any of the current state.
+    parsedIR.clear();
+    asmState = AsmParserState();
     return;
   }
+}
 
 //===----------------------------------------------------------------------===//
 // MLIRDocument: Definitions and References
 //===----------------------------------------------------------------------===//
 
 void MLIRDocument::getLocationsOf(const lsp::URIForFile &uri,
                                   const lsp::Position &defPos,
                                   std::vector<lsp::Location> &locations) {

 //===----------------------------------------------------------------------===//
 // MLIRServer
 //===----------------------------------------------------------------------===//
 
 lsp::MLIRServer::MLIRServer(DialectRegistry &registry)
     : impl(std::make_unique<Impl>(registry)) {}
 lsp::MLIRServer::~MLIRServer() {}
 
-void lsp::MLIRServer::addOrUpdateDocument(const URIForFile &uri,
-                                          StringRef contents) {
-  impl->documents[uri.file()] =
-      std::make_unique<MLIRDocument>(uri, contents, impl->registry);
+void lsp::MLIRServer::addOrUpdateDocument(
+    const URIForFile &uri, StringRef contents,
+    std::vector<Diagnostic> &diagnostics) {
+  impl->documents[uri.file()] = std::make_unique<MLIRDocument>(
+      uri, contents, impl->registry, diagnostics);
 }
 
 void lsp::MLIRServer::removeDocument(const URIForFile &uri) {
   impl->documents.erase(uri.file());
 }
 
 void lsp::MLIRServer::getLocationsOf(const URIForFile &uri,
                                      const Position &defPos,

mlir/lib/Tools/mlir-lsp-server/lsp/Protocol.h

 llvm::json::Value toJSON(const Position &value);
 raw_ostream &operator<<(raw_ostream &os, const Position &value);
 
 //===----------------------------------------------------------------------===//
 // Range
 //===----------------------------------------------------------------------===//
 
 struct Range {
+  Range() = default;
+  Range(Position start, Position end) : start(start), end(end) {}
+  Range(Position loc) : Range(loc, loc) {}
+
   /// The range's start position.
   Position start;
 
   /// The range's end position.
   Position end;
 
   friend bool operator==(const Range &lhs, const Range &rhs) {
     return std::tie(lhs.start, lhs.end) == std::tie(rhs.start, rhs.end);
struct Hover {
   MarkupContent contents;
 
   /// An optional range is a range inside a text document that is used to
   /// visualize a hover, e.g. by changing the background color.
   Optional<Range> range;
 };
 llvm::json::Value toJSON(const Hover &hover);
 
+//===----------------------------------------------------------------------===//
+// DiagnosticRelatedInformation
+//===----------------------------------------------------------------------===//
+
+/// Represents a related message and source code location for a diagnostic.
+/// This should be used to point to code locations that cause or related to a
+/// diagnostics, e.g. when duplicating a symbol in a scope.
+struct DiagnosticRelatedInformation {
+  DiagnosticRelatedInformation(Location location, std::string message)
+      : location(location), message(std::move(message)) {}
+
+  /// The location of this related diagnostic information.
+  Location location;
+  /// The message of this related diagnostic information.
+  std::string message;
+};
+llvm::json::Value toJSON(const DiagnosticRelatedInformation &info);
+
+//===----------------------------------------------------------------------===//
+// Diagnostic
+//===----------------------------------------------------------------------===//
+
+enum class DiagnosticSeverity {
+  /// It is up to the client to interpret diagnostics as error, warning, info or
+  /// hint.
+  Undetermined = 0,
+  Error = 1,
+  Warning = 2,
+  Information = 3,
+  Hint = 4
+};
+
+struct Diagnostic {
+  /// The source range where the message applies.
+  Range range;
+
+  /// The diagnostic's severity. Can be omitted. If omitted it is up to the
+  /// client to interpret diagnostics as error, warning, info or hint.
+  DiagnosticSeverity severity = DiagnosticSeverity::Undetermined;
+
+  /// A human-readable string describing the source of this diagnostic, e.g.
+  /// 'typescript' or 'super lint'.
+  std::string source;
+
+  /// The diagnostic's message.
+  std::string message;
+
+  /// An array of related diagnostic information, e.g. when symbol-names within
+  /// a scope collide all definitions can be marked via this property.
+  Optional<std::vector<DiagnosticRelatedInformation>> relatedInformation;
+
+  /// The diagnostic's category. Can be omitted.
+  /// An LSP extension that's used to send the name of the category over to the
+  /// client. The category typically describes the compilation stage during
+  /// which the issue was produced, e.g. "Semantic Issue" or "Parse Issue".
+  Optional<std::string> category;
+};
+llvm::json::Value toJSON(const Diagnostic &diag);
+
+//===----------------------------------------------------------------------===//
+// PublishDiagnosticsParams
+//===----------------------------------------------------------------------===//
+
+struct PublishDiagnosticsParams {
+  PublishDiagnosticsParams(URIForFile uri) : uri(uri) {}
+
+  /// The URI for which diagnostic information is reported.
+  URIForFile uri;
+  /// The list of reported diagnostics.
+  std::vector<Diagnostic> diagnostics;
+};
+llvm::json::Value toJSON(const PublishDiagnosticsParams &params);
+
 } // namespace lsp
 } // namespace mlir
 
 namespace llvm {
 template <> struct format_provider<mlir::lsp::Position> {
   static void format(const mlir::lsp::Position &pos, raw_ostream &os,
                      StringRef style) {
     assert(style.empty() && "style modifiers for this type are not supported");
     os << pos;
   }
 };
 } // namespace llvm
 
 #endif

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

 //===----------------------------------------------------------------------===//
 
 llvm::json::Value mlir::lsp::toJSON(const Hover &hover) {
   llvm::json::Object result{{"contents", toJSON(hover.contents)}};
   if (hover.range.hasValue())
     result["range"] = toJSON(*hover.range);
   return std::move(result);
 }
+
+//===----------------------------------------------------------------------===//
+// DiagnosticRelatedInformation
+//===----------------------------------------------------------------------===//
+
+llvm::json::Value mlir::lsp::toJSON(const DiagnosticRelatedInformation &info) {
+  return llvm::json::Object{
+      {"location", info.location},
+      {"message", info.message},
+  };
+}
+
+//===----------------------------------------------------------------------===//
+// Diagnostic
+//===----------------------------------------------------------------------===//
+
+llvm::json::Value mlir::lsp::toJSON(const Diagnostic &diag) {
+  llvm::json::Object result{
+      {"range", diag.range},
+      {"severity", (int)diag.severity},
+      {"message", diag.message},
+  };
+  if (diag.category)
+    result["category"] = *diag.category;
+  if (!diag.source.empty())
+    result["source"] = diag.source;
+  if (diag.relatedInformation)
+    result["relatedInformation"] = *diag.relatedInformation;
+  return std::move(result);
+}
+
+//===----------------------------------------------------------------------===//
+// PublishDiagnosticsParams
+//===----------------------------------------------------------------------===//
+
+llvm::json::Value mlir::lsp::toJSON(const PublishDiagnosticsParams &params) {
+  return llvm::json::Object{
+      {"uri", params.uri},
+      {"diagnostics", params.diagnostics},
+  };
+}

mlir/lib/Tools/mlir-lsp-server/lsp/Transport.h

 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/FormatAdapters.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_ostream.h"
 #include <atomic>
 
 namespace mlir {
 namespace lsp {
-class JSONTransport;
+class MessageHandler;
 
 //===----------------------------------------------------------------------===//
-// Reply
+// JSONTransport
 //===----------------------------------------------------------------------===//
 
-/// Function object to reply to an LSP call.
-/// Each instance must be called exactly once, otherwise:
-///  - if there was no reply, an error reply is sent
-///  - if there were multiple replies, only the first is sent
-class Reply {
+/// The encoding style of the JSON-RPC messages (both input and output).
+enum JSONStreamStyle {
+  /// Encoding per the LSP specification, with mandatory Content-Length header.
+  Standard,
+  /// Messages are delimited by a '// -----' line. Comment lines start with //.
+  Delimited
+};
+
+/// A transport class that performs the JSON-RPC communication with the LSP
+/// client.
+class JSONTransport {
 public:
-  Reply(const llvm::json::Value &id, StringRef method,
-        JSONTransport &transport);
-  Reply(Reply &&other);
-  Reply &operator=(Reply &&) = delete;
-  Reply(const Reply &) = delete;
-  Reply &operator=(const Reply &) = delete;
+  JSONTransport(std::FILE *in, raw_ostream &out,
+                JSONStreamStyle style = JSONStreamStyle::Standard,
+                bool prettyOutput = false)
+      : in(in), out(out), style(style), prettyOutput(prettyOutput) {}
 
-  void operator()(llvm::Expected<llvm::json::Value> reply);
+  /// The following methods are used to send a message to the LSP client.
+  void notify(StringRef method, llvm::json::Value params);
+  void call(StringRef method, llvm::json::Value params, llvm::json::Value id);
+  void reply(llvm::json::Value id, llvm::Expected<llvm::json::Value> result);
+
+  /// Start executing the JSON-RPC transport.
+  llvm::Error run(MessageHandler &handler);
 
 private:
-  StringRef method;
-  std::atomic<bool> replied = {false};
-  llvm::json::Value id;
-  JSONTransport *transport;
+  /// 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);
+
+  /// 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;
 };
 
 //===----------------------------------------------------------------------===//
 // MessageHandler
 //===----------------------------------------------------------------------===//
 
 /// A Callback<T> is a void function that accepts Expected<T>. This is
 /// accepted by functions that logically return T.
 template <typename T>
 using Callback = llvm::unique_function<void(llvm::Expected<T>)>;
 
+/// An OutgoingNotification<T> is a function used for outgoing notifications
+/// send to the client.
+template <typename T>
+using OutgoingNotification = llvm::unique_function<void(const T &)>;
+
 /// A handler used to process the incoming transport messages.
 class MessageHandler {
 public:
   MessageHandler(JSONTransport &transport) : transport(transport) {}
 
   bool onNotify(StringRef method, llvm::json::Value value);
   bool onCall(StringRef method, llvm::json::Value params, llvm::json::Value id);
   bool onReply(llvm::json::Value id, llvm::Expected<llvm::json::Value> result);
notificationHandlers[method] = [method, handler,
                                     thisPtr](llvm::json::Value rawParams) {
       llvm::Expected<Param> param = parse<Param>(rawParams, method, "request");
       if (!param)
         return llvm::consumeError(param.takeError());
       (thisPtr->*handler)(*param);
     };
   }
 
+  /// Create an OutgoingNotification object used for the given method.
+  template <typename T>
+  OutgoingNotification<T> outgoingNotification(llvm::StringLiteral method) {
+    return [&, method](const T &params) {
+      transport.notify(method, llvm::json::Value(params));
+    };
+  }
+
 private:
   template <typename HandlerT>
   using HandlerMap = llvm::StringMap<llvm::unique_function<HandlerT>>;
 
   HandlerMap<void(llvm::json::Value)> notificationHandlers;
   HandlerMap<void(llvm::json::Value, Callback<llvm::json::Value>)>
       methodHandlers;
 
   JSONTransport &transport;
 };
 
-//===----------------------------------------------------------------------===//
-// JSONTransport
-//===----------------------------------------------------------------------===//
-
-/// The encoding style of the JSON-RPC messages (both input and output).
-enum JSONStreamStyle {
-  /// Encoding per the LSP specification, with mandatory Content-Length header.
-  Standard,
-  /// Messages are delimited by a '// -----' line. Comment lines start with //.
-  Delimited
-};
-
-/// A transport class that performs the JSON-RPC communication with the LSP
-/// client.
-class JSONTransport {
-public:
-  JSONTransport(std::FILE *in, raw_ostream &out,
-                JSONStreamStyle style = JSONStreamStyle::Standard,
-                bool prettyOutput = false)
-      : in(in), out(out), style(style), prettyOutput(prettyOutput) {}
-
-  /// The following methods are used to send a message to the LSP client.
-  void notify(StringRef method, llvm::json::Value params);
-  void call(StringRef method, llvm::json::Value params, llvm::json::Value id);
-  void reply(llvm::json::Value id, llvm::Expected<llvm::json::Value> result);
-
-  /// 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);
-
-  /// 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;
-};
-
 } // namespace lsp
 } // namespace mlir
 
 #endif

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

 
 using namespace mlir;
 using namespace mlir::lsp;
 
 //===----------------------------------------------------------------------===//
 // Reply
 //===----------------------------------------------------------------------===//
 
+namespace {
+/// Function object to reply to an LSP call.
+/// Each instance must be called exactly once, otherwise:
+///  - if there was no reply, an error reply is sent
+///  - if there were multiple replies, only the first is sent
+class Reply {
+public:
+  Reply(const llvm::json::Value &id, StringRef method,
+        JSONTransport &transport);
+  Reply(Reply &&other);
+  Reply &operator=(Reply &&) = delete;
+  Reply(const Reply &) = delete;
+  Reply &operator=(const Reply &) = delete;
+
+  void operator()(llvm::Expected<llvm::json::Value> reply);
+
+private:
+  StringRef method;
+  std::atomic<bool> replied = {false};
+  llvm::json::Value id;
+  JSONTransport *transport;
+};
+} // namespace
+
 Reply::Reply(const llvm::json::Value &id, llvm::StringRef method,
              JSONTransport &transport)
     : id(id), transport(&transport) {}
 
 Reply::Reply(Reply &&other)
     : replied(other.replied.load()), id(std::move(other.id)),
       transport(other.transport) {
   other.transport = nullptr;

mlir/test/mlir-lsp-server/diagnostics.test

+// RUN: mlir-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
+{"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 ()"
+}}}
+// CHECK: "method": "textDocument/publishDiagnostics",
+// CHECK-NEXT: "params": {
+// CHECK-NEXT:     "diagnostics": [
+// CHECK-NEXT:       {
+// CHECK-NEXT:         "category": "Parse Error",
+// CHECK-NEXT:         "message": "custom op 'func' expected valid '@'-identifier for symbol name",
+// CHECK-NEXT:         "range": {
+// CHECK-NEXT:           "end": {
+// CHECK-NEXT:             "character": 6,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           },
+// CHECK-NEXT:           "start": {
+// CHECK-NEXT:             "character": 6,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           }
+// CHECK-NEXT:         },
+// CHECK-NEXT:         "severity": 1,
+// CHECK-NEXT:         "source": "mlir"
+// CHECK-NEXT:       }
+// CHECK-NEXT:     ],
+// CHECK-NEXT:     "uri": "test:///foo.mlir"
+// CHECK-NEXT:   }
+// -----
+{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}