mlir/lib/Tools/tblgen-lsp-server/TableGenServer.cpp - rust-lang/llvm-project - Git at Google

 //===- TableGenServer.cpp - TableGen Language Server ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//

 #include "TableGenServer.h"

 #include "../lsp-server-support/CompilationDatabase.h"
 #include "../lsp-server-support/Logging.h"
 #include "../lsp-server-support/Protocol.h"
 #include "../lsp-server-support/SourceMgrUtils.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/IntervalMap.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 #include "llvm/TableGen/Parser.h"
 #include "llvm/TableGen/Record.h"

 using namespace mlir;

 /// Returns a language server uri for the given source location. `mainFileURI`
 /// corresponds to the uri for the main file of the source manager.
 static lsp::URIForFile getURIFromLoc(const llvm::SourceMgr &mgr, SMLoc loc,
                                      const lsp::URIForFile &mainFileURI) {
   int bufferId = mgr.FindBufferContainingLoc(loc);
   if (bufferId == 0 || bufferId == static_cast<int>(mgr.getMainFileID()))
     return mainFileURI;
   llvm::Expected<lsp::URIForFile> fileForLoc = lsp::URIForFile::fromFile(
       mgr.getBufferInfo(bufferId).Buffer->getBufferIdentifier());
   if (fileForLoc)
     return *fileForLoc;
   lsp::Logger::error("Failed to create URI for include file: {0}",
                      llvm::toString(fileForLoc.takeError()));
   return mainFileURI;
 }

 /// Returns a language server location from the given source range.
 static lsp::Location getLocationFromLoc(llvm::SourceMgr &mgr, SMRange loc,
                                         const lsp::URIForFile &uri) {
   return lsp::Location(getURIFromLoc(mgr, loc.Start, uri),
                        lsp::Range(mgr, loc));
 }
 static lsp::Location getLocationFromLoc(llvm::SourceMgr &mgr, SMLoc loc,
                                         const lsp::URIForFile &uri) {
   return getLocationFromLoc(mgr, lsp::convertTokenLocToRange(loc), uri);
 }

 /// Convert the given TableGen diagnostic to the LSP form.
 static Optional<lsp::Diagnostic>
 getLspDiagnoticFromDiag(const llvm::SMDiagnostic &diag,
                         const lsp::URIForFile &uri) {
   auto *sourceMgr = const_cast<llvm::SourceMgr *>(diag.getSourceMgr());
   if (!sourceMgr || !diag.getLoc().isValid())
     return llvm::None;

   lsp::Diagnostic lspDiag;
   lspDiag.source = "tablegen";
   lspDiag.category = "Parse Error";

   // Try to grab a file location for this diagnostic.
   lsp::Location loc = getLocationFromLoc(*sourceMgr, diag.getLoc(), uri);
   lspDiag.range = loc.range;

   // Skip diagnostics that weren't emitted within the main file.
   if (loc.uri != uri)
     return llvm::None;

   // Convert the severity for the diagnostic.
   switch (diag.getKind()) {
   case llvm::SourceMgr::DK_Warning:
     lspDiag.severity = lsp::DiagnosticSeverity::Warning;
     break;
   case llvm::SourceMgr::DK_Error:
     lspDiag.severity = lsp::DiagnosticSeverity::Error;
     break;
   case llvm::SourceMgr::DK_Note:
     // Notes are emitted separately from the main diagnostic, so we just treat
     // them as remarks given that we can't determine the diagnostic to relate
     // them to.
   case llvm::SourceMgr::DK_Remark:
     lspDiag.severity = lsp::DiagnosticSeverity::Information;
     break;
   }
   lspDiag.message = diag.getMessage().str();

   return lspDiag;
 }

 //===----------------------------------------------------------------------===//
 // TableGenIndex
 //===----------------------------------------------------------------------===//

 namespace {
 /// This class represents a single symbol definition within a TableGen index. It
 /// contains the definition of the symbol, the location of the symbol, and any
 /// recorded references.
 struct TableGenIndexSymbol {
   TableGenIndexSymbol(const llvm::Record *record)
       : definition(record),
         defLoc(lsp::convertTokenLocToRange(record->getLoc().front())) {}
   TableGenIndexSymbol(const llvm::RecordVal *value)
       : definition(value),
         defLoc(lsp::convertTokenLocToRange(value->getLoc())) {}

   /// The main definition of the symbol.
   PointerUnion<const llvm::Record *, const llvm::RecordVal *> definition;

   /// The source location of the definition.
   SMRange defLoc;

   /// The source location of the references of the definition.
   SmallVector<SMRange> references;
 };

 /// This class provides an index for definitions/uses within a TableGen
 /// document. It provides efficient lookup of a definition given an input source
 /// range.
 class TableGenIndex {
 public:
   TableGenIndex() : intervalMap(allocator) {}

   /// Initialize the index with the given RecordKeeper.
   void initialize(const llvm::RecordKeeper &records);

   /// Lookup a symbol for the given location. Returns nullptr if no symbol could
   /// be found. If provided, `overlappedRange` is set to the range that the
   /// provided `loc` overlapped with.
   const TableGenIndexSymbol *lookup(SMLoc loc,
                                     SMRange *overlappedRange = nullptr) const;

 private:
   /// The type of interval map used to store source references. SMRange is
   /// half-open, so we also need to use a half-open interval map.
   using MapT = llvm::IntervalMap<
       const char *, const TableGenIndexSymbol *,
       llvm::IntervalMapImpl::NodeSizer<const char *,
                                        const TableGenIndexSymbol *>::LeafSize,
       llvm::IntervalMapHalfOpenInfo<const char *>>;

   /// An allocator for the interval map.
   MapT::Allocator allocator;

   /// An interval map containing a corresponding definition mapped to a source
   /// interval.
   MapT intervalMap;

   /// A mapping between definitions and their corresponding symbol.
   DenseMap<const void *, std::unique_ptr<TableGenIndexSymbol>> defToSymbol;
 };
 } // namespace

 void TableGenIndex::initialize(const llvm::RecordKeeper &records) {
   auto getOrInsertDef = [&](const auto *def) -> TableGenIndexSymbol * {
     auto it = defToSymbol.try_emplace(def, nullptr);
     if (it.second)
       it.first->second = std::make_unique<TableGenIndexSymbol>(def);
     return &*it.first->second;
   };
   auto insertRef = [&](TableGenIndexSymbol *sym, SMRange refLoc,
                        bool isDef = false) {
     const char *startLoc = refLoc.Start.getPointer();
     const char *endLoc = refLoc.End.getPointer();

     // If the location we got was empty, try to lex a token from the start
     // location.
     if (startLoc == endLoc) {
       refLoc = lsp::convertTokenLocToRange(SMLoc::getFromPointer(startLoc));
       startLoc = refLoc.Start.getPointer();
       endLoc = refLoc.End.getPointer();

       // If the location is still empty, bail on trying to use this reference
       // location.
       if (startLoc == endLoc)
         return;
     }

     // Check to see if a symbol is already attached to this location.
     // IntervalMap doesn't allow overlapping inserts, and we don't really
     // want multiple symbols attached to a source location anyways. This
     // shouldn't really happen in practice, but we should handle it gracefully.
     if (!intervalMap.overlaps(startLoc, endLoc))
       intervalMap.insert(startLoc, endLoc, sym);

     if (!isDef)
       sym->references.push_back(refLoc);
   };
   auto classes =
       llvm::make_pointee_range(llvm::make_second_range(records.getClasses()));
   auto defs =
       llvm::make_pointee_range(llvm::make_second_range(records.getDefs()));
   for (const llvm::Record &def : llvm::concat<llvm::Record>(classes, defs)) {
     auto *sym = getOrInsertDef(&def);
     insertRef(sym, sym->defLoc, /*isDef=*/true);

     // Add references to the definition.
     for (SMLoc loc : def.getLoc().drop_front())
       insertRef(sym, lsp::convertTokenLocToRange(loc));

     // Add references to any super classes.
     for (auto &it : def.getSuperClasses())
       insertRef(getOrInsertDef(it.first),
                 lsp::convertTokenLocToRange(it.second.Start));

     // Add definitions for any values.
     for (const llvm::RecordVal &value : def.getValues()) {
       auto *sym = getOrInsertDef(&value);
       insertRef(sym, sym->defLoc, /*isDef=*/true);
     }
   }
 }

 const TableGenIndexSymbol *
 TableGenIndex::lookup(SMLoc loc, SMRange *overlappedRange) const {
   auto it = intervalMap.find(loc.getPointer());
   if (!it.valid() || loc.getPointer() < it.start())
     return nullptr;

   if (overlappedRange) {
     *overlappedRange = SMRange(SMLoc::getFromPointer(it.start()),
                                SMLoc::getFromPointer(it.stop()));
   }
   return it.value();
 }

 //===----------------------------------------------------------------------===//
 // TableGenTextFile
 //===----------------------------------------------------------------------===//

 namespace {
 /// This class represents a text file containing one or more TableGen documents.
 class TableGenTextFile {
 public:
   TableGenTextFile(const lsp::URIForFile &uri, StringRef fileContents,
                    int64_t version,
                    const std::vector<std::string> &extraIncludeDirs,
                    std::vector<lsp::Diagnostic> &diagnostics);

   /// Return the current version of this text file.
   int64_t getVersion() const { return version; }

   /// Update the file to the new version using the provided set of content
   /// changes. Returns failure if the update was unsuccessful.
   LogicalResult update(const lsp::URIForFile &uri, int64_t newVersion,
                        ArrayRef<lsp::TextDocumentContentChangeEvent> changes,
                        std::vector<lsp::Diagnostic> &diagnostics);

   //===--------------------------------------------------------------------===//
   // Definitions and References
   //===--------------------------------------------------------------------===//

   void getLocationsOf(const lsp::URIForFile &uri, const lsp::Position &defPos,
                       std::vector<lsp::Location> &locations);
   void findReferencesOf(const lsp::URIForFile &uri, const lsp::Position &pos,
                         std::vector<lsp::Location> &references);

   //===--------------------------------------------------------------------===//
   // Document Links
   //===--------------------------------------------------------------------===//

   void getDocumentLinks(const lsp::URIForFile &uri,
                         std::vector<lsp::DocumentLink> &links);

   //===--------------------------------------------------------------------===//
   // Hover
   //===--------------------------------------------------------------------===//

   Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
                                  const lsp::Position &hoverPos);

 private:
   /// Initialize the text file from the given file contents.
   void initialize(const lsp::URIForFile &uri, int64_t newVersion,
                   std::vector<lsp::Diagnostic> &diagnostics);

   /// The full string contents of the file.
   std::string contents;

   /// The version of this file.
   int64_t version;

   /// The include directories for this file.
   std::vector<std::string> includeDirs;

   /// The source manager containing the contents of the input file.
   llvm::SourceMgr sourceMgr;

   /// The record keeper containing the parsed tablegen constructs.
   std::unique_ptr<llvm::RecordKeeper> recordKeeper;

   /// The index of the parsed file.
   TableGenIndex index;

   /// The set of includes of the parsed file.
   SmallVector<lsp::SourceMgrInclude> parsedIncludes;
 };
 } // namespace

 TableGenTextFile::TableGenTextFile(
     const lsp::URIForFile &uri, StringRef fileContents, int64_t version,
     const std::vector<std::string> &extraIncludeDirs,
     std::vector<lsp::Diagnostic> &diagnostics)
     : contents(fileContents.str()), version(version) {
   // Build the set of include directories for this file.
   llvm::SmallString<32> uriDirectory(uri.file());
   llvm::sys::path::remove_filename(uriDirectory);
   includeDirs.push_back(uriDirectory.str().str());
   includeDirs.insert(includeDirs.end(), extraIncludeDirs.begin(),
                      extraIncludeDirs.end());

   // Initialize the file.
   initialize(uri, version, diagnostics);
 }

 LogicalResult
 TableGenTextFile::update(const lsp::URIForFile &uri, int64_t newVersion,
                          ArrayRef<lsp::TextDocumentContentChangeEvent> changes,
                          std::vector<lsp::Diagnostic> &diagnostics) {
   if (failed(lsp::TextDocumentContentChangeEvent::applyTo(changes, contents))) {
     lsp::Logger::error("Failed to update contents of {0}", uri.file());
     return failure();
   }

   // If the file contents were properly changed, reinitialize the text file.
   initialize(uri, newVersion, diagnostics);
   return success();
 }

 void TableGenTextFile::initialize(const lsp::URIForFile &uri,
                                   int64_t newVersion,
                                   std::vector<lsp::Diagnostic> &diagnostics) {
   version = newVersion;
   sourceMgr = llvm::SourceMgr();
   recordKeeper = std::make_unique<llvm::RecordKeeper>();

   // Build a buffer for this file.
   auto memBuffer = llvm::MemoryBuffer::getMemBuffer(contents, uri.file());
   if (!memBuffer) {
     lsp::Logger::error("Failed to create memory buffer for file", uri.file());
     return;
   }
   sourceMgr.setIncludeDirs(includeDirs);
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());

   // This class provides a context argument for the llvm::SourceMgr diagnostic
   // handler.
   struct DiagHandlerContext {
     std::vector<lsp::Diagnostic> &diagnostics;
     const lsp::URIForFile &uri;
   } handlerContext{diagnostics, uri};

   // Set the diagnostic handler for the tablegen source manager.
   sourceMgr.setDiagHandler(
       [](const llvm::SMDiagnostic &diag, void *rawHandlerContext) {
         auto *ctx = reinterpret_cast<DiagHandlerContext *>(rawHandlerContext);
         if (auto lspDiag = getLspDiagnoticFromDiag(diag, ctx->uri))
           ctx->diagnostics.push_back(*lspDiag);
       },
       &handlerContext);
   bool failedToParse = llvm::TableGenParseFile(sourceMgr, *recordKeeper);

   // Process all of the include files.
   lsp::gatherIncludeFiles(sourceMgr, parsedIncludes);
   if (failedToParse)
     return;

   // If we successfully parsed the file, we can now build the index.
   index.initialize(*recordKeeper);
 }

 //===----------------------------------------------------------------------===//
 // TableGenTextFile: Definitions and References
 //===----------------------------------------------------------------------===//

 void TableGenTextFile::getLocationsOf(const lsp::URIForFile &uri,
                                       const lsp::Position &defPos,
                                       std::vector<lsp::Location> &locations) {
   SMLoc posLoc = defPos.getAsSMLoc(sourceMgr);
   const TableGenIndexSymbol *symbol = index.lookup(posLoc);
   if (!symbol)
     return;

   locations.push_back(getLocationFromLoc(sourceMgr, symbol->defLoc, uri));
 }

 void TableGenTextFile::findReferencesOf(
     const lsp::URIForFile &uri, const lsp::Position &pos,
     std::vector<lsp::Location> &references) {
   SMLoc posLoc = pos.getAsSMLoc(sourceMgr);
   const TableGenIndexSymbol *symbol = index.lookup(posLoc);
   if (!symbol)
     return;

   references.push_back(getLocationFromLoc(sourceMgr, symbol->defLoc, uri));
   for (SMRange refLoc : symbol->references)
     references.push_back(getLocationFromLoc(sourceMgr, refLoc, uri));
 }

 //===--------------------------------------------------------------------===//
 // TableGenTextFile: Document Links
 //===--------------------------------------------------------------------===//

 void TableGenTextFile::getDocumentLinks(const lsp::URIForFile &uri,
                                         std::vector<lsp::DocumentLink> &links) {
   for (const lsp::SourceMgrInclude &include : parsedIncludes)
     links.emplace_back(include.range, include.uri);
 }

 //===----------------------------------------------------------------------===//
 // TableGenTextFile: Hover
 //===----------------------------------------------------------------------===//

 Optional<lsp::Hover>
 TableGenTextFile::findHover(const lsp::URIForFile &uri,
                             const lsp::Position &hoverPos) {
   // Check for a reference to an include.
   for (const lsp::SourceMgrInclude &include : parsedIncludes)
     if (include.range.contains(hoverPos))
       return include.buildHover();
   return llvm::None;
 }

 //===----------------------------------------------------------------------===//
 // TableGenServer::Impl
 //===----------------------------------------------------------------------===//

 struct lsp::TableGenServer::Impl {
   explicit Impl(const Options &options)
       : options(options), compilationDatabase(options.compilationDatabases) {}

   /// TableGen LSP options.
   const Options &options;

   /// The compilation database containing additional information for files
   /// passed to the server.
   lsp::CompilationDatabase compilationDatabase;

   /// The files held by the server, mapped by their URI file name.
   llvm::StringMap<std::unique_ptr<TableGenTextFile>> files;
 };

 //===----------------------------------------------------------------------===//
 // TableGenServer
 //===----------------------------------------------------------------------===//

 lsp::TableGenServer::TableGenServer(const Options &options)
     : impl(std::make_unique<Impl>(options)) {}
 lsp::TableGenServer::~TableGenServer() = default;

 void lsp::TableGenServer::addDocument(const URIForFile &uri, StringRef contents,
                                       int64_t version,
                                       std::vector<Diagnostic> &diagnostics) {
   // Build the set of additional include directories.
   std::vector<std::string> additionalIncludeDirs = impl->options.extraDirs;
   const auto &fileInfo = impl->compilationDatabase.getFileInfo(uri.file());
   llvm::append_range(additionalIncludeDirs, fileInfo.includeDirs);

   impl->files[uri.file()] = std::make_unique<TableGenTextFile>(
       uri, contents, version, additionalIncludeDirs, diagnostics);
 }

 void lsp::TableGenServer::updateDocument(
     const URIForFile &uri, ArrayRef<TextDocumentContentChangeEvent> changes,
     int64_t version, std::vector<Diagnostic> &diagnostics) {
   // Check that we actually have a document for this uri.
   auto it = impl->files.find(uri.file());
   if (it == impl->files.end())
     return;

   // Try to update the document. If we fail, erase the file from the server. A
   // failed updated generally means we've fallen out of sync somewhere.
   if (failed(it->second->update(uri, version, changes, diagnostics)))
     impl->files.erase(it);
 }

 Optional<int64_t> lsp::TableGenServer::removeDocument(const URIForFile &uri) {
   auto it = impl->files.find(uri.file());
   if (it == impl->files.end())
     return llvm::None;

   int64_t version = it->second->getVersion();
   impl->files.erase(it);
   return version;
 }

 void lsp::TableGenServer::getLocationsOf(const URIForFile &uri,
                                          const Position &defPos,
                                          std::vector<Location> &locations) {
   auto fileIt = impl->files.find(uri.file());
   if (fileIt != impl->files.end())
     fileIt->second->getLocationsOf(uri, defPos, locations);
 }

 void lsp::TableGenServer::findReferencesOf(const URIForFile &uri,
                                            const Position &pos,
                                            std::vector<Location> &references) {
   auto fileIt = impl->files.find(uri.file());
   if (fileIt != impl->files.end())
     fileIt->second->findReferencesOf(uri, pos, references);
 }

 void lsp::TableGenServer::getDocumentLinks(
     const URIForFile &uri, std::vector<DocumentLink> &documentLinks) {
   auto fileIt = impl->files.find(uri.file());
   if (fileIt != impl->files.end())
     return fileIt->second->getDocumentLinks(uri, documentLinks);
 }

 Optional<lsp::Hover> lsp::TableGenServer::findHover(const URIForFile &uri,
                                                     const Position &hoverPos) {
   auto fileIt = impl->files.find(uri.file());
   if (fileIt != impl->files.end())
     return fileIt->second->findHover(uri, hoverPos);
   return llvm::None;
 }
	//===- TableGenServer.cpp - TableGen Language Server ----------------------===//
	//
	// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	//
	//===----------------------------------------------------------------------===//

	#include "TableGenServer.h"

	#include "../lsp-server-support/CompilationDatabase.h"
	#include "../lsp-server-support/Logging.h"
	#include "../lsp-server-support/Protocol.h"
	#include "../lsp-server-support/SourceMgrUtils.h"
	#include "mlir/Support/LogicalResult.h"
	#include "llvm/ADT/IntervalMap.h"
	#include "llvm/ADT/PointerUnion.h"
	#include "llvm/ADT/StringMap.h"
	#include "llvm/ADT/StringSet.h"
	#include "llvm/ADT/TypeSwitch.h"
	#include "llvm/Support/FileSystem.h"
	#include "llvm/Support/Path.h"
	#include "llvm/TableGen/Parser.h"
	#include "llvm/TableGen/Record.h"

	using namespace mlir;

	/// Returns a language server uri for the given source location. `mainFileURI`
	/// corresponds to the uri for the main file of the source manager.
	static lsp::URIForFile getURIFromLoc(const llvm::SourceMgr &mgr, SMLoc loc,
	const lsp::URIForFile &mainFileURI) {
	int bufferId = mgr.FindBufferContainingLoc(loc);
	if (bufferId == 0 \|\| bufferId == static_cast<int>(mgr.getMainFileID()))
	return mainFileURI;
	llvm::Expected<lsp::URIForFile> fileForLoc = lsp::URIForFile::fromFile(
	mgr.getBufferInfo(bufferId).Buffer->getBufferIdentifier());
	if (fileForLoc)
	return *fileForLoc;
	lsp::Logger::error("Failed to create URI for include file: {0}",
	llvm::toString(fileForLoc.takeError()));
	return mainFileURI;
	}

	/// Returns a language server location from the given source range.
	static lsp::Location getLocationFromLoc(llvm::SourceMgr &mgr, SMRange loc,
	const lsp::URIForFile &uri) {
	return lsp::Location(getURIFromLoc(mgr, loc.Start, uri),
	lsp::Range(mgr, loc));
	}
	static lsp::Location getLocationFromLoc(llvm::SourceMgr &mgr, SMLoc loc,
	const lsp::URIForFile &uri) {
	return getLocationFromLoc(mgr, lsp::convertTokenLocToRange(loc), uri);
	}

	/// Convert the given TableGen diagnostic to the LSP form.
	static Optional<lsp::Diagnostic>
	getLspDiagnoticFromDiag(const llvm::SMDiagnostic &diag,
	const lsp::URIForFile &uri) {
	auto sourceMgr = const_cast<llvm::SourceMgr >(diag.getSourceMgr());
	if (!sourceMgr \|\| !diag.getLoc().isValid())
	return llvm::None;

	lsp::Diagnostic lspDiag;
	lspDiag.source = "tablegen";
	lspDiag.category = "Parse Error";

	// Try to grab a file location for this diagnostic.
	lsp::Location loc = getLocationFromLoc(*sourceMgr, diag.getLoc(), uri);
	lspDiag.range = loc.range;

	// Skip diagnostics that weren't emitted within the main file.
	if (loc.uri != uri)
	return llvm::None;

	// Convert the severity for the diagnostic.
	switch (diag.getKind()) {
	case llvm::SourceMgr::DK_Warning:
	lspDiag.severity = lsp::DiagnosticSeverity::Warning;
	break;
	case llvm::SourceMgr::DK_Error:
	lspDiag.severity = lsp::DiagnosticSeverity::Error;
	break;
	case llvm::SourceMgr::DK_Note:
	// Notes are emitted separately from the main diagnostic, so we just treat
	// them as remarks given that we can't determine the diagnostic to relate
	// them to.
	case llvm::SourceMgr::DK_Remark:
	lspDiag.severity = lsp::DiagnosticSeverity::Information;
	break;
	}
	lspDiag.message = diag.getMessage().str();

	return lspDiag;
	}

	//===----------------------------------------------------------------------===//
	// TableGenIndex
	//===----------------------------------------------------------------------===//

	namespace {
	/// This class represents a single symbol definition within a TableGen index. It
	/// contains the definition of the symbol, the location of the symbol, and any
	/// recorded references.
	struct TableGenIndexSymbol {
	TableGenIndexSymbol(const llvm::Record *record)
	: definition(record),
	defLoc(lsp::convertTokenLocToRange(record->getLoc().front())) {}
	TableGenIndexSymbol(const llvm::RecordVal *value)
	: definition(value),
	defLoc(lsp::convertTokenLocToRange(value->getLoc())) {}

	/// The main definition of the symbol.
	PointerUnion<const llvm::Record , const llvm::RecordVal > definition;

	/// The source location of the definition.
	SMRange defLoc;

	/// The source location of the references of the definition.
	SmallVector<SMRange> references;
	};

	/// This class provides an index for definitions/uses within a TableGen
	/// document. It provides efficient lookup of a definition given an input source
	/// range.
	class TableGenIndex {
	public:
	TableGenIndex() : intervalMap(allocator) {}

	/// Initialize the index with the given RecordKeeper.
	void initialize(const llvm::RecordKeeper &records);

	/// Lookup a symbol for the given location. Returns nullptr if no symbol could
	/// be found. If provided, `overlappedRange` is set to the range that the
	/// provided `loc` overlapped with.
	const TableGenIndexSymbol *lookup(SMLoc loc,
	SMRange *overlappedRange = nullptr) const;

	private:
	/// The type of interval map used to store source references. SMRange is
	/// half-open, so we also need to use a half-open interval map.
	using MapT = llvm::IntervalMap<
	const char , const TableGenIndexSymbol ,
	llvm::IntervalMapImpl::NodeSizer<const char *,
	const TableGenIndexSymbol *>::LeafSize,
	llvm::IntervalMapHalfOpenInfo<const char *>>;

	/// An allocator for the interval map.
	MapT::Allocator allocator;

	/// An interval map containing a corresponding definition mapped to a source
	/// interval.
	MapT intervalMap;

	/// A mapping between definitions and their corresponding symbol.
	DenseMap<const void *, std::unique_ptr<TableGenIndexSymbol>> defToSymbol;
	};
	} // namespace

	void TableGenIndex::initialize(const llvm::RecordKeeper &records) {
	auto getOrInsertDef = [&](const auto def) -> TableGenIndexSymbol {
	auto it = defToSymbol.try_emplace(def, nullptr);
	if (it.second)
	it.first->second = std::make_unique<TableGenIndexSymbol>(def);
	return &*it.first->second;
	};
	auto insertRef = [&](TableGenIndexSymbol *sym, SMRange refLoc,
	bool isDef = false) {
	const char *startLoc = refLoc.Start.getPointer();
	const char *endLoc = refLoc.End.getPointer();

	// If the location we got was empty, try to lex a token from the start
	// location.
	if (startLoc == endLoc) {
	refLoc = lsp::convertTokenLocToRange(SMLoc::getFromPointer(startLoc));
	startLoc = refLoc.Start.getPointer();
	endLoc = refLoc.End.getPointer();

	// If the location is still empty, bail on trying to use this reference
	// location.
	if (startLoc == endLoc)
	return;
	}

	// Check to see if a symbol is already attached to this location.
	// IntervalMap doesn't allow overlapping inserts, and we don't really
	// want multiple symbols attached to a source location anyways. This
	// shouldn't really happen in practice, but we should handle it gracefully.
	if (!intervalMap.overlaps(startLoc, endLoc))
	intervalMap.insert(startLoc, endLoc, sym);

	if (!isDef)
	sym->references.push_back(refLoc);
	};
	auto classes =
	llvm::make_pointee_range(llvm::make_second_range(records.getClasses()));
	auto defs =
	llvm::make_pointee_range(llvm::make_second_range(records.getDefs()));
	for (const llvm::Record &def : llvm::concat<llvm::Record>(classes, defs)) {
	auto *sym = getOrInsertDef(&def);
	insertRef(sym, sym->defLoc, /isDef=/true);

	// Add references to the definition.
	for (SMLoc loc : def.getLoc().drop_front())
	insertRef(sym, lsp::convertTokenLocToRange(loc));

	// Add references to any super classes.
	for (auto &it : def.getSuperClasses())
	insertRef(getOrInsertDef(it.first),
	lsp::convertTokenLocToRange(it.second.Start));

	// Add definitions for any values.
	for (const llvm::RecordVal &value : def.getValues()) {
	auto *sym = getOrInsertDef(&value);
	insertRef(sym, sym->defLoc, /isDef=/true);
	}
	}
	}

	const TableGenIndexSymbol *
	TableGenIndex::lookup(SMLoc loc, SMRange *overlappedRange) const {
	auto it = intervalMap.find(loc.getPointer());
	if (!it.valid() \|\| loc.getPointer() < it.start())
	return nullptr;

	if (overlappedRange) {
	*overlappedRange = SMRange(SMLoc::getFromPointer(it.start()),
	SMLoc::getFromPointer(it.stop()));
	}
	return it.value();
	}

	//===----------------------------------------------------------------------===//
	// TableGenTextFile
	//===----------------------------------------------------------------------===//

	namespace {
	/// This class represents a text file containing one or more TableGen documents.
	class TableGenTextFile {
	public:
	TableGenTextFile(const lsp::URIForFile &uri, StringRef fileContents,
	int64_t version,
	const std::vector<std::string> &extraIncludeDirs,
	std::vector<lsp::Diagnostic> &diagnostics);

	/// Return the current version of this text file.
	int64_t getVersion() const { return version; }

	/// Update the file to the new version using the provided set of content
	/// changes. Returns failure if the update was unsuccessful.
	LogicalResult update(const lsp::URIForFile &uri, int64_t newVersion,
	ArrayRef<lsp::TextDocumentContentChangeEvent> changes,
	std::vector<lsp::Diagnostic> &diagnostics);

	//===--------------------------------------------------------------------===//
	// Definitions and References
	//===--------------------------------------------------------------------===//

	void getLocationsOf(const lsp::URIForFile &uri, const lsp::Position &defPos,
	std::vector<lsp::Location> &locations);
	void findReferencesOf(const lsp::URIForFile &uri, const lsp::Position &pos,
	std::vector<lsp::Location> &references);

	//===--------------------------------------------------------------------===//
	// Document Links
	//===--------------------------------------------------------------------===//

	void getDocumentLinks(const lsp::URIForFile &uri,
	std::vector<lsp::DocumentLink> &links);

	//===--------------------------------------------------------------------===//
	// Hover
	//===--------------------------------------------------------------------===//

	Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
	const lsp::Position &hoverPos);

	private:
	/// Initialize the text file from the given file contents.
	void initialize(const lsp::URIForFile &uri, int64_t newVersion,
	std::vector<lsp::Diagnostic> &diagnostics);

	/// The full string contents of the file.
	std::string contents;

	/// The version of this file.
	int64_t version;

	/// The include directories for this file.
	std::vector<std::string> includeDirs;

	/// The source manager containing the contents of the input file.
	llvm::SourceMgr sourceMgr;

	/// The record keeper containing the parsed tablegen constructs.
	std::unique_ptr<llvm::RecordKeeper> recordKeeper;

	/// The index of the parsed file.
	TableGenIndex index;

	/// The set of includes of the parsed file.
	SmallVector<lsp::SourceMgrInclude> parsedIncludes;
	};
	} // namespace

	TableGenTextFile::TableGenTextFile(
	const lsp::URIForFile &uri, StringRef fileContents, int64_t version,
	const std::vector<std::string> &extraIncludeDirs,
	std::vector<lsp::Diagnostic> &diagnostics)
	: contents(fileContents.str()), version(version) {
	// Build the set of include directories for this file.
	llvm::SmallString<32> uriDirectory(uri.file());
	llvm::sys::path::remove_filename(uriDirectory);
	includeDirs.push_back(uriDirectory.str().str());
	includeDirs.insert(includeDirs.end(), extraIncludeDirs.begin(),
	extraIncludeDirs.end());

	// Initialize the file.
	initialize(uri, version, diagnostics);
	}

	LogicalResult
	TableGenTextFile::update(const lsp::URIForFile &uri, int64_t newVersion,
	ArrayRef<lsp::TextDocumentContentChangeEvent> changes,
	std::vector<lsp::Diagnostic> &diagnostics) {
	if (failed(lsp::TextDocumentContentChangeEvent::applyTo(changes, contents))) {
	lsp::Logger::error("Failed to update contents of {0}", uri.file());
	return failure();
	}

	// If the file contents were properly changed, reinitialize the text file.
	initialize(uri, newVersion, diagnostics);
	return success();
	}

	void TableGenTextFile::initialize(const lsp::URIForFile &uri,
	int64_t newVersion,
	std::vector<lsp::Diagnostic> &diagnostics) {
	version = newVersion;
	sourceMgr = llvm::SourceMgr();
	recordKeeper = std::make_unique<llvm::RecordKeeper>();

	// Build a buffer for this file.
	auto memBuffer = llvm::MemoryBuffer::getMemBuffer(contents, uri.file());
	if (!memBuffer) {
	lsp::Logger::error("Failed to create memory buffer for file", uri.file());
	return;
	}
	sourceMgr.setIncludeDirs(includeDirs);
	sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());

	// This class provides a context argument for the llvm::SourceMgr diagnostic
	// handler.
	struct DiagHandlerContext {
	std::vector<lsp::Diagnostic> &diagnostics;
	const lsp::URIForFile &uri;
	} handlerContext{diagnostics, uri};

	// Set the diagnostic handler for the tablegen source manager.
	sourceMgr.setDiagHandler(
	[](const llvm::SMDiagnostic &diag, void *rawHandlerContext) {
	auto ctx = reinterpret_cast<DiagHandlerContext >(rawHandlerContext);
	if (auto lspDiag = getLspDiagnoticFromDiag(diag, ctx->uri))
	ctx->diagnostics.push_back(*lspDiag);
	},
	&handlerContext);
	bool failedToParse = llvm::TableGenParseFile(sourceMgr, *recordKeeper);

	// Process all of the include files.
	lsp::gatherIncludeFiles(sourceMgr, parsedIncludes);
	if (failedToParse)
	return;

	// If we successfully parsed the file, we can now build the index.
	index.initialize(*recordKeeper);
	}

	//===----------------------------------------------------------------------===//
	// TableGenTextFile: Definitions and References
	//===----------------------------------------------------------------------===//

	void TableGenTextFile::getLocationsOf(const lsp::URIForFile &uri,
	const lsp::Position &defPos,
	std::vector<lsp::Location> &locations) {
	SMLoc posLoc = defPos.getAsSMLoc(sourceMgr);
	const TableGenIndexSymbol *symbol = index.lookup(posLoc);
	if (!symbol)
	return;

	locations.push_back(getLocationFromLoc(sourceMgr, symbol->defLoc, uri));
	}

	void TableGenTextFile::findReferencesOf(
	const lsp::URIForFile &uri, const lsp::Position &pos,
	std::vector<lsp::Location> &references) {
	SMLoc posLoc = pos.getAsSMLoc(sourceMgr);
	const TableGenIndexSymbol *symbol = index.lookup(posLoc);
	if (!symbol)
	return;

	references.push_back(getLocationFromLoc(sourceMgr, symbol->defLoc, uri));
	for (SMRange refLoc : symbol->references)
	references.push_back(getLocationFromLoc(sourceMgr, refLoc, uri));
	}

	//===--------------------------------------------------------------------===//
	// TableGenTextFile: Document Links
	//===--------------------------------------------------------------------===//

	void TableGenTextFile::getDocumentLinks(const lsp::URIForFile &uri,
	std::vector<lsp::DocumentLink> &links) {
	for (const lsp::SourceMgrInclude &include : parsedIncludes)
	links.emplace_back(include.range, include.uri);
	}

	//===----------------------------------------------------------------------===//
	// TableGenTextFile: Hover
	//===----------------------------------------------------------------------===//

	Optional<lsp::Hover>
	TableGenTextFile::findHover(const lsp::URIForFile &uri,
	const lsp::Position &hoverPos) {
	// Check for a reference to an include.
	for (const lsp::SourceMgrInclude &include : parsedIncludes)
	if (include.range.contains(hoverPos))
	return include.buildHover();
	return llvm::None;
	}

	//===----------------------------------------------------------------------===//
	// TableGenServer::Impl
	//===----------------------------------------------------------------------===//

	struct lsp::TableGenServer::Impl {
	explicit Impl(const Options &options)
	: options(options), compilationDatabase(options.compilationDatabases) {}

	/// TableGen LSP options.
	const Options &options;

	/// The compilation database containing additional information for files
	/// passed to the server.
	lsp::CompilationDatabase compilationDatabase;

	/// The files held by the server, mapped by their URI file name.
	llvm::StringMap<std::unique_ptr<TableGenTextFile>> files;
	};

	//===----------------------------------------------------------------------===//
	// TableGenServer
	//===----------------------------------------------------------------------===//

	lsp::TableGenServer::TableGenServer(const Options &options)
	: impl(std::make_unique<Impl>(options)) {}
	lsp::TableGenServer::~TableGenServer() = default;

	void lsp::TableGenServer::addDocument(const URIForFile &uri, StringRef contents,
	int64_t version,
	std::vector<Diagnostic> &diagnostics) {
	// Build the set of additional include directories.
	std::vector<std::string> additionalIncludeDirs = impl->options.extraDirs;
	const auto &fileInfo = impl->compilationDatabase.getFileInfo(uri.file());
	llvm::append_range(additionalIncludeDirs, fileInfo.includeDirs);

	impl->files[uri.file()] = std::make_unique<TableGenTextFile>(
	uri, contents, version, additionalIncludeDirs, diagnostics);
	}

	void lsp::TableGenServer::updateDocument(
	const URIForFile &uri, ArrayRef<TextDocumentContentChangeEvent> changes,
	int64_t version, std::vector<Diagnostic> &diagnostics) {
	// Check that we actually have a document for this uri.
	auto it = impl->files.find(uri.file());
	if (it == impl->files.end())
	return;

	// Try to update the document. If we fail, erase the file from the server. A
	// failed updated generally means we've fallen out of sync somewhere.
	if (failed(it->second->update(uri, version, changes, diagnostics)))
	impl->files.erase(it);
	}

	Optional<int64_t> lsp::TableGenServer::removeDocument(const URIForFile &uri) {
	auto it = impl->files.find(uri.file());
	if (it == impl->files.end())
	return llvm::None;

	int64_t version = it->second->getVersion();
	impl->files.erase(it);
	return version;
	}

	void lsp::TableGenServer::getLocationsOf(const URIForFile &uri,
	const Position &defPos,
	std::vector<Location> &locations) {
	auto fileIt = impl->files.find(uri.file());
	if (fileIt != impl->files.end())
	fileIt->second->getLocationsOf(uri, defPos, locations);
	}

	void lsp::TableGenServer::findReferencesOf(const URIForFile &uri,
	const Position &pos,
	std::vector<Location> &references) {
	auto fileIt = impl->files.find(uri.file());
	if (fileIt != impl->files.end())
	fileIt->second->findReferencesOf(uri, pos, references);
	}

	void lsp::TableGenServer::getDocumentLinks(
	const URIForFile &uri, std::vector<DocumentLink> &documentLinks) {
	auto fileIt = impl->files.find(uri.file());
	if (fileIt != impl->files.end())
	return fileIt->second->getDocumentLinks(uri, documentLinks);
	}

	Optional<lsp::Hover> lsp::TableGenServer::findHover(const URIForFile &uri,
	const Position &hoverPos) {
	auto fileIt = impl->files.find(uri.file());
	if (fileIt != impl->files.end())
	return fileIt->second->findHover(uri, hoverPos);
	return llvm::None;
	}