1 //===- TableGenServer.cpp - TableGen 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 "TableGenServer.h"
10 
11 #include "../lsp-server-support/CompilationDatabase.h"
12 #include "../lsp-server-support/Logging.h"
13 #include "../lsp-server-support/Protocol.h"
14 #include "../lsp-server-support/SourceMgrUtils.h"
15 #include "llvm/ADT/IntervalMap.h"
16 #include "llvm/ADT/PointerUnion.h"
17 #include "llvm/ADT/StringMap.h"
18 #include "llvm/ADT/StringSet.h"
19 #include "llvm/ADT/TypeSwitch.h"
20 #include "llvm/Support/FileSystem.h"
21 #include "llvm/Support/Path.h"
22 #include "llvm/TableGen/Parser.h"
23 #include "llvm/TableGen/Record.h"
24 
25 using namespace mlir;
26 
27 /// Returns a language server uri for the given source location. `mainFileURI`
28 /// corresponds to the uri for the main file of the source manager.
29 static lsp::URIForFile getURIFromLoc(const llvm::SourceMgr &mgr, SMLoc loc,
30                                      const lsp::URIForFile &mainFileURI) {
31   int bufferId = mgr.FindBufferContainingLoc(loc);
32   if (bufferId == 0 || bufferId == static_cast<int>(mgr.getMainFileID()))
33     return mainFileURI;
34   llvm::Expected<lsp::URIForFile> fileForLoc = lsp::URIForFile::fromFile(
35       mgr.getBufferInfo(bufferId).Buffer->getBufferIdentifier());
36   if (fileForLoc)
37     return *fileForLoc;
38   lsp::Logger::error("Failed to create URI for include file: {0}",
39                      llvm::toString(fileForLoc.takeError()));
40   return mainFileURI;
41 }
42 
43 /// Returns a language server location from the given source range.
44 static lsp::Location getLocationFromLoc(llvm::SourceMgr &mgr, SMRange loc,
45                                         const lsp::URIForFile &uri) {
46   return lsp::Location(getURIFromLoc(mgr, loc.Start, uri),
47                        lsp::Range(mgr, loc));
48 }
49 static lsp::Location getLocationFromLoc(llvm::SourceMgr &mgr, SMLoc loc,
50                                         const lsp::URIForFile &uri) {
51   return getLocationFromLoc(mgr, lsp::convertTokenLocToRange(loc), uri);
52 }
53 
54 /// Convert the given TableGen diagnostic to the LSP form.
55 static Optional<lsp::Diagnostic>
56 getLspDiagnoticFromDiag(const llvm::SMDiagnostic &diag,
57                         const lsp::URIForFile &uri) {
58   auto *sourceMgr = const_cast<llvm::SourceMgr *>(diag.getSourceMgr());
59   if (!sourceMgr || !diag.getLoc().isValid())
60     return llvm::None;
61 
62   lsp::Diagnostic lspDiag;
63   lspDiag.source = "tablegen";
64   lspDiag.category = "Parse Error";
65 
66   // Try to grab a file location for this diagnostic.
67   lsp::Location loc = getLocationFromLoc(*sourceMgr, diag.getLoc(), uri);
68   lspDiag.range = loc.range;
69 
70   // Skip diagnostics that weren't emitted within the main file.
71   if (loc.uri != uri)
72     return llvm::None;
73 
74   // Convert the severity for the diagnostic.
75   switch (diag.getKind()) {
76   case llvm::SourceMgr::DK_Warning:
77     lspDiag.severity = lsp::DiagnosticSeverity::Warning;
78     break;
79   case llvm::SourceMgr::DK_Error:
80     lspDiag.severity = lsp::DiagnosticSeverity::Error;
81     break;
82   case llvm::SourceMgr::DK_Note:
83     // Notes are emitted separately from the main diagnostic, so we just treat
84     // them as remarks given that we can't determine the diagnostic to relate
85     // them to.
86   case llvm::SourceMgr::DK_Remark:
87     lspDiag.severity = lsp::DiagnosticSeverity::Information;
88     break;
89   }
90   lspDiag.message = diag.getMessage().str();
91 
92   return lspDiag;
93 }
94 
95 //===----------------------------------------------------------------------===//
96 // TableGenIndex
97 //===----------------------------------------------------------------------===//
98 
99 namespace {
100 /// This class represents a single symbol definition within a TableGen index. It
101 /// contains the definition of the symbol, the location of the symbol, and any
102 /// recorded references.
103 struct TableGenIndexSymbol {
104   TableGenIndexSymbol(const llvm::Record *record)
105       : definition(record),
106         defLoc(lsp::convertTokenLocToRange(record->getLoc().front())) {}
107   TableGenIndexSymbol(const llvm::RecordVal *value)
108       : definition(value),
109         defLoc(lsp::convertTokenLocToRange(value->getLoc())) {}
110 
111   /// The main definition of the symbol.
112   PointerUnion<const llvm::Record *, const llvm::RecordVal *> definition;
113 
114   /// The source location of the definition.
115   SMRange defLoc;
116 
117   /// The source location of the references of the definition.
118   SmallVector<SMRange> references;
119 };
120 
121 /// This class provides an index for definitions/uses within a TableGen
122 /// document. It provides efficient lookup of a definition given an input source
123 /// range.
124 class TableGenIndex {
125 public:
126   TableGenIndex() : intervalMap(allocator) {}
127 
128   /// Initialize the index with the given RecordKeeper.
129   void initialize(const llvm::RecordKeeper &records);
130 
131   /// Lookup a symbol for the given location. Returns nullptr if no symbol could
132   /// be found. If provided, `overlappedRange` is set to the range that the
133   /// provided `loc` overlapped with.
134   const TableGenIndexSymbol *lookup(SMLoc loc,
135                                     SMRange *overlappedRange = nullptr) const;
136 
137 private:
138   /// The type of interval map used to store source references. SMRange is
139   /// half-open, so we also need to use a half-open interval map.
140   using MapT = llvm::IntervalMap<
141       const char *, const TableGenIndexSymbol *,
142       llvm::IntervalMapImpl::NodeSizer<const char *,
143                                        const TableGenIndexSymbol *>::LeafSize,
144       llvm::IntervalMapHalfOpenInfo<const char *>>;
145 
146   /// An allocator for the interval map.
147   MapT::Allocator allocator;
148 
149   /// An interval map containing a corresponding definition mapped to a source
150   /// interval.
151   MapT intervalMap;
152 
153   /// A mapping between definitions and their corresponding symbol.
154   DenseMap<const void *, std::unique_ptr<TableGenIndexSymbol>> defToSymbol;
155 };
156 } // namespace
157 
158 void TableGenIndex::initialize(const llvm::RecordKeeper &records) {
159   auto getOrInsertDef = [&](const auto *def) -> TableGenIndexSymbol * {
160     auto it = defToSymbol.try_emplace(def, nullptr);
161     if (it.second)
162       it.first->second = std::make_unique<TableGenIndexSymbol>(def);
163     return &*it.first->second;
164   };
165   auto insertRef = [&](TableGenIndexSymbol *sym, SMRange refLoc,
166                        bool isDef = false) {
167     const char *startLoc = refLoc.Start.getPointer();
168     const char *endLoc = refLoc.End.getPointer();
169 
170     // If the location we got was empty, try to lex a token from the start
171     // location.
172     if (startLoc == endLoc) {
173       refLoc = lsp::convertTokenLocToRange(SMLoc::getFromPointer(startLoc));
174       startLoc = refLoc.Start.getPointer();
175       endLoc = refLoc.End.getPointer();
176 
177       // If the location is still empty, bail on trying to use this reference
178       // location.
179       if (startLoc == endLoc)
180         return;
181     }
182 
183     // Check to see if a symbol is already attached to this location.
184     // IntervalMap doesn't allow overlapping inserts, and we don't really
185     // want multiple symbols attached to a source location anyways. This
186     // shouldn't really happen in practice, but we should handle it gracefully.
187     if (!intervalMap.overlaps(startLoc, endLoc))
188       intervalMap.insert(startLoc, endLoc, sym);
189 
190     if (!isDef)
191       sym->references.push_back(refLoc);
192   };
193   auto classes =
194       llvm::make_pointee_range(llvm::make_second_range(records.getClasses()));
195   auto defs =
196       llvm::make_pointee_range(llvm::make_second_range(records.getDefs()));
197   for (const llvm::Record &def : llvm::concat<llvm::Record>(classes, defs)) {
198     auto *sym = getOrInsertDef(&def);
199     insertRef(sym, sym->defLoc, /*isDef=*/true);
200 
201     // Add references to the definition.
202     for (SMLoc loc : def.getLoc().drop_front())
203       insertRef(sym, lsp::convertTokenLocToRange(loc));
204 
205     // Add references to any super classes.
206     for (auto &it : def.getSuperClasses())
207       insertRef(getOrInsertDef(it.first),
208                 lsp::convertTokenLocToRange(it.second.Start));
209 
210     // Add definitions for any values.
211     for (const llvm::RecordVal &value : def.getValues()) {
212       auto *sym = getOrInsertDef(&value);
213       insertRef(sym, sym->defLoc, /*isDef=*/true);
214     }
215   }
216 }
217 
218 const TableGenIndexSymbol *
219 TableGenIndex::lookup(SMLoc loc, SMRange *overlappedRange) const {
220   auto it = intervalMap.find(loc.getPointer());
221   if (!it.valid() || loc.getPointer() < it.start())
222     return nullptr;
223 
224   if (overlappedRange) {
225     *overlappedRange = SMRange(SMLoc::getFromPointer(it.start()),
226                                SMLoc::getFromPointer(it.stop()));
227   }
228   return it.value();
229 }
230 
231 //===----------------------------------------------------------------------===//
232 // TableGenTextFile
233 //===----------------------------------------------------------------------===//
234 
235 namespace {
236 /// This class represents a text file containing one or more TableGen documents.
237 class TableGenTextFile {
238 public:
239   TableGenTextFile(const lsp::URIForFile &uri, StringRef fileContents,
240                    int64_t version,
241                    const std::vector<std::string> &extraIncludeDirs,
242                    std::vector<lsp::Diagnostic> &diagnostics);
243 
244   /// Return the current version of this text file.
245   int64_t getVersion() const { return version; }
246 
247   //===--------------------------------------------------------------------===//
248   // Definitions and References
249   //===--------------------------------------------------------------------===//
250 
251   void getLocationsOf(const lsp::URIForFile &uri, const lsp::Position &defPos,
252                       std::vector<lsp::Location> &locations);
253   void findReferencesOf(const lsp::URIForFile &uri, const lsp::Position &pos,
254                         std::vector<lsp::Location> &references);
255 
256   //===--------------------------------------------------------------------===//
257   // Document Links
258   //===--------------------------------------------------------------------===//
259 
260   void getDocumentLinks(const lsp::URIForFile &uri,
261                         std::vector<lsp::DocumentLink> &links);
262 
263   //===--------------------------------------------------------------------===//
264   // Hover
265   //===--------------------------------------------------------------------===//
266 
267   Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
268                                  const lsp::Position &hoverPos);
269 
270 private:
271   /// The full string contents of the file.
272   std::string contents;
273 
274   /// The version of this file.
275   int64_t version;
276 
277   /// The include directories for this file.
278   std::vector<std::string> includeDirs;
279 
280   /// The source manager containing the contents of the input file.
281   llvm::SourceMgr sourceMgr;
282 
283   /// The record keeper containing the parsed tablegen constructs.
284   llvm::RecordKeeper recordKeeper;
285 
286   /// The index of the parsed file.
287   TableGenIndex index;
288 
289   /// The set of includes of the parsed file.
290   SmallVector<lsp::SourceMgrInclude> parsedIncludes;
291 };
292 } // namespace
293 
294 TableGenTextFile::TableGenTextFile(
295     const lsp::URIForFile &uri, StringRef fileContents, int64_t version,
296     const std::vector<std::string> &extraIncludeDirs,
297     std::vector<lsp::Diagnostic> &diagnostics)
298     : contents(fileContents.str()), version(version) {
299   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
300   if (!memBuffer) {
301     lsp::Logger::error("Failed to create memory buffer for file", uri.file());
302     return;
303   }
304 
305   // Build the set of include directories for this file.
306   llvm::SmallString<32> uriDirectory(uri.file());
307   llvm::sys::path::remove_filename(uriDirectory);
308   includeDirs.push_back(uriDirectory.str().str());
309   includeDirs.insert(includeDirs.end(), extraIncludeDirs.begin(),
310                      extraIncludeDirs.end());
311 
312   sourceMgr.setIncludeDirs(includeDirs);
313   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
314 
315   // This class provides a context argument for the llvm::SourceMgr diagnostic
316   // handler.
317   struct DiagHandlerContext {
318     std::vector<lsp::Diagnostic> &diagnostics;
319     const lsp::URIForFile &uri;
320   } handlerContext{diagnostics, uri};
321 
322   // Set the diagnostic handler for the tablegen source manager.
323   sourceMgr.setDiagHandler(
324       [](const llvm::SMDiagnostic &diag, void *rawHandlerContext) {
325         auto *ctx = reinterpret_cast<DiagHandlerContext *>(rawHandlerContext);
326         if (auto lspDiag = getLspDiagnoticFromDiag(diag, ctx->uri))
327           ctx->diagnostics.push_back(*lspDiag);
328       },
329       &handlerContext);
330   bool failedToParse = llvm::TableGenParseFile(sourceMgr, recordKeeper);
331 
332   // Process all of the include files.
333   lsp::gatherIncludeFiles(sourceMgr, parsedIncludes);
334   if (failedToParse)
335     return;
336 
337   // If we successfully parsed the file, we can now build the index.
338   index.initialize(recordKeeper);
339 }
340 
341 //===----------------------------------------------------------------------===//
342 // TableGenTextFile: Definitions and References
343 //===----------------------------------------------------------------------===//
344 
345 void TableGenTextFile::getLocationsOf(const lsp::URIForFile &uri,
346                                       const lsp::Position &defPos,
347                                       std::vector<lsp::Location> &locations) {
348   SMLoc posLoc = defPos.getAsSMLoc(sourceMgr);
349   const TableGenIndexSymbol *symbol = index.lookup(posLoc);
350   if (!symbol)
351     return;
352 
353   locations.push_back(getLocationFromLoc(sourceMgr, symbol->defLoc, uri));
354 }
355 
356 void TableGenTextFile::findReferencesOf(
357     const lsp::URIForFile &uri, const lsp::Position &pos,
358     std::vector<lsp::Location> &references) {
359   SMLoc posLoc = pos.getAsSMLoc(sourceMgr);
360   const TableGenIndexSymbol *symbol = index.lookup(posLoc);
361   if (!symbol)
362     return;
363 
364   references.push_back(getLocationFromLoc(sourceMgr, symbol->defLoc, uri));
365   for (SMRange refLoc : symbol->references)
366     references.push_back(getLocationFromLoc(sourceMgr, refLoc, uri));
367 }
368 
369 //===--------------------------------------------------------------------===//
370 // TableGenTextFile: Document Links
371 //===--------------------------------------------------------------------===//
372 
373 void TableGenTextFile::getDocumentLinks(const lsp::URIForFile &uri,
374                                         std::vector<lsp::DocumentLink> &links) {
375   for (const lsp::SourceMgrInclude &include : parsedIncludes)
376     links.emplace_back(include.range, include.uri);
377 }
378 
379 //===----------------------------------------------------------------------===//
380 // TableGenTextFile: Hover
381 //===----------------------------------------------------------------------===//
382 
383 Optional<lsp::Hover>
384 TableGenTextFile::findHover(const lsp::URIForFile &uri,
385                             const lsp::Position &hoverPos) {
386   // Check for a reference to an include.
387   for (const lsp::SourceMgrInclude &include : parsedIncludes)
388     if (include.range.contains(hoverPos))
389       return include.buildHover();
390   return llvm::None;
391 }
392 
393 //===----------------------------------------------------------------------===//
394 // TableGenServer::Impl
395 //===----------------------------------------------------------------------===//
396 
397 struct lsp::TableGenServer::Impl {
398   explicit Impl(const Options &options)
399       : options(options), compilationDatabase(options.compilationDatabases) {}
400 
401   /// TableGen LSP options.
402   const Options &options;
403 
404   /// The compilation database containing additional information for files
405   /// passed to the server.
406   lsp::CompilationDatabase compilationDatabase;
407 
408   /// The files held by the server, mapped by their URI file name.
409   llvm::StringMap<std::unique_ptr<TableGenTextFile>> files;
410 };
411 
412 //===----------------------------------------------------------------------===//
413 // TableGenServer
414 //===----------------------------------------------------------------------===//
415 
416 lsp::TableGenServer::TableGenServer(const Options &options)
417     : impl(std::make_unique<Impl>(options)) {}
418 lsp::TableGenServer::~TableGenServer() = default;
419 
420 void lsp::TableGenServer::addOrUpdateDocument(
421     const URIForFile &uri, StringRef contents, int64_t version,
422     std::vector<Diagnostic> &diagnostics) {
423   // Build the set of additional include directories.
424   std::vector<std::string> additionalIncludeDirs = impl->options.extraDirs;
425   const auto &fileInfo = impl->compilationDatabase.getFileInfo(uri.file());
426   llvm::append_range(additionalIncludeDirs, fileInfo.includeDirs);
427 
428   impl->files[uri.file()] = std::make_unique<TableGenTextFile>(
429       uri, contents, version, additionalIncludeDirs, diagnostics);
430 }
431 
432 Optional<int64_t> lsp::TableGenServer::removeDocument(const URIForFile &uri) {
433   auto it = impl->files.find(uri.file());
434   if (it == impl->files.end())
435     return llvm::None;
436 
437   int64_t version = it->second->getVersion();
438   impl->files.erase(it);
439   return version;
440 }
441 
442 void lsp::TableGenServer::getLocationsOf(const URIForFile &uri,
443                                          const Position &defPos,
444                                          std::vector<Location> &locations) {
445   auto fileIt = impl->files.find(uri.file());
446   if (fileIt != impl->files.end())
447     fileIt->second->getLocationsOf(uri, defPos, locations);
448 }
449 
450 void lsp::TableGenServer::findReferencesOf(const URIForFile &uri,
451                                            const Position &pos,
452                                            std::vector<Location> &references) {
453   auto fileIt = impl->files.find(uri.file());
454   if (fileIt != impl->files.end())
455     fileIt->second->findReferencesOf(uri, pos, references);
456 }
457 
458 void lsp::TableGenServer::getDocumentLinks(
459     const URIForFile &uri, std::vector<DocumentLink> &documentLinks) {
460   auto fileIt = impl->files.find(uri.file());
461   if (fileIt != impl->files.end())
462     return fileIt->second->getDocumentLinks(uri, documentLinks);
463 }
464 
465 Optional<lsp::Hover> lsp::TableGenServer::findHover(const URIForFile &uri,
466                                                     const Position &hoverPos) {
467   auto fileIt = impl->files.find(uri.file());
468   if (fileIt != impl->files.end())
469     return fileIt->second->findHover(uri, hoverPos);
470   return llvm::None;
471 }
472