MLIRServer.cpp | Explore in Territory

//===- MLIRServer.cpp - MLIR Generic 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 "MLIRServer.h"
#include "Protocol.h"
#include "mlir/AsmParser/AsmParser.h"
#include "mlir/AsmParser/AsmParserState.h"
#include "mlir/AsmParser/CodeComplete.h"
#include "mlir/Bytecode/BytecodeWriter.h"
#include "mlir/IR/Operation.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Support/ToolUtilities.h"
#include "mlir/Tools/lsp-server-support/Logging.h"
#include "mlir/Tools/lsp-server-support/SourceMgrUtils.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/Support/Base64.h"
#include "llvm/Support/SourceMgr.h"
#include <optional>

usingnamespacemlir;

/// Returns the range of a lexical token given a SMLoc corresponding to the
/// start of an token location. The range is computed heuristically, and
/// supports identifier-like tokens, strings, etc.
static SMRange convertTokenLocToRange(SMLoc loc) { … }

/// Returns a language server location from the given MLIR file location.
/// `uriScheme` is the scheme to use when building new uris.
static std::optional<lsp::Location> getLocationFromLoc(StringRef uriScheme,
                                                       FileLineColLoc loc) { … }

/// Returns a language server location from the given MLIR location, or
/// std::nullopt if one couldn't be created. `uriScheme` is the scheme to use
/// when building new uris. `uri` is an optional additional filter that, when
/// present, is used to filter sub locations that do not share the same uri.
static std::optional<lsp::Location>
getLocationFromLoc(llvm::SourceMgr &sourceMgr, Location loc,
                   StringRef uriScheme, const lsp::URIForFile *uri = nullptr) { … }

/// Collect all of the locations from the given MLIR location that are not
/// contained within the given URI.
static void collectLocationsFromLoc(Location loc,
                                    std::vector<lsp::Location> &locations,
                                    const lsp::URIForFile &uri) { … }

/// Returns true if the given range contains the given source location. Note
/// that this has slightly different behavior than SMRange because it is
/// inclusive of the end location.
static bool contains(SMRange range, SMLoc loc) { … }

/// Returns true if the given location is contained by the definition or one of
/// the uses of the given SMDefinition. If provided, `overlappedRange` is set to
/// the range within `def` that the provided `loc` overlapped with.
static bool isDefOrUse(const AsmParserState::SMDefinition &def, SMLoc loc,
                       SMRange *overlappedRange = nullptr) { … }

/// Given a location pointing to a result, return the result number it refers
/// to or std::nullopt if it refers to all of the results.
static std::optional<unsigned> getResultNumberFromLoc(SMLoc loc) { … }

/// Given a source location range, return the text covered by the given range.
/// If the range is invalid, returns std::nullopt.
static std::optional<StringRef> getTextFromRange(SMRange range) { … }

/// Given a block, return its position in its parent region.
static unsigned getBlockNumber(Block *block) { … }

/// Given a block and source location, print the source name of the block to the
/// given output stream.
static void printDefBlockName(raw_ostream &os, Block *block, SMRange loc = { … }
static void printDefBlockName(raw_ostream &os,
                              const AsmParserState::BlockDefinition &def) { … }

/// Convert the given MLIR diagnostic to the LSP form.
static lsp::Diagnostic getLspDiagnoticFromDiag(llvm::SourceMgr &sourceMgr,
                                               Diagnostic &diag,
                                               const lsp::URIForFile &uri) { … }

//===----------------------------------------------------------------------===//
// MLIRDocument
//===----------------------------------------------------------------------===//

namespace {
/// This class represents all of the information pertaining to a specific MLIR
/// document.
struct MLIRDocument { … };
} // namespace

MLIRDocument::MLIRDocument(MLIRContext &context, const lsp::URIForFile &uri,
                           StringRef contents,
                           std::vector<lsp::Diagnostic> &diagnostics) { … }

//===----------------------------------------------------------------------===//
// MLIRDocument: Definitions and References
//===----------------------------------------------------------------------===//

void MLIRDocument::getLocationsOf(const lsp::URIForFile &uri,
                                  const lsp::Position &defPos,
                                  std::vector<lsp::Location> &locations) { … }

void MLIRDocument::findReferencesOf(const lsp::URIForFile &uri,
                                    const lsp::Position &pos,
                                    std::vector<lsp::Location> &references) { … }

//===----------------------------------------------------------------------===//
// MLIRDocument: Hover
//===----------------------------------------------------------------------===//

std::optional<lsp::Hover>
MLIRDocument::findHover(const lsp::URIForFile &uri,
                        const lsp::Position &hoverPos) { … }

std::optional<lsp::Hover> MLIRDocument::buildHoverForOperation(
    SMRange hoverRange, const AsmParserState::OperationDefinition &op) { … }

lsp::Hover MLIRDocument::buildHoverForOperationResult(SMRange hoverRange,
                                                      Operation *op,
                                                      unsigned resultStart,
                                                      unsigned resultEnd,
                                                      SMLoc posLoc) { … }

lsp::Hover
MLIRDocument::buildHoverForBlock(SMRange hoverRange,
                                 const AsmParserState::BlockDefinition &block) { … }

lsp::Hover MLIRDocument::buildHoverForBlockArgument(
    SMRange hoverRange, BlockArgument arg,
    const AsmParserState::BlockDefinition &block) { … }

lsp::Hover MLIRDocument::buildHoverForAttributeAlias(
    SMRange hoverRange, const AsmParserState::AttributeAliasDefinition &attr) { … }

lsp::Hover MLIRDocument::buildHoverForTypeAlias(
    SMRange hoverRange, const AsmParserState::TypeAliasDefinition &type) { … }

//===----------------------------------------------------------------------===//
// MLIRDocument: Document Symbols
//===----------------------------------------------------------------------===//

void MLIRDocument::findDocumentSymbols(
    std::vector<lsp::DocumentSymbol> &symbols) { … }

void MLIRDocument::findDocumentSymbols(
    Operation *op, std::vector<lsp::DocumentSymbol> &symbols) { … }

//===----------------------------------------------------------------------===//
// MLIRDocument: Code Completion
//===----------------------------------------------------------------------===//

namespace {
class LSPCodeCompleteContext : public AsmParserCodeCompleteContext { … };
} // namespace

lsp::CompletionList
MLIRDocument::getCodeCompletion(const lsp::URIForFile &uri,
                                const lsp::Position &completePos,
                                const DialectRegistry &registry) { … }

//===----------------------------------------------------------------------===//
// MLIRDocument: Code Action
//===----------------------------------------------------------------------===//

void MLIRDocument::getCodeActionForDiagnostic(
    const lsp::URIForFile &uri, lsp::Position &pos, StringRef severity,
    StringRef message, std::vector<lsp::TextEdit> &edits) { … }

//===----------------------------------------------------------------------===//
// MLIRDocument: Bytecode
//===----------------------------------------------------------------------===//

llvm::Expected<lsp::MLIRConvertBytecodeResult>
MLIRDocument::convertToBytecode() { … }

//===----------------------------------------------------------------------===//
// MLIRTextFileChunk
//===----------------------------------------------------------------------===//

namespace {
/// This class represents a single chunk of an MLIR text file.
struct MLIRTextFileChunk { … };
} // namespace

//===----------------------------------------------------------------------===//
// MLIRTextFile
//===----------------------------------------------------------------------===//

namespace {
/// This class represents a text file containing one or more MLIR documents.
class MLIRTextFile { … };
} // namespace

MLIRTextFile::MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents,
                           int64_t version, DialectRegistry &registry,
                           std::vector<lsp::Diagnostic> &diagnostics)
    : … { … }

