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                                       params.textDocument.version);
108   server.addOrUpdateDocument(params.textDocument.uri, params.textDocument.text,
109                              params.textDocument.version,
110                              diagParams.diagnostics);
111 
112   // Publish any recorded diagnostics.
113   publishDiagnostics(diagParams);
114 }
115 void LSPServer::Impl::onDocumentDidClose(
116     const DidCloseTextDocumentParams &params) {
117   Optional<int64_t> version = server.removeDocument(params.textDocument.uri);
118   if (!version)
119     return;
120 
121   // Empty out the diagnostics shown for this document. This will clear out
122   // anything currently displayed by the client for this document (e.g. in the
123   // "Problems" pane of VSCode).
124   publishDiagnostics(
125       PublishDiagnosticsParams(params.textDocument.uri, *version));
126 }
127 void LSPServer::Impl::onDocumentDidChange(
128     const DidChangeTextDocumentParams &params) {
129   // TODO: We currently only support full document updates, we should refactor
130   // to avoid this.
131   if (params.contentChanges.size() != 1)
132     return;
133   PublishDiagnosticsParams diagParams(params.textDocument.uri,
134                                       params.textDocument.version);
135   server.addOrUpdateDocument(
136       params.textDocument.uri, params.contentChanges.front().text,
137       params.textDocument.version, diagParams.diagnostics);
138 
139   // Publish any recorded diagnostics.
140   publishDiagnostics(diagParams);
141 }
142 
143 //===----------------------------------------------------------------------===//
144 // Definitions and References
145 
146 void LSPServer::Impl::onGoToDefinition(const TextDocumentPositionParams &params,
147                                        Callback<std::vector<Location>> reply) {
148   std::vector<Location> locations;
149   server.getLocationsOf(params.textDocument.uri, params.position, locations);
150   reply(std::move(locations));
151 }
152 
153 void LSPServer::Impl::onReference(const ReferenceParams &params,
154                                   Callback<std::vector<Location>> reply) {
155   std::vector<Location> locations;
156   server.findReferencesOf(params.textDocument.uri, params.position, locations);
157   reply(std::move(locations));
158 }
159 
160 //===----------------------------------------------------------------------===//
161 // Hover
162 
163 void LSPServer::Impl::onHover(const TextDocumentPositionParams &params,
164                               Callback<Optional<Hover>> reply) {
165   reply(server.findHover(params.textDocument.uri, params.position));
166 }
167 
168 //===----------------------------------------------------------------------===//
169 // LSPServer
170 //===----------------------------------------------------------------------===//
171 
172 LSPServer::LSPServer(MLIRServer &server, JSONTransport &transport)
173     : impl(std::make_unique<Impl>(server, transport)) {}
174 LSPServer::~LSPServer() {}
175 
176 LogicalResult LSPServer::run() {
177   MessageHandler messageHandler(impl->transport);
178 
179   // Initialization
180   messageHandler.method("initialize", impl.get(), &Impl::onInitialize);
181   messageHandler.notification("initialized", impl.get(), &Impl::onInitialized);
182   messageHandler.method("shutdown", impl.get(), &Impl::onShutdown);
183 
184   // Document Changes
185   messageHandler.notification("textDocument/didOpen", impl.get(),
186                               &Impl::onDocumentDidOpen);
187   messageHandler.notification("textDocument/didClose", impl.get(),
188                               &Impl::onDocumentDidClose);
189   messageHandler.notification("textDocument/didChange", impl.get(),
190                               &Impl::onDocumentDidChange);
191 
192   // Definitions and References
193   messageHandler.method("textDocument/definition", impl.get(),
194                         &Impl::onGoToDefinition);
195   messageHandler.method("textDocument/references", impl.get(),
196                         &Impl::onReference);
197 
198   // Hover
199   messageHandler.method("textDocument/hover", impl.get(), &Impl::onHover);
200 
201   // Diagnostics
202   impl->publishDiagnostics =
203       messageHandler.outgoingNotification<PublishDiagnosticsParams>(
204           "textDocument/publishDiagnostics");
205 
206   // Run the main loop of the transport.
207   LogicalResult result = success();
208   if (llvm::Error error = impl->transport.run(messageHandler)) {
209     Logger::error("Transport error: {0}", error);
210     llvm::consumeError(std::move(error));
211     result = failure();
212   } else {
213     result = success(impl->shutdownRequestReceived);
214   }
215   return result;
216 }
217