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