void MLIRTextFile::getLocationsOf(const lsp::URIForFile &uri,
                                  lsp::Position defPos,
                                  std::vector<lsp::Location> &locations) { … }

void MLIRTextFile::findReferencesOf(const lsp::URIForFile &uri,
                                    lsp::Position pos,
                                    std::vector<lsp::Location> &references) { … }

std::optional<lsp::Hover> MLIRTextFile::findHover(const lsp::URIForFile &uri,
                                                  lsp::Position hoverPos) { … }

void MLIRTextFile::findDocumentSymbols(
    std::vector<lsp::DocumentSymbol> &symbols) { … }

lsp::CompletionList MLIRTextFile::getCodeCompletion(const lsp::URIForFile &uri,
                                                    lsp::Position completePos) { … }

void MLIRTextFile::getCodeActions(const lsp::URIForFile &uri,
                                  const lsp::Range &pos,
                                  const lsp::CodeActionContext &context,
                                  std::vector<lsp::CodeAction> &actions) { … }

llvm::Expected<lsp::MLIRConvertBytecodeResult>
MLIRTextFile::convertToBytecode() { … }

MLIRTextFileChunk &MLIRTextFile::getChunkFor(lsp::Position &pos) { … }

//===----------------------------------------------------------------------===//
// MLIRServer::Impl
//===----------------------------------------------------------------------===//

struct lsp::MLIRServer::Impl { … };

//===----------------------------------------------------------------------===//
// MLIRServer
//===----------------------------------------------------------------------===//

lsp::MLIRServer::MLIRServer(DialectRegistry &registry)
    : … { … }
lsp::MLIRServer::~MLIRServer() = default;

void lsp::MLIRServer::addOrUpdateDocument(
    const URIForFile &uri, StringRef contents, int64_t version,
    std::vector<Diagnostic> &diagnostics) { … }

std::optional<int64_t> lsp::MLIRServer::removeDocument(const URIForFile &uri) { … }

void lsp::MLIRServer::getLocationsOf(const URIForFile &uri,
                                     const Position &defPos,
                                     std::vector<Location> &locations) { … }

void lsp::MLIRServer::findReferencesOf(const URIForFile &uri,
                                       const Position &pos,
                                       std::vector<Location> &references) { … }

std::optional<lsp::Hover> lsp::MLIRServer::findHover(const URIForFile &uri,
                                                     const Position &hoverPos) { … }

void lsp::MLIRServer::findDocumentSymbols(
    const URIForFile &uri, std::vector<DocumentSymbol> &symbols) { … }

lsp::CompletionList
lsp::MLIRServer::getCodeCompletion(const URIForFile &uri,
                                   const Position &completePos) { … }

void lsp::MLIRServer::getCodeActions(const URIForFile &uri, const Range &pos,
                                     const CodeActionContext &context,
                                     std::vector<CodeAction> &actions) { … }

llvm::Expected<lsp::MLIRConvertBytecodeResult>
lsp::MLIRServer::convertFromBytecode(const URIForFile &uri) { … }

llvm::Expected<lsp::MLIRConvertBytecodeResult>
lsp::MLIRServer::convertToBytecode(const URIForFile &uri) { … }
llvm/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
