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   MLIRServer &server;
54   JSONTransport &transport;
55 
56   /// Used to indicate that the 'shutdown' request was received from the
57   /// Language Server client.
58   bool shutdownRequestReceived = false;
59 };
60 
61 //===----------------------------------------------------------------------===//
62 // Initialization
63 
64 void LSPServer::Impl::onInitialize(const InitializeParams &params,
65                                    Callback<llvm::json::Value> reply) {
66   llvm::json::Object serverCaps{
67       {"textDocumentSync",
68        llvm::json::Object{
69            {"openClose", true},
70            {"change", (int)TextDocumentSyncKind::Full},
71            {"save", true},
72        }},
73       {"definitionProvider", true},
74       {"referencesProvider", true},
75   };
76 
77   llvm::json::Object result{
78       {{"serverInfo",
79         llvm::json::Object{{"name", "mlir-lsp-server"}, {"version", "0.0.0"}}},
80        {"capabilities", std::move(serverCaps)}}};
81   reply(std::move(result));
82 }
83 void LSPServer::Impl::onInitialized(const InitializedParams &) {}
84 void LSPServer::Impl::onShutdown(const NoParams &,
85                                  Callback<std::nullptr_t> reply) {
86   shutdownRequestReceived = true;
87   reply(nullptr);
88 }
89 
90 //===----------------------------------------------------------------------===//
91 // Document Change
92 
93 void LSPServer::Impl::onDocumentDidOpen(
94     const DidOpenTextDocumentParams &params) {
95   server.addOrUpdateDocument(params.textDocument.uri, params.textDocument.text);
96 }
97 void LSPServer::Impl::onDocumentDidClose(
98     const DidCloseTextDocumentParams &params) {
99   server.removeDocument(params.textDocument.uri);
100 }
101 void LSPServer::Impl::onDocumentDidChange(
102     const DidChangeTextDocumentParams &params) {
103   // TODO: We currently only support full document updates, we should refactor
104   // to avoid this.
105   if (params.contentChanges.size() != 1)
106     return;
107   server.addOrUpdateDocument(params.textDocument.uri,
108                              params.contentChanges.front().text);
109 }
110 
111 //===----------------------------------------------------------------------===//
112 // Definitions and References
113 
114 void LSPServer::Impl::onGoToDefinition(const TextDocumentPositionParams &params,
115                                        Callback<std::vector<Location>> reply) {
116   std::vector<Location> locations;
117   server.getLocationsOf(params.textDocument.uri, params.position, locations);
118   reply(std::move(locations));
119 }
120 
121 void LSPServer::Impl::onReference(const ReferenceParams &params,
122                                   Callback<std::vector<Location>> reply) {
123   std::vector<Location> locations;
124   server.findReferencesOf(params.textDocument.uri, params.position, locations);
125   reply(std::move(locations));
126 }
127 
128 //===----------------------------------------------------------------------===//
129 // LSPServer
130 //===----------------------------------------------------------------------===//
131 
132 LSPServer::LSPServer(MLIRServer &server, JSONTransport &transport)
133     : impl(std::make_unique<Impl>(server, transport)) {}
134 LSPServer::~LSPServer() {}
135 
136 LogicalResult LSPServer::run() {
137   MessageHandler messageHandler(impl->transport);
138 
139   // Initialization
140   messageHandler.method("initialize", impl.get(), &Impl::onInitialize);
141   messageHandler.notification("initialized", impl.get(), &Impl::onInitialized);
142   messageHandler.method("shutdown", impl.get(), &Impl::onShutdown);
143 
144   // Document Changes
145   messageHandler.notification("textDocument/didOpen", impl.get(),
146                               &Impl::onDocumentDidOpen);
147   messageHandler.notification("textDocument/didClose", impl.get(),
148                               &Impl::onDocumentDidClose);
149   messageHandler.notification("textDocument/didChange", impl.get(),
150                               &Impl::onDocumentDidChange);
151 
152   // Definitions and References
153   messageHandler.method("textDocument/definition", impl.get(),
154                         &Impl::onGoToDefinition);
155   messageHandler.method("textDocument/references", impl.get(),
156                         &Impl::onReference);
157 
158   LogicalResult result = success();
159   if (llvm::Error error = impl->transport.run(messageHandler)) {
160     Logger::error("Transport error: {0}", error);
161     llvm::consumeError(std::move(error));
162     result = failure();
163   } else {
164     result = success(impl->shutdownRequestReceived);
165   }
166   return result;
167 }
168