//===- AsmState.h - Assembly State Utilities --------------------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file defines various classes and utilites for interacting with the MLIR
// assembly formats.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_IR_ASMSTATE_H_
#define MLIR_IR_ASMSTATE_H_

#include "mlir/Bytecode/BytecodeReaderConfig.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/StringMap.h"

#include <memory>
#include <variant>

namespace mlir {
class AsmResourcePrinter;
class AsmDialectResourceHandle;
class Operation;

namespace detail {
class AsmStateImpl;
} // namespace detail

//===----------------------------------------------------------------------===//
// Resources
//===----------------------------------------------------------------------===//

/// The following classes enable support for parsing and printing resources
/// within MLIR assembly formats. Resources are a mechanism by which dialects,
/// and external clients, may attach additional information when parsing or
/// printing IR without that information being encoded in the IR itself.
/// Resources are not uniqued within the MLIR context, are not attached directly
/// to any operation, and are solely intended to live and be processed outside
/// of the immediate IR.
///
/// Resources are encoded using a key-value pair nested within dictionaries
/// anchored either on a dialect, or an externally registered entity.
/// Dictionaries anchored on dialects use the dialect namespace directly, and
/// dictionaries anchored on external entities use a provided unique identifier.
/// The resource key is an identifier used to disambiguate the data. The
/// resource value may be stored in various limited forms, but general encodings
/// use a string (human readable) or blob format (binary). Within the textual
/// format, an example may be of the form:
///
/// {-#
///   // The `dialect_resources` section within the file-level metadata
///   // dictionary is used to contain any dialect resource entries.
///   dialect_resources: {
///     // Here is a dictionary anchored on "foo_dialect", which is a dialect
///     // namespace.
///     foo_dialect: {
///       // `some_dialect_resource` is a key to be interpreted by the dialect,
///       // and used to initialize/configure/etc.
///       some_dialect_resource: "Some important resource value"
///     }
///   },
///   // The `external_resources` section within the file-level metadata
///   // dictionary is used to contain any non-dialect resource entries.
///   external_resources: {
///     // Here is a dictionary anchored on "mlir_reproducer", which is an
///     // external entity representing MLIR's crash reproducer functionality.
///     mlir_reproducer: {
///       // `pipeline` is an entry that holds a crash reproducer pipeline
///       // resource.
///       pipeline: "func.func(canonicalize,cse)"
///     }
///   }
/// #-}
///

//===----------------------------------------------------------------------===//
// Resource Entry
//===----------------------------------------------------------------------===//

class HeapAsmResourceBlob;

/// This class represents a processed binary blob of data. A resource blob is
/// essentially a collection of data, potentially mutable, with an associated
/// deleter function (used if the data needs to be destroyed).
class AsmResourceBlob {
public:
  /// A deleter function that frees a blob given the data, allocation size, and
  /// allocation aligment.
  using DeleterFn =
      llvm::unique_function<void(void *data, size_t size, size_t align)>;

  //===--------------------------------------------------------------------===//
  // Construction
  //===--------------------------------------------------------------------===//

  AsmResourceBlob() = default;
  AsmResourceBlob(ArrayRef<char> data, size_t dataAlignment, DeleterFn deleter,
                  bool dataIsMutable)
      : data(data), dataAlignment(dataAlignment), deleter(std::move(deleter)),
        dataIsMutable(dataIsMutable) {}
  /// Utility constructor that initializes a blob with a non-char type T.
  template <typename T, typename DelT>
  AsmResourceBlob(ArrayRef<T> data, DelT &&deleteFn, bool dataIsMutable)
      : data((const char *)data.data(), data.size() * sizeof(T)),
        dataAlignment(alignof(T)),
        deleter([deleteFn = std::forward<DelT>(deleteFn)](
                    void *data, size_t size, size_t align) {
          return deleteFn((T *)data, size, align);
        }),
        dataIsMutable(dataIsMutable) {}
  AsmResourceBlob(AsmResourceBlob &&) = default;
  AsmResourceBlob &operator=(AsmResourceBlob &&rhs) {
    // Delete the current blob if necessary.
    if (deleter)
      deleter(const_cast<char *>(data.data()), data.size(), dataAlignment);

    // Take the data entries from rhs.
    data = rhs.data;
    dataAlignment = rhs.dataAlignment;
    deleter = std::move(rhs.deleter);
    dataIsMutable = rhs.dataIsMutable;
    return *this;
  }
  AsmResourceBlob(const AsmResourceBlob &) = delete;
  AsmResourceBlob &operator=(const AsmResourceBlob &) = delete;
  ~AsmResourceBlob() {
    if (deleter)
      deleter(const_cast<char *>(data.data()), data.size(), dataAlignment);
  }

