1 //===- MLIRServer.cpp - MLIR Generic 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 "MLIRServer.h"
10 #include "lsp/Logging.h"
11 #include "lsp/Protocol.h"
12 #include "mlir/IR/Operation.h"
13 #include "mlir/Parser.h"
14 #include "mlir/Parser/AsmParserState.h"
15 #include "llvm/Support/SourceMgr.h"
16 
17 using namespace mlir;
18 
19 /// Returns a language server position for the given source location.
20 static lsp::Position getPosFromLoc(llvm::SourceMgr &mgr, llvm::SMLoc loc) {
21   std::pair<unsigned, unsigned> lineAndCol = mgr.getLineAndColumn(loc);
22   lsp::Position pos;
23   pos.line = lineAndCol.first - 1;
24   pos.character = lineAndCol.second;
25   return pos;
26 }
27 
28 /// Returns a source location from the given language server position.
29 static llvm::SMLoc getPosFromLoc(llvm::SourceMgr &mgr, lsp::Position pos) {
30   return mgr.FindLocForLineAndColumn(mgr.getMainFileID(), pos.line + 1,
31                                      pos.character);
32 }
33 
34 /// Returns a language server range for the given source range.
35 static lsp::Range getRangeFromLoc(llvm::SourceMgr &mgr, llvm::SMRange range) {
36   // lsp::Range is an inclusive range, SMRange is half-open.
37   llvm::SMLoc inclusiveEnd =
38       llvm::SMLoc::getFromPointer(range.End.getPointer() - 1);
39   return {getPosFromLoc(mgr, range.Start), getPosFromLoc(mgr, inclusiveEnd)};
40 }
41 
42 /// Returns a language server location from the given source range.
43 static lsp::Location getLocationFromLoc(llvm::SourceMgr &mgr,
44                                         llvm::SMRange range,
45                                         const lsp::URIForFile &uri) {
46   return lsp::Location{uri, getRangeFromLoc(mgr, range)};
47 }
48 
49 /// Returns a language server location from the given MLIR file location.
50 static Optional<lsp::Location> getLocationFromLoc(FileLineColLoc loc) {
51   llvm::Expected<lsp::URIForFile> sourceURI =
52       lsp::URIForFile::fromFile(loc.getFilename());
53   if (!sourceURI) {
54     lsp::Logger::error("Failed to create URI for file `{0}`: {1}",
55                        loc.getFilename(),
56                        llvm::toString(sourceURI.takeError()));
57     return llvm::None;
58   }
59 
60   lsp::Position position;
61   position.line = loc.getLine() - 1;
62   position.character = loc.getColumn();
63   return lsp::Location{*sourceURI, lsp::Range{position, position}};
64 }
65 
66 /// Collect all of the locations from the given MLIR location that are not
67 /// contained within the given URI.
68 static void collectLocationsFromLoc(Location loc,
69                                     std::vector<lsp::Location> &locations,
70                                     const lsp::URIForFile &uri) {
71   SetVector<Location> visitedLocs;
72   loc->walk([&](Location nestedLoc) {
73     FileLineColLoc fileLoc = nestedLoc.dyn_cast<FileLineColLoc>();
74     if (!fileLoc || !visitedLocs.insert(nestedLoc))
75       return WalkResult::advance();
76 
77     Optional<lsp::Location> sourceLoc = getLocationFromLoc(fileLoc);
78     if (sourceLoc && sourceLoc->uri != uri)
79       locations.push_back(*sourceLoc);
80     return WalkResult::advance();
81   });
82 }
83 
84 /// Returns true if the given range contains the given source location. Note
85 /// that this has slightly different behavior than SMRange because it is
86 /// inclusive of the end location.
87 static bool contains(llvm::SMRange range, llvm::SMLoc loc) {
88   return range.Start.getPointer() <= loc.getPointer() &&
89          loc.getPointer() <= range.End.getPointer();
90 }
91 
92 /// Returns true if the given location is contained by the definition or one of
93 /// the uses of the given SMDefinition.
94 static bool isDefOrUse(const AsmParserState::SMDefinition &def,
95                        llvm::SMLoc loc) {
96   auto isUseFn = [&](const llvm::SMRange &range) {
97     return contains(range, loc);
98   };
99   return contains(def.loc, loc) || llvm::any_of(def.uses, isUseFn);
100 }
101 
102 //===----------------------------------------------------------------------===//
103 // MLIRDocument
104 //===----------------------------------------------------------------------===//
105 
106 namespace {
107 /// This class represents all of the information pertaining to a specific MLIR
108 /// document.
109 struct MLIRDocument {
110   MLIRDocument(const lsp::URIForFile &uri, StringRef contents,
111                DialectRegistry &registry);
112 
113   void getLocationsOf(const lsp::URIForFile &uri, const lsp::Position &defPos,
114                       std::vector<lsp::Location> &locations);
115   void findReferencesOf(const lsp::URIForFile &uri, const lsp::Position &pos,
116                         std::vector<lsp::Location> &references);
117 
118   /// The context used to hold the state contained by the parsed document.
119   MLIRContext context;
120 
121   /// The high level parser state used to find definitions and references within
122   /// the source file.
123   AsmParserState asmState;
124 
125   /// The container for the IR parsed from the input file.
126   Block parsedIR;
127 
128   /// The source manager containing the contents of the input file.
129   llvm::SourceMgr sourceMgr;
130 };
131 } // namespace
132 
133 MLIRDocument::MLIRDocument(const lsp::URIForFile &uri, StringRef contents,
134                            DialectRegistry &registry)
135     : context(registry) {
136   context.allowUnregisteredDialects();
137   ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diag) {
138     // TODO: What should we do with these diagnostics?
139     //  * Cache and show to the user?
140     //  * Ignore?
141     lsp::Logger::error("Error when parsing MLIR document `{0}`: `{1}`",
142                        uri.file(), diag.str());
143   });
144 
145   // Try to parsed the given IR string.
146   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
147   if (!memBuffer) {
148     lsp::Logger::error("Failed to create memory buffer for file", uri.file());
149     return;
150   }
151 
152   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), llvm::SMLoc());
153   if (failed(
154           parseSourceFile(sourceMgr, &parsedIR, &context, nullptr, &asmState)))
155     return;
156 }
157 
158 void MLIRDocument::getLocationsOf(const lsp::URIForFile &uri,
159                                   const lsp::Position &defPos,
160                                   std::vector<lsp::Location> &locations) {
161   llvm::SMLoc posLoc = getPosFromLoc(sourceMgr, defPos);
162 
163   // Functor used to check if an SM definition contains the position.
164   auto containsPosition = [&](const AsmParserState::SMDefinition &def) {
165     if (!isDefOrUse(def, posLoc))
166       return false;
167     locations.push_back(getLocationFromLoc(sourceMgr, def.loc, uri));
168     return true;
169   };
170 
171   // Check all definitions related to operations.
172   for (const AsmParserState::OperationDefinition &op : asmState.getOpDefs()) {
173     if (contains(op.loc, posLoc))
174       return collectLocationsFromLoc(op.op->getLoc(), locations, uri);
175     for (const auto &result : op.resultGroups)
176       if (containsPosition(result.second))
177         return collectLocationsFromLoc(op.op->getLoc(), locations, uri);
178   }
179 
180   // Check all definitions related to blocks.
181   for (const AsmParserState::BlockDefinition &block : asmState.getBlockDefs()) {
182     if (containsPosition(block.definition))
183       return;
184     for (const AsmParserState::SMDefinition &arg : block.arguments)
185       if (containsPosition(arg))
186         return;
187   }
188 }
189 
190 void MLIRDocument::findReferencesOf(const lsp::URIForFile &uri,
191                                     const lsp::Position &pos,
192                                     std::vector<lsp::Location> &references) {
193   // Functor used to append all of the definitions/uses of the given SM
194   // definition to the reference list.
195   auto appendSMDef = [&](const AsmParserState::SMDefinition &def) {
196     references.push_back(getLocationFromLoc(sourceMgr, def.loc, uri));
197     for (const llvm::SMRange &use : def.uses)
198       references.push_back(getLocationFromLoc(sourceMgr, use, uri));
199   };
200 
201   llvm::SMLoc posLoc = getPosFromLoc(sourceMgr, pos);
202 
203   // Check all definitions related to operations.
204   for (const AsmParserState::OperationDefinition &op : asmState.getOpDefs()) {
205     if (contains(op.loc, posLoc)) {
206       for (const auto &result : op.resultGroups)
207         appendSMDef(result.second);
208       return;
209     }
210     for (const auto &result : op.resultGroups)
211       if (isDefOrUse(result.second, posLoc))
212         return appendSMDef(result.second);
213   }
214 
215   // Check all definitions related to blocks.
216   for (const AsmParserState::BlockDefinition &block : asmState.getBlockDefs()) {
217     if (isDefOrUse(block.definition, posLoc))
218       return appendSMDef(block.definition);
219 
220     for (const AsmParserState::SMDefinition &arg : block.arguments)
221       if (isDefOrUse(arg, posLoc))
222         return appendSMDef(arg);
223   }
224 }
225 
226 //===----------------------------------------------------------------------===//
227 // MLIRServer::Impl
228 //===----------------------------------------------------------------------===//
229 
230 struct lsp::MLIRServer::Impl {
231   Impl(DialectRegistry &registry) : registry(registry) {}
232 
233   /// The registry containing dialects that can be recognized in parsed .mlir
234   /// files.
235   DialectRegistry &registry;
236 
237   /// The documents held by the server, mapped by their URI file name.
238   llvm::StringMap<std::unique_ptr<MLIRDocument>> documents;
239 };
240 
241 //===----------------------------------------------------------------------===//
242 // MLIRServer
243 //===----------------------------------------------------------------------===//
244 
245 lsp::MLIRServer::MLIRServer(DialectRegistry &registry)
246     : impl(std::make_unique<Impl>(registry)) {}
247 lsp::MLIRServer::~MLIRServer() {}
248 
249 void lsp::MLIRServer::addOrUpdateDocument(const URIForFile &uri,
250                                           StringRef contents) {
251   impl->documents[uri.file()] =
252       std::make_unique<MLIRDocument>(uri, contents, impl->registry);
253 }
254 
255 void lsp::MLIRServer::removeDocument(const URIForFile &uri) {
256   impl->documents.erase(uri.file());
257 }
258 
259 void lsp::MLIRServer::getLocationsOf(const URIForFile &uri,
260                                      const Position &defPos,
261                                      std::vector<Location> &locations) {
262   auto fileIt = impl->documents.find(uri.file());
263   if (fileIt != impl->documents.end())
264     fileIt->second->getLocationsOf(uri, defPos, locations);
265 }
266 
267 void lsp::MLIRServer::findReferencesOf(const URIForFile &uri,
268                                        const Position &pos,
269                                        std::vector<Location> &references) {
270   auto fileIt = impl->documents.find(uri.file());
271   if (fileIt != impl->documents.end())
272     fileIt->second->findReferencesOf(uri, pos, references);
273 }
274