1 files changed, 22 insertions, 45 deletions

diff --git a/clang-tools-extra/clangd/JSONRPCDispatcher.h b/clang-tools-extra/clangd/JSONRPCDispatcher.h
index b223ceb5f45..43c466f7014 100644
--- a/clang-tools-extra/clangd/JSONRPCDispatcher.h
+++ b/clang-tools-extra/clangd/JSONRPCDispatcher.h
@@ -14,6 +14,7 @@
 #include "Logger.h"
 #include "Protocol.h"
 #include "Trace.h"
+#include "Transport.h"
 #include "clang/Basic/LLVM.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringMap.h"
@@ -24,37 +25,19 @@
 namespace clang {
 namespace clangd {
 
-/// Encapsulates output and logs streams and provides thread-safe access to
-/// them.
+// Logs to an output stream, such as stderr.
+// FIXME: Rename to StreamLogger or such, and move to Logger.h.
 class JSONOutput : public Logger {
-  // FIXME(ibiryukov): figure out if we can shrink the public interface of
-  // JSONOutput now that we pass Context everywhere.
 public:
-  JSONOutput(llvm::raw_ostream &Outs, llvm::raw_ostream &Logs,
-             Logger::Level MinLevel, llvm::raw_ostream *InputMirror = nullptr,
-             bool Pretty = false)
-      : Pretty(Pretty), MinLevel(MinLevel), Outs(Outs), Logs(Logs),
-        InputMirror(InputMirror) {}
-
-  /// Emit a JSONRPC message.
-  void writeMessage(const llvm::json::Value &Result);
+  JSONOutput(llvm::raw_ostream &Logs, Logger::Level MinLevel)
+      : MinLevel(MinLevel), Logs(Logs) {}
 
   /// Write a line to the logging stream.
   void log(Level, const llvm::formatv_object_base &Message) override;
 
-  /// Mirror \p Message into InputMirror stream. Does nothing if InputMirror is
-  /// null.
-  /// Unlike other methods of JSONOutput, mirrorInput is not thread-safe.
-  void mirrorInput(const Twine &Message);
-
-  // Whether output should be pretty-printed.
-  const bool Pretty;
-
 private:
   Logger::Level MinLevel;
-  llvm::raw_ostream &Outs;
   llvm::raw_ostream &Logs;
-  llvm::raw_ostream *InputMirror;
 
   std::mutex StreamMutex;
 };
@@ -81,14 +64,15 @@ void call(llvm::StringRef Method, llvm::json::Value &&Params);
 ///
 /// The `$/cancelRequest` notification is handled by the dispatcher itself.
 /// It marks the matching request as cancelled, if it's still running.
-class JSONRPCDispatcher {
+class JSONRPCDispatcher : private Transport::MessageHandler {
 public:
   /// A handler responds to requests for a particular method name.
+  /// It returns false if the server should now shut down.
   ///
   /// JSONRPCDispatcher will mark the handler's context as cancelled if a
   /// matching cancellation request is received. Handlers are encouraged to
   /// check for cancellation and fail quickly in this case.
-  using Handler = std::function<void(const llvm::json::Value &)>;
+  using Handler = std::function<bool(const llvm::json::Value &)>;
 
   /// Create a new JSONRPCDispatcher. UnknownHandler is called when an unknown
   /// method is received.
@@ -97,10 +81,22 @@ public:
   /// Registers a Handler for the specified Method.
   void registerHandler(StringRef Method, Handler H);
 
-  /// Parses a JSONRPC message and calls the Handler for it.
-  bool call(const llvm::json::Value &Message, JSONOutput &Out);
+  /// Parses input queries from LSP client (coming from \p In) and runs call
+  /// method for each query.
+  ///
+  /// Input stream(\p In) must be opened in binary mode to avoid
+  /// preliminary replacements of \r\n with \n. We use C-style FILE* for reading
+  /// as std::istream has unclear interaction with signals, which are sent by
+  /// debuggers on some OSs.
+  llvm::Error runLanguageServerLoop(Transport &);
 
 private:
+  bool onReply(llvm::json::Value ID,
+               llvm::Expected<llvm::json::Value> Result) override;
+  bool onNotify(llvm::StringRef Method, llvm::json::Value Message) override;
+  bool onCall(llvm::StringRef Method, llvm::json::Value Message,
+              llvm::json::Value ID) override;
+
   // Tracking cancellations needs a mutex: handlers may finish on a different
   // thread, and that's when we clean up entries in the map.
   mutable std::mutex RequestCancelersMutex;
@@ -113,25 +109,6 @@ private:
   Handler UnknownHandler;
 };
 
-/// Controls the way JSON-RPC messages are encoded (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
-};
-
-/// Parses input queries from LSP client (coming from \p In) and runs call
-/// method of \p Dispatcher for each query.
-/// After handling each query checks if \p IsDone is set true and exits the loop
-/// if it is.
-/// Input stream(\p In) must be opened in binary mode to avoid preliminary
-/// replacements of \r\n with \n.
-/// We use C-style FILE* for reading as std::istream has unclear interaction
-/// with signals, which are sent by debuggers on some OSs.
-void runLanguageServerLoop(std::FILE *In, JSONOutput &Out,
-                           JSONStreamStyle InputStyle,
-                           JSONRPCDispatcher &Dispatcher, bool &IsDone);
 } // namespace clangd
 } // namespace clang