  //===--------------------------------------------------------------------===//
  // Data Access
  //===--------------------------------------------------------------------===//

  /// Return the alignment of the underlying data.
  size_t getDataAlignment() const { return dataAlignment; }

  /// Return the raw underlying data of this blob.
  ArrayRef<char> getData() const { return data; }

  /// Return the underlying data as an array of the given type. This is an
  /// inherrently unsafe operation, and should only be used when the data is
  /// known to be of the correct type.
  template <typename T>
  ArrayRef<T> getDataAs() const {
    return llvm::ArrayRef<T>((const T *)data.data(), data.size() / sizeof(T));
  }

  /// Return a mutable reference to the raw underlying data of this blob.
  /// Asserts that the blob `isMutable`.
  MutableArrayRef<char> getMutableData() {
    assert(isMutable() &&
           "cannot access mutable reference to non-mutable data");
    return MutableArrayRef<char>(const_cast<char *>(data.data()), data.size());
  }

  /// Return if the data of this blob is mutable.
  bool isMutable() const { return dataIsMutable; }

  /// Return the deleter function of this blob.
  DeleterFn &getDeleter() { return deleter; }
  const DeleterFn &getDeleter() const { return deleter; }

private:
  /// The raw, properly aligned, blob data.
  ArrayRef<char> data;

  /// The alignment of the data.
  size_t dataAlignment = 0;

  /// An optional deleter function used to deallocate the underlying data when
  /// necessary.
  DeleterFn deleter;

  /// Whether the data is mutable.
  bool dataIsMutable;

  friend class HeapAsmResourceBlob;
};

/// This class provides a simple utility wrapper for creating heap allocated
/// AsmResourceBlobs.
class HeapAsmResourceBlob {
public:
  /// Create a new heap allocated blob with the given size and alignment.
  /// `dataIsMutable` indicates if the allocated data can be mutated. By
  /// default, we treat heap allocated blobs as mutable.
  static AsmResourceBlob allocate(size_t size, size_t align,
                                  bool dataIsMutable = true) {
    return AsmResourceBlob(
        ArrayRef<char>((char *)llvm::allocate_buffer(size, align), size), align,
        llvm::deallocate_buffer, dataIsMutable);
  }
  /// Create a new heap allocated blob and copy the provided data into it.
  static AsmResourceBlob allocateAndCopyWithAlign(ArrayRef<char> data,
                                                  size_t align,
                                                  bool dataIsMutable = true) {
    // This sets the blob to be mutable initially to allow writing
    // (getMutableData) below.
    AsmResourceBlob blob = allocate(data.size(), align, /*dataIsMutable=*/true);
    std::memcpy(blob.getMutableData().data(), data.data(), data.size());
    blob.dataIsMutable = dataIsMutable;
    return blob;
  }
  template <typename T>
  static AsmResourceBlob allocateAndCopyInferAlign(ArrayRef<T> data,
                                                   bool dataIsMutable = true) {
    return allocateAndCopyWithAlign(
        ArrayRef<char>((const char *)data.data(), data.size() * sizeof(T)),
        alignof(T), dataIsMutable);
  }
};
/// This class provides a simple utility wrapper for creating "unmanaged"
/// AsmResourceBlobs. The lifetime of the data provided to these blobs is
/// guaranteed to persist beyond the lifetime of this reference.
class UnmanagedAsmResourceBlob {
public:
  /// Create a new unmanaged resource directly referencing the provided data.
  /// `dataIsMutable` indicates if the allocated data can be mutated. By
  /// default, we treat unmanaged blobs as immutable.
  static AsmResourceBlob
  allocateWithAlign(ArrayRef<char> data, size_t align,
                    AsmResourceBlob::DeleterFn deleter = {},
                    bool dataIsMutable = false) {
    return AsmResourceBlob(data, align, std::move(deleter), dataIsMutable);
  }
  template <typename T>
  static AsmResourceBlob
  allocateInferAlign(ArrayRef<T> data, AsmResourceBlob::DeleterFn deleter = {},
                     bool dataIsMutable = false) {
    return allocateWithAlign(
        ArrayRef<char>((const char *)data.data(), data.size() * sizeof(T)),
        alignof(T), std::move(deleter), dataIsMutable);
  }
};

/// This class is used to build resource entries for use by the printer. Each
/// resource entry is represented using a key/value pair. The provided key must
/// be unique within the current context, which allows for a client to provide
/// resource entries without worrying about overlap with other clients.
class AsmResourceBuilder {
public:
  virtual ~AsmResourceBuilder();

