#ifndef TRITON_DIALECT_TRITON_TRANSFORMS_FUNCTION_TYPE_CONVERSION_H_
#define TRITON_DIALECT_TRITON_TRANSFORMS_FUNCTION_TYPE_CONVERSION_H_
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Transforms/DialectConversion.h"

namespace mlir::triton {

/**
 * @brief Rename the aggregated function arguments that were generated by the
 * type converter.
 *
 * @note The callback should return a std::optional<LogicalResult>. If the
 * callback returns a std::nullopt, the argument is not renamed and the
 * conversion continues. If the callback returns a LogicalResult, the out_suffix
 * list is populated with the suffixes for each aggregated argument.
 *
 * @param delimiter The delimiter to use between the base name and the suffixes.
 */
class FuncArgRenamer {
public:
  FuncArgRenamer(const char *delimiter = ".") : delimiter(delimiter) {}
  template <typename FnT, typename T = typename llvm::function_traits<
                              std::decay_t<FnT>>::template arg_t<0>>
  void addRenamer(FnT &&callback) {
    renamers.emplace_back(wrapCallback<T>(std::forward<FnT>(callback)));
  }
  LogicalResult apply(Type type, FunctionOpInterface funcOp, int index,
                      TypeConverter::SignatureConversion &conversion) const;

private:
  using RenamerCallbackFn = std::function<std::optional<LogicalResult>(
      Type, llvm::SmallVectorImpl<std::string> &)>;

  /**
   * @brief Wraps a callback of form `std::optional<LogicalResult>(T,
   * llvm::SmallVectorImpl<std::string> &)` into a RenamerCallbackFn.
   *
   * @tparam T The type of the argument.
   * @tparam FnT The type of the callback.
   * @param callback The callback to wrap.
   * @return A RenamerCallbackFn.
   */
  template <typename T, typename FnT>
  std::enable_if_t<
      std::is_invocable_v<FnT, T, llvm::SmallVectorImpl<std::string> &> &&
          std::is_base_of_v<Type, T>,
      RenamerCallbackFn>
  wrapCallback(FnT &&callback) const {
    return [callback = std::forward<FnT>(callback)](
               Type type, llvm::SmallVectorImpl<std::string> &out_suffix)
               -> std::optional<LogicalResult> {
      if (auto t = dyn_cast<T>(type)) {
        return callback(t, out_suffix);
      }
      return std::nullopt;
    };
  }

  llvm::SmallVector<RenamerCallbackFn> renamers;
  const char *delimiter;
};

/**
 * @brief Provides helper patterns for converting triton function operations
 * using a type converter.
 *
 * Note we cannot use upstream passes for this because they are unaware of
 * tt.call and tt.return.
 */
void populateFunctionTypeConversions(const TypeConverter &converter,
                                     const FuncArgRenamer &renamer,
                                     RewritePatternSet &patterns);

} // namespace mlir::triton

#endif // TRITON_DIALECT_TRITON_TRANSFORMS_FUNCTION_TYPE_CONVERSION_H_
