//===----------------------------------------------------------------------===//
//
// 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 common utilities for generating C++ from tablegen
// structures.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_TABLEGEN_CODEGENHELPERS_H
#define MLIR_TABLEGEN_CODEGENHELPERS_H

#include "mlir/TableGen/Constraint.h"
#include "mlir/TableGen/Dialect.h"
#include "mlir/TableGen/Format.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/TableGen/CodeGenHelpers.h"
#include <utility>

namespace llvm {
class RecordKeeper;
} // namespace llvm

namespace mlir {
namespace tblgen {
class Constraint;
class DagLeaf;

// Format into a std::string
template <typename... Parameters>
std::string strfmt(const char *fmt, Parameters &&...parameters) {
  return llvm::formatv(fmt, std::forward<Parameters>(parameters)...).str();
}

// A helper RAII class to emit nested namespaces for a dialect.
class DialectNamespaceEmitter {
public:
  DialectNamespaceEmitter(raw_ostream &os, const Dialect &dialect) {
    if (!dialect)
      return;
    nsEmitter.emplace(os, dialect.getCppNamespace());
  }

private:
  std::optional<llvm::NamespaceEmitter> nsEmitter;
};

/// This class represents how an error stream string being constructed will be
/// consumed.
enum class ErrorStreamType {
  // Inside a string that's streamed into an InflightDiagnostic.
  InString,
  // Inside a string inside an OpError.
  InsideOpError,
};

/// This class deduplicates shared operation verification code by emitting
/// static functions alongside the op definitions. These methods are local to
/// the definition file, and are invoked within the operation verify methods.
/// An example is shown below:
///
/// static LogicalResult localVerify(...)
///
/// LogicalResult OpA::verify(...) {
///  if (failed(localVerify(...)))
///    return failure();
///  ...
/// }
///
/// LogicalResult OpB::verify(...) {
///  if (failed(localVerify(...)))
///    return failure();
///  ...
/// }
///
class StaticVerifierFunctionEmitter {
public:
  /// Create a constraint uniquer with a unique prefix derived from the record
  /// keeper with an optional tag.
  StaticVerifierFunctionEmitter(raw_ostream &os,
                                const llvm::RecordKeeper &records,
                                StringRef tag = "");

  /// Collect and unique all the constraints used by operations.
  void collectOpConstraints(ArrayRef<const llvm::Record *> opDefs);

  /// Collect and unique all compatible type, attribute, successor, and region
  /// constraints from the operations in the file and emit them at the top of
  /// the generated file.
  ///
  /// Constraints that do not meet the restriction that they can only reference
  /// `$_self` and `$_op` are not uniqued.
  void emitOpConstraints();

  /// Unique all compatible type and attribute constraints from a pattern file
  /// and emit them at the top of the generated file.
  ///
  /// Constraints that do not meet the restriction that they can only reference
  /// `$_self`, `$_op`, and `$_builder` are not uniqued.
  void emitPatternConstraints(const ArrayRef<DagLeaf> constraints);

  /// Get the name of the static function used for the given type constraint.
  /// These functions are used for operand and result constraints and have the
  /// form:
  ///
  ///   LogicalResult(Operation *op, Type type, StringRef valueKind,
  ///                 unsigned valueIndex);
  ///
  /// Pattern constraints have the form:
  ///
  ///   LogicalResult(PatternRewriter &rewriter, Operation *op, Type type,
  ///                 StringRef failureStr);
  ///
  StringRef getTypeConstraintFn(const Constraint &constraint) const;

  /// Get the name of the static function used for the given attribute
  /// constraint. These functions are in the form:
  ///
  ///   LogicalResult(Operation *op, Attribute attr, StringRef attrName);
  ///
  /// If a uniqued constraint was not found, this function returns std::nullopt.
  /// The uniqued constraints cannot be used in the context of an OpAdaptor.
  ///
  /// Pattern constraints have the form:
  ///
  ///   LogicalResult(PatternRewriter &rewriter, Operation *op, Attribute attr,
  ///                 StringRef failureStr);
  ///
  std::optional<StringRef>
  getAttrConstraintFn(const Constraint &constraint) const;