  /// Build a resource entry represented by the given bool.
  virtual void buildBool(StringRef key, bool data) = 0;

  /// Build a resource entry represented by the given human-readable string
  /// value.
  virtual void buildString(StringRef key, StringRef data) = 0;

  /// Build an resource entry represented by the given binary blob data.
  virtual void buildBlob(StringRef key, ArrayRef<char> data,
                         uint32_t dataAlignment) = 0;
  /// Build an resource entry represented by the given binary blob data. This is
  /// a useful overload if the data type is known. Note that this does not
  /// support `char` element types to avoid accidentally not providing the
  /// expected alignment of data in situations that treat blobs generically.
  template <typename T>
  std::enable_if_t<!std::is_same<T, char>::value> buildBlob(StringRef key,
                                                            ArrayRef<T> data) {
    buildBlob(
        key, ArrayRef<char>((const char *)data.data(), data.size() * sizeof(T)),
        alignof(T));
  }
  /// Build an resource entry represented by the given resource blob. This is
  /// a useful overload if a blob already exists in-memory.
  void buildBlob(StringRef key, const AsmResourceBlob &blob) {
    buildBlob(key, blob.getData(), blob.getDataAlignment());
  }
};

/// This enum represents the different kinds of resource values.
enum class AsmResourceEntryKind {
  /// A blob of data with an accompanying alignment.
  Blob,
  /// A boolean value.
  Bool,
  /// A string value.
  String,
};
StringRef toString(AsmResourceEntryKind kind);

/// This class represents a single parsed resource entry.
class AsmParsedResourceEntry {
public:
  virtual ~AsmParsedResourceEntry();

  /// Return the key of the resource entry.
  virtual StringRef getKey() const = 0;

  /// Emit an error at the location of this entry.
  virtual InFlightDiagnostic emitError() const = 0;

  /// Return the kind of this value.
  virtual AsmResourceEntryKind getKind() const = 0;

  /// Parse the resource entry represented by a boolean. Returns failure if the
  /// entry does not correspond to a bool.
  virtual FailureOr<bool> parseAsBool() const = 0;

  /// Parse the resource entry represented by a human-readable string. Returns
  /// failure if the entry does not correspond to a string.
  virtual FailureOr<std::string> parseAsString() const = 0;

  /// An allocator function used to allocate memory for a blob when required.
  /// The function is provided a size and alignment, and should return an
  /// aligned allocation buffer.
  using BlobAllocatorFn =
      function_ref<AsmResourceBlob(size_t size, size_t align)>;

  /// Parse the resource entry represented by a binary blob. Returns failure if
  /// the entry does not correspond to a blob. If the blob needed to be
  /// allocated, the given allocator function is invoked.
  virtual FailureOr<AsmResourceBlob>
  parseAsBlob(BlobAllocatorFn allocator) const = 0;
  /// Parse the resource entry represented by a binary blob using heap
  /// allocation.
  FailureOr<AsmResourceBlob> parseAsBlob() const {
    return parseAsBlob([](size_t size, size_t align) {
      return HeapAsmResourceBlob::allocate(size, align);
    });
  }
};

//===----------------------------------------------------------------------===//
// Resource Parser/Printer
//===----------------------------------------------------------------------===//

/// This class represents an instance of a resource parser. This class should be
/// implemented by non-dialect clients that want to inject additional resources
/// into MLIR assembly formats.
class AsmResourceParser {
public:
  /// Create a new parser with the given identifying name. This name uniquely
  /// identifies the entries of this parser, and differentiates them from other
  /// contexts.
  AsmResourceParser(StringRef name) : name(name.str()) {}
  virtual ~AsmResourceParser();

  /// Return the name of this parser.
  StringRef getName() const { return name; }

  /// Parse the given resource entry. Returns failure if the key/data were not
  /// valid, or could otherwise not be processed correctly. Any necessary errors
  /// should be emitted with the provided entry.
  virtual LogicalResult parseResource(AsmParsedResourceEntry &entry) = 0;

