1 //===--- Transport.h - Sending and Receiving LSP messages -------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // The language server protocol is usually implemented by writing messages as
10 // JSON-RPC over the stdin/stdout of a subprocess. This file contains a JSON
11 // transport interface that handles this communication.
12 //
13 //===----------------------------------------------------------------------===//
14 
15 #ifndef LIB_MLIR_TOOLS_LSPSERVERSUPPORT_TRANSPORT_H_
16 #define LIB_MLIR_TOOLS_LSPSERVERSUPPORT_TRANSPORT_H_
17 
18 #include "Logging.h"
19 #include "Protocol.h"
20 #include "mlir/Support/LLVM.h"
21 #include "mlir/Support/LogicalResult.h"
22 #include "llvm/ADT/FunctionExtras.h"
23 #include "llvm/ADT/StringMap.h"
24 #include "llvm/ADT/StringRef.h"
25 #include "llvm/Support/FormatAdapters.h"
26 #include "llvm/Support/JSON.h"
27 #include "llvm/Support/raw_ostream.h"
28 #include <atomic>
29 
30 namespace mlir {
31 namespace lsp {
32 class MessageHandler;
33 
34 //===----------------------------------------------------------------------===//
35 // JSONTransport
36 //===----------------------------------------------------------------------===//
37 
38 /// The encoding style of the JSON-RPC messages (both input and output).
39 enum JSONStreamStyle {
40   /// Encoding per the LSP specification, with mandatory Content-Length header.
41   Standard,
42   /// Messages are delimited by a '// -----' line. Comment lines start with //.
43   Delimited
44 };
45 
46 /// A transport class that performs the JSON-RPC communication with the LSP
47 /// client.
48 class JSONTransport {
49 public:
50   JSONTransport(std::FILE *in, raw_ostream &out,
51                 JSONStreamStyle style = JSONStreamStyle::Standard,
52                 bool prettyOutput = false)
in(in)53       : in(in), out(out), style(style), prettyOutput(prettyOutput) {}
54 
55   /// The following methods are used to send a message to the LSP client.
56   void notify(StringRef method, llvm::json::Value params);
57   void call(StringRef method, llvm::json::Value params, llvm::json::Value id);
58   void reply(llvm::json::Value id, llvm::Expected<llvm::json::Value> result);
59 
60   /// Start executing the JSON-RPC transport.
61   llvm::Error run(MessageHandler &handler);
62 
63 private:
64   /// Dispatches the given incoming json message to the message handler.
65   bool handleMessage(llvm::json::Value msg, MessageHandler &handler);
66   /// Writes the given message to the output stream.
67   void sendMessage(llvm::json::Value msg);
68 
69   /// Read in a message from the input stream.
readMessage(std::string & json)70   LogicalResult readMessage(std::string &json) {
71     return style == JSONStreamStyle::Delimited ? readDelimitedMessage(json)
72                                                : readStandardMessage(json);
73   }
74   LogicalResult readDelimitedMessage(std::string &json);
75   LogicalResult readStandardMessage(std::string &json);
76 
77   /// An output buffer used when building output messages.
78   SmallVector<char, 0> outputBuffer;
79   /// The input file stream.
80   std::FILE *in;
81   /// The output file stream.
82   raw_ostream &out;
83   /// The JSON stream style to use.
84   JSONStreamStyle style;
85   /// If the output JSON should be formatted for easier readability.
86   bool prettyOutput;
87 };
88 
89 //===----------------------------------------------------------------------===//
90 // MessageHandler
91 //===----------------------------------------------------------------------===//
92 
93 /// A Callback<T> is a void function that accepts Expected<T>. This is
94 /// accepted by functions that logically return T.
95 template <typename T>
96 using Callback = llvm::unique_function<void(llvm::Expected<T>)>;
97 
98 /// An OutgoingNotification<T> is a function used for outgoing notifications
99 /// send to the client.
100 template <typename T>
101 using OutgoingNotification = llvm::unique_function<void(const T &)>;
102 
103 /// A handler used to process the incoming transport messages.
104 class MessageHandler {
105 public:
MessageHandler(JSONTransport & transport)106   MessageHandler(JSONTransport &transport) : transport(transport) {}
107 
108   bool onNotify(StringRef method, llvm::json::Value value);
109   bool onCall(StringRef method, llvm::json::Value params, llvm::json::Value id);
110   bool onReply(llvm::json::Value id, llvm::Expected<llvm::json::Value> result);
111 
112   template <typename T>
parse(const llvm::json::Value & raw,StringRef payloadName,StringRef payloadKind)113   static llvm::Expected<T> parse(const llvm::json::Value &raw,
114                                  StringRef payloadName, StringRef payloadKind) {
115     T result;
116     llvm::json::Path::Root root;
117     if (fromJSON(raw, result, root))
118       return std::move(result);
119 
120     // Dump the relevant parts of the broken message.
121     std::string context;
122     llvm::raw_string_ostream os(context);
123     root.printErrorContext(raw, os);
124 
125     // Report the error (e.g. to the client).
126     return llvm::make_error<LSPError>(
127         llvm::formatv("failed to decode {0} {1}: {2}", payloadName, payloadKind,
128                       fmt_consume(root.getError())),
129         ErrorCode::InvalidParams);
130   }
131 
132   template <typename Param, typename Result, typename ThisT>
method(llvm::StringLiteral method,ThisT * thisPtr,void (ThisT::* handler)(const Param &,Callback<Result>))133   void method(llvm::StringLiteral method, ThisT *thisPtr,
134               void (ThisT::*handler)(const Param &, Callback<Result>)) {
135     methodHandlers[method] = [method, handler,
136                               thisPtr](llvm::json::Value rawParams,
137                                        Callback<llvm::json::Value> reply) {
138       llvm::Expected<Param> param = parse<Param>(rawParams, method, "request");
139       if (!param)
140         return reply(param.takeError());
141       (thisPtr->*handler)(*param, std::move(reply));
142     };
143   }
144 
145   template <typename Param, typename ThisT>
notification(llvm::StringLiteral method,ThisT * thisPtr,void (ThisT::* handler)(const Param &))146   void notification(llvm::StringLiteral method, ThisT *thisPtr,
147                     void (ThisT::*handler)(const Param &)) {
148     notificationHandlers[method] = [method, handler,
149                                     thisPtr](llvm::json::Value rawParams) {
150       llvm::Expected<Param> param = parse<Param>(rawParams, method, "request");
151       if (!param)
152         return llvm::consumeError(param.takeError());
153       (thisPtr->*handler)(*param);
154     };
155   }
156 
157   /// Create an OutgoingNotification object used for the given method.
158   template <typename T>
outgoingNotification(llvm::StringLiteral method)159   OutgoingNotification<T> outgoingNotification(llvm::StringLiteral method) {
160     return [&, method](const T &params) {
161       Logger::info("--> {0}", method);
162       transport.notify(method, llvm::json::Value(params));
163     };
164   }
165 
166 private:
167   template <typename HandlerT>
168   using HandlerMap = llvm::StringMap<llvm::unique_function<HandlerT>>;
169 
170   HandlerMap<void(llvm::json::Value)> notificationHandlers;
171   HandlerMap<void(llvm::json::Value, Callback<llvm::json::Value>)>
172       methodHandlers;
173 
174   JSONTransport &transport;
175 };
176 
177 } // namespace lsp
178 } // namespace mlir
179 
180 #endif
181