1 //===- LSPServer.cpp - MLIR Language Server -------------------------------===//
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 #include "LSPServer.h"
10 #include "MLIRServer.h"
11 #include "lsp/Logging.h"
12 #include "lsp/Protocol.h"
13 #include "lsp/Transport.h"
14 #include "llvm/ADT/FunctionExtras.h"
15 #include "llvm/ADT/StringMap.h"
16 
17 #define DEBUG_TYPE "mlir-lsp-server"
18 
19 using namespace mlir;
20 using namespace mlir::lsp;
21 
22 //===----------------------------------------------------------------------===//
23 // LSPServer::Impl
24 //===----------------------------------------------------------------------===//
25 
26 struct LSPServer::Impl {
27   Impl(MLIRServer &server, JSONTransport &transport)
28       : server(server), transport(transport) {}
29 
30   //===--------------------------------------------------------------------===//
31   // Initialization
32 
33   void onInitialize(const InitializeParams &params,
34                     Callback<llvm::json::Value> reply);
35   void onInitialized(const InitializedParams &params);
36   void onShutdown(const NoParams &params, Callback<std::nullptr_t> reply);
37 
38   //===--------------------------------------------------------------------===//
39   // Document Change
40 
41   void onDocumentDidOpen(const DidOpenTextDocumentParams &params);
42   void onDocumentDidClose(const DidCloseTextDocumentParams &params);
43   void onDocumentDidChange(const DidChangeTextDocumentParams &params);
44 
45   //===--------------------------------------------------------------------===//
46   // Definitions and References
47 
48   void onGoToDefinition(const TextDocumentPositionParams &params,
49                         Callback<std::vector<Location>> reply);
50   void onReference(const ReferenceParams &params,
51                    Callback<std::vector<Location>> reply);
52 
53   //===--------------------------------------------------------------------===//
54   // Hover
55 
56   void onHover(const TextDocumentPositionParams &params,
57                Callback<Optional<Hover>> reply);
58 
59   MLIRServer &server;
60   JSONTransport &transport;
61 
62   /// An outgoing notification used to send diagnostics to the client when they
63   /// are ready to be processed.
64   OutgoingNotification<PublishDiagnosticsParams> publishDiagnostics;
65 
66   /// Used to indicate that the 'shutdown' request was received from the
67   /// Language Server client.
68   bool shutdownRequestReceived = false;
69 };
70 
71 //===----------------------------------------------------------------------===//
72 // Initialization
73 
74 void LSPServer::Impl::onInitialize(const InitializeParams &params,
75                                    Callback<llvm::json::Value> reply) {
76   llvm::json::Object serverCaps{
77       {"textDocumentSync",
78        llvm::json::Object{
79            {"openClose", true},
80            {"change", (int)TextDocumentSyncKind::Full},
81            {"save", true},
82        }},
83       {"definitionProvider", true},
84       {"referencesProvider", true},
85       {"hoverProvider", true},
86   };
87 
88   llvm::json::Object result{
89       {{"serverInfo",
90         llvm::json::Object{{"name", "mlir-lsp-server"}, {"version", "0.0.0"}}},
91        {"capabilities", std::move(serverCaps)}}};
92   reply(std::move(result));
93 }
94 void LSPServer::Impl::onInitialized(const InitializedParams &) {}
95 void LSPServer::Impl::onShutdown(const NoParams &,
96                                  Callback<std::nullptr_t> reply) {
97   shutdownRequestReceived = true;
98   reply(nullptr);
99 }
100 
101 //===----------------------------------------------------------------------===//
102 // Document Change
103 
104 void LSPServer::Impl::onDocumentDidOpen(
105     const DidOpenTextDocumentParams &params) {
106   PublishDiagnosticsParams diagParams(params.textDocument.uri);
107   server.addOrUpdateDocument(params.textDocument.uri, params.textDocument.text,
108                              diagParams.diagnostics);
109 
110   // Publish any recorded diagnostics.
111   publishDiagnostics(diagParams);
112 }
113 void LSPServer::Impl::onDocumentDidClose(
114     const DidCloseTextDocumentParams &params) {
115   server.removeDocument(params.textDocument.uri);
116 
117   // Empty out the diagnostics shown for this document. This will clear out
118   // anything currently displayed by the client for this document (e.g. in the
119   // "Problems" pane of VSCode).
120   publishDiagnostics(PublishDiagnosticsParams(params.textDocument.uri));
121 }
122 void LSPServer::Impl::onDocumentDidChange(
123     const DidChangeTextDocumentParams &params) {
124   // TODO: We currently only support full document updates, we should refactor
125   // to avoid this.
126   if (params.contentChanges.size() != 1)
127     return;
128   PublishDiagnosticsParams diagParams(params.textDocument.uri);
129   server.addOrUpdateDocument(params.textDocument.uri,
130                              params.contentChanges.front().text,
131                              diagParams.diagnostics);
132 
133   // Publish any recorded diagnostics.
134   publishDiagnostics(diagParams);
135 }
136 
137 //===----------------------------------------------------------------------===//
138 // Definitions and References
139 
140 void LSPServer::Impl::onGoToDefinition(const TextDocumentPositionParams &params,
141                                        Callback<std::vector<Location>> reply) {
142   std::vector<Location> locations;
143   server.getLocationsOf(params.textDocument.uri, params.position, locations);
144   reply(std::move(locations));
145 }
146 
147 void LSPServer::Impl::onReference(const ReferenceParams &params,
148                                   Callback<std::vector<Location>> reply) {
149   std::vector<Location> locations;
150   server.findReferencesOf(params.textDocument.uri, params.position, locations);
151   reply(std::move(locations));
152 }
153 
154 //===----------------------------------------------------------------------===//
155 // Hover
156 
157 void LSPServer::Impl::onHover(const TextDocumentPositionParams &params,
158                               Callback<Optional<Hover>> reply) {
159   reply(server.findHover(params.textDocument.uri, params.position));
160 }
161 
162 //===----------------------------------------------------------------------===//
163 // LSPServer
164 //===----------------------------------------------------------------------===//
165 
166 LSPServer::LSPServer(MLIRServer &server, JSONTransport &transport)
167     : impl(std::make_unique<Impl>(server, transport)) {}
168 LSPServer::~LSPServer() {}
169 
170 LogicalResult LSPServer::run() {
171   MessageHandler messageHandler(impl->transport);
172 
173   // Initialization
174   messageHandler.method("initialize", impl.get(), &Impl::onInitialize);
175   messageHandler.notification("initialized", impl.get(), &Impl::onInitialized);
176   messageHandler.method("shutdown", impl.get(), &Impl::onShutdown);
177 
178   // Document Changes
179   messageHandler.notification("textDocument/didOpen", impl.get(),
180                               &Impl::onDocumentDidOpen);
181   messageHandler.notification("textDocument/didClose", impl.get(),
182                               &Impl::onDocumentDidClose);
183   messageHandler.notification("textDocument/didChange", impl.get(),
184                               &Impl::onDocumentDidChange);
185 
186   // Definitions and References
187   messageHandler.method("textDocument/definition", impl.get(),
188                         &Impl::onGoToDefinition);
189   messageHandler.method("textDocument/references", impl.get(),
190                         &Impl::onReference);
191 
192   // Hover
193   messageHandler.method("textDocument/hover", impl.get(), &Impl::onHover);
194 
195   // Diagnostics
196   impl->publishDiagnostics =
197       messageHandler.outgoingNotification<PublishDiagnosticsParams>(
198           "textDocument/publishDiagnostics");
199 
200   // Run the main loop of the transport.
201   LogicalResult result = success();
202   if (llvm::Error error = impl->transport.run(messageHandler)) {
203     Logger::error("Transport error: {0}", error);
204     llvm::consumeError(std::move(error));
205     result = failure();
206   } else {
207     result = success(impl->shutdownRequestReceived);
208   }
209   return result;
210 }
211