  /// Return a resource parser implemented via the given callable, whose form
  /// should match that of `parseResource` above.
  template <typename CallableT>
  static std::unique_ptr<AsmResourceParser> fromCallable(StringRef name,
                                                         CallableT &&parseFn) {
    struct Processor : public AsmResourceParser {
      Processor(StringRef name, CallableT &&parseFn)
          : AsmResourceParser(name), parseFn(std::move(parseFn)) {}
      LogicalResult parseResource(AsmParsedResourceEntry &entry) override {
        return parseFn(entry);
      }

      std::decay_t<CallableT> parseFn;
    };
    return std::make_unique<Processor>(name, std::forward<CallableT>(parseFn));
  }

private:
  std::string name;
};

/// This class represents an instance of a resource printer. This class should
/// be implemented by non-dialect clients that want to inject additional
/// resources into MLIR assembly formats.
class AsmResourcePrinter {
public:
  /// Create a new printer with the given identifying name. This name uniquely
  /// identifies the entries of this printer, and differentiates them from
  /// other contexts.
  AsmResourcePrinter(StringRef name) : name(name.str()) {}
  virtual ~AsmResourcePrinter();

  /// Return the name of this printer.
  StringRef getName() const { return name; }

  /// Build any resources to include during printing, utilizing the given
  /// top-level root operation to help determine what information to include.
  /// Provided data should be registered in the form of a key/data pair, to the
  /// given builder.
  virtual void buildResources(Operation *op,
                              AsmResourceBuilder &builder) const = 0;

  /// Return a resource printer implemented via the given callable, whose form
  /// should match that of `buildResources` above.
  template <typename CallableT>
  static std::unique_ptr<AsmResourcePrinter> fromCallable(StringRef name,
                                                          CallableT &&printFn) {
    struct Printer : public AsmResourcePrinter {
      Printer(StringRef name, CallableT &&printFn)
          : AsmResourcePrinter(name), printFn(std::move(printFn)) {}
      void buildResources(Operation *op,
                          AsmResourceBuilder &builder) const override {
        printFn(op, builder);
      }

      std::decay_t<CallableT> printFn;
    };
    return std::make_unique<Printer>(name, std::forward<CallableT>(printFn));
  }

private:
  std::string name;
};

/// A fallback map containing external resources not explicitly handled by
/// another parser/printer.
class FallbackAsmResourceMap {
public:
  /// This class represents an opaque resource.
  struct OpaqueAsmResource {
    OpaqueAsmResource(StringRef key,
                      std::variant<AsmResourceBlob, bool, std::string> value)
        : key(key.str()), value(std::move(value)) {}

    /// The key identifying the resource.
    std::string key;
    /// An opaque value for the resource, whose variant values align 1-1 with
    /// the kinds defined in AsmResourceEntryKind.
    std::variant<AsmResourceBlob, bool, std::string> value;
  };

  /// Return a parser than can be used for parsing entries for the given
  /// identifier key.
  AsmResourceParser &getParserFor(StringRef key);

  /// Build a set of resource printers to print the resources within this map.
  std::vector<std::unique_ptr<AsmResourcePrinter>> getPrinters();

private:
  struct ResourceCollection : public AsmResourceParser {
    ResourceCollection(StringRef name) : AsmResourceParser(name) {}

    /// Parse a resource into this collection.
    LogicalResult parseResource(AsmParsedResourceEntry &entry) final;

    /// Build the resources held by this collection.
    void buildResources(Operation *op, AsmResourceBuilder &builder) const;

    /// The set of resources parsed into this collection.
    SmallVector<OpaqueAsmResource> resources;
  };

  /// The set of opaque resources.
  llvm::MapVector<std::string, std::unique_ptr<ResourceCollection>,
                  llvm::StringMap<unsigned>>
      keyToResources;
};

//===----------------------------------------------------------------------===//
// ParserConfig
//===----------------------------------------------------------------------===//

/// This class represents a configuration for the MLIR assembly parser. It
/// contains all of the necessary state to parse a MLIR source file.
class ParserConfig {
public:
  /// Construct a parser configuration with the given context.
  /// `verifyAfterParse` indicates if the IR should be verified after parsing.
  /// `fallbackResourceMap` is an optional fallback handler that can be used to
  /// parse external resources not explicitly handled by another parser.
  ParserConfig(MLIRContext *context, bool verifyAfterParse = true,
               FallbackAsmResourceMap *fallbackResourceMap = nullptr)
      : context(context), verifyAfterParse(verifyAfterParse),
        fallbackResourceMap(fallbackResourceMap) {
    assert(context && "expected valid MLIR context");
  }

  /// Return the MLIRContext to be used when parsing.
  MLIRContext *getContext() const { return context; }

  /// Returns if the parser should verify the IR after parsing.
  bool shouldVerifyAfterParse() const { return verifyAfterParse; }

  /// Returns the parsing configurations associated to the bytecode read.
  BytecodeReaderConfig &getBytecodeReaderConfig() const {
    return const_cast<BytecodeReaderConfig &>(bytecodeReaderConfig);
  }