  /// Get the name of the static function used for the given property
  /// constraint. These functions are in the form:
  ///
  ///   LogicalResult(Operation *op, T property, StringRef propName);
  ///
  /// where T is the interface type specified in the constraint.
  /// If a uniqued constraint was not found, this function returns std::nullopt.
  /// The uniqued constraints cannot be used in the context of an OpAdaptor.
  ///
  /// Pattern constraints have the form:
  ///
  ///   LogicalResult(PatternRewriter &rewriter, Operation *op, T property,
  ///                 StringRef failureStr);
  ///
  std::optional<StringRef>
  getPropConstraintFn(const Constraint &constraint) const;

  /// Get the name of the static function used for the given successor
  /// constraint. These functions are in the form:
  ///
  ///   LogicalResult(Operation *op, Block *successor, StringRef successorName,
  ///                 unsigned successorIndex);
  ///
  StringRef getSuccessorConstraintFn(const Constraint &constraint) const;

  /// Get the name of the static function used for the given region constraint.
  /// These functions are in the form:
  ///
  ///   LogicalResult(Operation *op, Region &region, StringRef regionName,
  ///                 unsigned regionIndex);
  ///
  /// The region name may be empty.
  StringRef getRegionConstraintFn(const Constraint &constraint) const;

private:
  /// Emit static type constraint functions.
  void emitTypeConstraints();
  /// Emit static attribute constraint functions.
  void emitAttrConstraints();
  /// Emit static property constraint functions.
  void emitPropConstraints();
  /// Emit static successor constraint functions.
  void emitSuccessorConstraints();
  /// Emit static region constraint functions.
  void emitRegionConstraints();

  /// Emit pattern constraints.
  void emitPatternConstraints();

  /// Collect and unique all pattern constraints.
  void collectPatternConstraints(ArrayRef<DagLeaf> constraints);

  /// The output stream.
  raw_ostream &os;

  /// A unique label for the file currently being generated. This is used to
  /// ensure that the static functions have a unique name.
  std::string uniqueOutputLabel;

  /// Use a MapVector to ensure that functions are generated deterministically.
  using ConstraintMap = llvm::MapVector<Constraint, std::string,
                                        llvm::DenseMap<Constraint, unsigned>>;

  /// A generic function to emit constraints
  void emitConstraints(const ConstraintMap &constraints, StringRef selfName,
                       const char *codeTemplate,
                       ErrorStreamType errorStreamType);

  /// Assign a unique name to a unique constraint.
  std::string getUniqueName(StringRef kind, unsigned index);
  /// Unique a constraint in the map.
  void collectConstraint(ConstraintMap &map, StringRef kind,
                         Constraint constraint);

  /// The set of type constraints used for operand and result verification in
  /// the current file.
  ConstraintMap typeConstraints;
  /// The set of attribute constraints used in the current file.
  ConstraintMap attrConstraints;
  /// The set of property constraints used in the current file.
  ConstraintMap propConstraints;
  /// The set of successor constraints used in the current file.
  ConstraintMap successorConstraints;
  /// The set of region constraints used in the current file.
  ConstraintMap regionConstraints;
};

/// Escape a string using C++ encoding. E.g. foo"bar -> foo\x22bar.
std::string escapeString(StringRef value);

namespace detail {
template <typename>
struct stringifier {
  template <typename T>
  static std::string apply(T &&t) {
    return std::string(std::forward<T>(t));
  }
};
template <>
struct stringifier<Twine> {
  static std::string apply(const Twine &twine) { return twine.str(); }
};
template <typename OptionalT>
struct stringifier<std::optional<OptionalT>> {
  static std::string apply(std::optional<OptionalT> optional) {
    return optional ? stringifier<OptionalT>::apply(*optional) : std::string();
  }
};
} // namespace detail

/// Generically convert a value to a std::string.
template <typename T>
std::string stringify(T &&t) {
  return detail::stringifier<std::remove_reference_t<std::remove_const_t<T>>>::
      apply(std::forward<T>(t));
}

/// Helper to generate a C++ streaming error message from a given message.
/// Message can contain '{{...}}' placeholders that are substituted with
/// C-expressions via tgfmt. It would effectively convert:
///   "failed to verify {{foo}}"
/// into:
///   "failed to verify " << bar
/// where bar is the result of evaluating 'tgfmt("foo", &ctx)' at compile
/// time.
std::string buildErrorStreamingString(
    StringRef message, const FmtContext &ctx,
    ErrorStreamType errorStreamType = ErrorStreamType::InString);

} // namespace tblgen
} // namespace mlir

#endif // MLIR_TABLEGEN_CODEGENHELPERS_H