1 //===- LSPServer.cpp - PDLL 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 
11 #include "../lsp-server-support/Logging.h"
12 #include "../lsp-server-support/Protocol.h"
13 #include "../lsp-server-support/Transport.h"
14 #include "PDLLServer.h"
15 #include "llvm/ADT/FunctionExtras.h"
16 #include "llvm/ADT/StringMap.h"
17 
18 #define DEBUG_TYPE "pdll-lsp-server"
19 
20 using namespace mlir;
21 using namespace mlir::lsp;
22 
23 //===----------------------------------------------------------------------===//
24 // LSPServer
25 //===----------------------------------------------------------------------===//
26 
27 namespace {
28 struct LSPServer {
29   LSPServer(PDLLServer &server, JSONTransport &transport)
30       : server(server), transport(transport) {}
31 
32   //===--------------------------------------------------------------------===//
33   // Initialization
34 
35   void onInitialize(const InitializeParams &params,
36                     Callback<llvm::json::Value> reply);
37   void onInitialized(const InitializedParams &params);
38   void onShutdown(const NoParams &params, Callback<std::nullptr_t> reply);
39 
40   //===--------------------------------------------------------------------===//
41   // Document Change
42 
43   void onDocumentDidOpen(const DidOpenTextDocumentParams &params);
44   void onDocumentDidClose(const DidCloseTextDocumentParams &params);
45   void onDocumentDidChange(const DidChangeTextDocumentParams &params);
46 
47   //===--------------------------------------------------------------------===//
48   // Definitions and References
49 
50   void onGoToDefinition(const TextDocumentPositionParams &params,
51                         Callback<std::vector<Location>> reply);
52   void onReference(const ReferenceParams &params,
53                    Callback<std::vector<Location>> reply);
54 
55   //===----------------------------------------------------------------------===//
56   // DocumentLink
57 
58   void onDocumentLink(const DocumentLinkParams &params,
59                       Callback<std::vector<DocumentLink>> reply);
60 
61   //===--------------------------------------------------------------------===//
62   // Hover
63 
64   void onHover(const TextDocumentPositionParams &params,
65                Callback<Optional<Hover>> reply);
66 
67   //===--------------------------------------------------------------------===//
68   // Document Symbols
69 
70   void onDocumentSymbol(const DocumentSymbolParams &params,
71                         Callback<std::vector<DocumentSymbol>> reply);
72 
73   //===--------------------------------------------------------------------===//
74   // Code Completion
75 
76   void onCompletion(const CompletionParams &params,
77                     Callback<CompletionList> reply);
78 
79   //===--------------------------------------------------------------------===//
80   // Signature Help
81 
82   void onSignatureHelp(const TextDocumentPositionParams &params,
83                        Callback<SignatureHelp> reply);
84 
85   //===--------------------------------------------------------------------===//
86   // Fields
87   //===--------------------------------------------------------------------===//
88 
89   PDLLServer &server;
90   JSONTransport &transport;
91 
92   /// An outgoing notification used to send diagnostics to the client when they
93   /// are ready to be processed.
94   OutgoingNotification<PublishDiagnosticsParams> publishDiagnostics;
95 
96   /// Used to indicate that the 'shutdown' request was received from the
97   /// Language Server client.
98   bool shutdownRequestReceived = false;
99 };
100 } // namespace
101 
102 //===----------------------------------------------------------------------===//
103 // Initialization
104 
105 void LSPServer::onInitialize(const InitializeParams &params,
106                              Callback<llvm::json::Value> reply) {
107   // Send a response with the capabilities of this server.
108   llvm::json::Object serverCaps{
109       {"textDocumentSync",
110        llvm::json::Object{
111            {"openClose", true},
112            {"change", (int)TextDocumentSyncKind::Full},
113            {"save", true},
114        }},
115       {"completionProvider",
116        llvm::json::Object{
117            {"allCommitCharacters",
118             {" ", "\t", "(", ")", "[", "]", "{",  "}", "<",
119              ">", ":",  ";", ",", "+", "-", "/",  "*", "%",
120              "^", "&",  "#", "?", ".", "=", "\"", "'", "|"}},
121            {"resolveProvider", false},
122            {"triggerCharacters",
123             {".", ">", "(", "{", ",", "<", ":", "[", " ", "\"", "/"}},
124        }},
125       {"signatureHelpProvider",
126        llvm::json::Object{
127            {"triggerCharacters", {"(", ","}},
128        }},
129       {"definitionProvider", true},
130       {"referencesProvider", true},
131       {"documentLinkProvider",
132        llvm::json::Object{
133            {"resolveProvider", false},
134        }},
135       {"hoverProvider", true},
136       {"documentSymbolProvider", true},
137   };
138 
139   llvm::json::Object result{
140       {{"serverInfo", llvm::json::Object{{"name", "mlir-pdll-lsp-server"},
141                                          {"version", "0.0.1"}}},
142        {"capabilities", std::move(serverCaps)}}};
143   reply(std::move(result));
144 }
145 void LSPServer::onInitialized(const InitializedParams &) {}
146 void LSPServer::onShutdown(const NoParams &, Callback<std::nullptr_t> reply) {
147   shutdownRequestReceived = true;
148   reply(nullptr);
149 }
150 
151 //===----------------------------------------------------------------------===//
152 // Document Change
153 
154 void LSPServer::onDocumentDidOpen(const DidOpenTextDocumentParams &params) {
155   PublishDiagnosticsParams diagParams(params.textDocument.uri,
156                                       params.textDocument.version);
157   server.addOrUpdateDocument(params.textDocument.uri, params.textDocument.text,
158                              params.textDocument.version,
159                              diagParams.diagnostics);
160 
161   // Publish any recorded diagnostics.
162   publishDiagnostics(diagParams);
163 }
164 void LSPServer::onDocumentDidClose(const DidCloseTextDocumentParams &params) {
165   Optional<int64_t> version = server.removeDocument(params.textDocument.uri);
166   if (!version)
167     return;
168 
169   // Empty out the diagnostics shown for this document. This will clear out
170   // anything currently displayed by the client for this document (e.g. in the
171   // "Problems" pane of VSCode).
172   publishDiagnostics(
173       PublishDiagnosticsParams(params.textDocument.uri, *version));
174 }
175 void LSPServer::onDocumentDidChange(const DidChangeTextDocumentParams &params) {
176   // TODO: We currently only support full document updates, we should refactor
177   // to avoid this.
178   if (params.contentChanges.size() != 1)
179     return;
180   PublishDiagnosticsParams diagParams(params.textDocument.uri,
181                                       params.textDocument.version);
182   server.addOrUpdateDocument(
183       params.textDocument.uri, params.contentChanges.front().text,
184       params.textDocument.version, diagParams.diagnostics);
185 
186   // Publish any recorded diagnostics.
187   publishDiagnostics(diagParams);
188 }
189 
190 //===----------------------------------------------------------------------===//
191 // Definitions and References
192 
193 void LSPServer::onGoToDefinition(const TextDocumentPositionParams &params,
194                                  Callback<std::vector<Location>> reply) {
195   std::vector<Location> locations;
196   server.getLocationsOf(params.textDocument.uri, params.position, locations);
197   reply(std::move(locations));
198 }
199 
200 void LSPServer::onReference(const ReferenceParams &params,
201                             Callback<std::vector<Location>> reply) {
202   std::vector<Location> locations;
203   server.findReferencesOf(params.textDocument.uri, params.position, locations);
204   reply(std::move(locations));
205 }
206 
207 //===----------------------------------------------------------------------===//
208 // DocumentLink
209 
210 void LSPServer::onDocumentLink(const DocumentLinkParams &params,
211                                Callback<std::vector<DocumentLink>> reply) {
212   std::vector<DocumentLink> links;
213   server.getDocumentLinks(params.textDocument.uri, links);
214   reply(std::move(links));
215 }
216 
217 //===----------------------------------------------------------------------===//
218 // Hover
219 
220 void LSPServer::onHover(const TextDocumentPositionParams &params,
221                         Callback<Optional<Hover>> reply) {
222   reply(server.findHover(params.textDocument.uri, params.position));
223 }
224 
225 //===----------------------------------------------------------------------===//
226 // Document Symbols
227 
228 void LSPServer::onDocumentSymbol(const DocumentSymbolParams &params,
229                                  Callback<std::vector<DocumentSymbol>> reply) {
230   std::vector<DocumentSymbol> symbols;
231   server.findDocumentSymbols(params.textDocument.uri, symbols);
232   reply(std::move(symbols));
233 }
234 
235 //===----------------------------------------------------------------------===//
236 // Code Completion
237 
238 void LSPServer::onCompletion(const CompletionParams &params,
239                              Callback<CompletionList> reply) {
240   reply(server.getCodeCompletion(params.textDocument.uri, params.position));
241 }
242 
243 //===----------------------------------------------------------------------===//
244 // Signature Help
245 
246 void LSPServer::onSignatureHelp(const TextDocumentPositionParams &params,
247                                 Callback<SignatureHelp> reply) {
248   reply(server.getSignatureHelp(params.textDocument.uri, params.position));
249 }
250 
251 //===----------------------------------------------------------------------===//
252 // Entry Point
253 //===----------------------------------------------------------------------===//
254 
255 LogicalResult mlir::lsp::runPdllLSPServer(PDLLServer &server,
256                                           JSONTransport &transport) {
257   LSPServer lspServer(server, transport);
258   MessageHandler messageHandler(transport);
259 
260   // Initialization
261   messageHandler.method("initialize", &lspServer, &LSPServer::onInitialize);
262   messageHandler.notification("initialized", &lspServer,
263                               &LSPServer::onInitialized);
264   messageHandler.method("shutdown", &lspServer, &LSPServer::onShutdown);
265 
266   // Document Changes
267   messageHandler.notification("textDocument/didOpen", &lspServer,
268                               &LSPServer::onDocumentDidOpen);
269   messageHandler.notification("textDocument/didClose", &lspServer,
270                               &LSPServer::onDocumentDidClose);
271   messageHandler.notification("textDocument/didChange", &lspServer,
272                               &LSPServer::onDocumentDidChange);
273 
274   // Definitions and References
275   messageHandler.method("textDocument/definition", &lspServer,
276                         &LSPServer::onGoToDefinition);
277   messageHandler.method("textDocument/references", &lspServer,
278                         &LSPServer::onReference);
279 
280   // Document Link
281   messageHandler.method("textDocument/documentLink", &lspServer,
282                         &LSPServer::onDocumentLink);
283 
284   // Hover
285   messageHandler.method("textDocument/hover", &lspServer, &LSPServer::onHover);
286 
287   // Document Symbols
288   messageHandler.method("textDocument/documentSymbol", &lspServer,
289                         &LSPServer::onDocumentSymbol);
290 
291   // Code Completion
292   messageHandler.method("textDocument/completion", &lspServer,
293                         &LSPServer::onCompletion);
294 
295   // Signature Help
296   messageHandler.method("textDocument/signatureHelp", &lspServer,
297                         &LSPServer::onSignatureHelp);
298 
299   // Diagnostics
300   lspServer.publishDiagnostics =
301       messageHandler.outgoingNotification<PublishDiagnosticsParams>(
302           "textDocument/publishDiagnostics");
303 
304   // Run the main loop of the transport.
305   if (llvm::Error error = transport.run(messageHandler)) {
306     Logger::error("Transport error: {0}", error);
307     llvm::consumeError(std::move(error));
308     return failure();
309   }
310   return success(lspServer.shutdownRequestReceived);
311 }
312