  /// Return the resource parser registered to the given name, or nullptr if no
  /// parser with `name` is registered.
  AsmResourceParser *getResourceParser(StringRef name) const {
    auto it = resourceParsers.find(name);
    if (it != resourceParsers.end())
      return it->second.get();
    if (fallbackResourceMap)
      return &fallbackResourceMap->getParserFor(name);
    return nullptr;
  }

  /// Attach the given resource parser.
  void attachResourceParser(std::unique_ptr<AsmResourceParser> parser) {
    StringRef name = parser->getName();
    auto it = resourceParsers.try_emplace(name, std::move(parser));
    (void)it;
    assert(it.second &&
           "resource parser already registered with the given name");
  }

  /// Attach the given callable resource parser with the given name.
  template <typename CallableT>
  std::enable_if_t<std::is_convertible<
      CallableT, function_ref<LogicalResult(AsmParsedResourceEntry &)>>::value>
  attachResourceParser(StringRef name, CallableT &&parserFn) {
    attachResourceParser(AsmResourceParser::fromCallable(
        name, std::forward<CallableT>(parserFn)));
  }

private:
  MLIRContext *context;
  bool verifyAfterParse;
  DenseMap<StringRef, std::unique_ptr<AsmResourceParser>> resourceParsers;
  FallbackAsmResourceMap *fallbackResourceMap;
  BytecodeReaderConfig bytecodeReaderConfig;
};

//===----------------------------------------------------------------------===//
// AsmState
//===----------------------------------------------------------------------===//

/// This class provides management for the lifetime of the state used when
/// printing the IR. It allows for alleviating the cost of recomputing the
/// internal state of the asm printer.
///
/// The IR should not be mutated in-between invocations using this state, and
/// the IR being printed must not be an parent of the IR originally used to
/// initialize this state. This means that if a child operation is provided, a
/// parent operation cannot reuse this state.
class AsmState {
public:
  /// This map represents the raw locations of operations within the output
  /// stream. This maps the original pointer to the operation, to a pair of line
  /// and column in the output stream.
  using LocationMap = DenseMap<Operation *, std::pair<unsigned, unsigned>>;

  /// Initialize the asm state at the level of the given operation. A location
  /// map may optionally be provided to be populated when printing. `map` is an
  /// optional fallback resource map, which when provided will attach resource
  /// printers for the fallback resources within the map.
  AsmState(Operation *op,
           const OpPrintingFlags &printerFlags = OpPrintingFlags(),
           LocationMap *locationMap = nullptr,
           FallbackAsmResourceMap *map = nullptr);
  AsmState(MLIRContext *ctx,
           const OpPrintingFlags &printerFlags = OpPrintingFlags(),
           LocationMap *locationMap = nullptr,
           FallbackAsmResourceMap *map = nullptr);
  ~AsmState();

  /// Get the printer flags.
  const OpPrintingFlags &getPrinterFlags() const;

  /// Return an instance of the internal implementation. Returns nullptr if the
  /// state has not been initialized.
  detail::AsmStateImpl &getImpl() { return *impl; }

  //===--------------------------------------------------------------------===//
  // Resources
  //===--------------------------------------------------------------------===//

  /// Attach the given resource printer to the AsmState.
  void attachResourcePrinter(std::unique_ptr<AsmResourcePrinter> printer);

  /// Attach an resource printer, in the form of a callable, to the AsmState.
  template <typename CallableT>
  std::enable_if_t<std::is_convertible<
      CallableT, function_ref<void(Operation *, AsmResourceBuilder &)>>::value>
  attachResourcePrinter(StringRef name, CallableT &&printFn) {
    attachResourcePrinter(AsmResourcePrinter::fromCallable(
        name, std::forward<CallableT>(printFn)));
  }

  /// Attach resource printers to the AsmState for the fallback resources
  /// in the given map.
  void attachFallbackResourcePrinter(FallbackAsmResourceMap &map) {
    for (auto &printer : map.getPrinters())
      attachResourcePrinter(std::move(printer));
  }

  /// Returns a map of dialect resources that were referenced when using this
  /// state to print IR.
  DenseMap<Dialect *, SetVector<AsmDialectResourceHandle>> &
  getDialectResources() const;

private:
  AsmState() = delete;

  /// A pointer to allocated storage for the impl state.
  std::unique_ptr<detail::AsmStateImpl> impl;
};

//===----------------------------------------------------------------------===//
// AsmPrinter CommandLine Options
//===----------------------------------------------------------------------===//

/// Register a set of useful command-line options that can be used to configure
/// various flags within the AsmPrinter.
void registerAsmPrinterCLOptions();

} // namespace mlir

#endif // MLIR_IR_ASMSTATE_H_