#ifndef LIB_TRANSFORMS_CONVERTTOCIPHERTEXTSEMANTICS_ASSIGNLAYOUT_H_
#define LIB_TRANSFORMS_CONVERTTOCIPHERTEXTSEMANTICS_ASSIGNLAYOUT_H_

#include <cstdint>
#include <functional>

#include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
#include "mlir/include/mlir/IR/Attributes.h"            // from @llvm-project
#include "mlir/include/mlir/IR/Builders.h"              // from @llvm-project
#include "mlir/include/mlir/IR/ImplicitLocOpBuilder.h"  // from @llvm-project
#include "mlir/include/mlir/IR/Operation.h"             // from @llvm-project
#include "mlir/include/mlir/IR/PatternMatch.h"          // from @llvm-project
#include "mlir/include/mlir/IR/Value.h"                 // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"             // from @llvm-project
#include "mlir/include/mlir/Support/LogicalResult.h"    // from @llvm-project

namespace mlir {
namespace heir {

enum class CodegenStrategy {
  AUTO = 0,
  NEVER_FOLD = 1,
  FOLD_WHEN_POSSIBLE = 2,
};

// Lower tensor_ext.assign_layout. Returns the final value produced by the
// encoding implementation. Applies createdOpCallback to each created
// operation.
FailureOr<Value> implementAssignLayout(
    Value input, Attribute layout, int64_t minSlotCount,
    ImplicitLocOpBuilder& builder,
    const std::function<void(Operation*)>& createdOpCallback,
    ArrayRef<int64_t> domainSchedule = {},
    CodegenStrategy strategy = CodegenStrategy::AUTO);

// Lower tensor_ext.unpack. Returns the final value produced by the unpacking
// implementation. Applies createdOpCallback to each created operation.
FailureOr<Value> implementUnpackOp(
    tensor_ext::UnpackOp op, ImplicitLocOpBuilder& builder,
    const std::function<void(Operation*)>& createdOpCallback =
        [](Operation* op) {});

// A pattern to wrap implementUnpackOp
struct LowerUnpackOp : public OpRewritePattern<tensor_ext::UnpackOp> {
  using OpRewritePattern<tensor_ext::UnpackOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(tensor_ext::UnpackOp op,
                                PatternRewriter& rewriter) const override {
    ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
    auto res = implementUnpackOp(op, builder);
    if (failed(res)) return failure();
    rewriter.replaceOp(op, res.value());
    return success();
  }
};

}  // namespace heir
}  // namespace mlir

#endif  // LIB_TRANSFORMS_CONVERTTOCIPHERTEXTSEMANTICS_ASSIGNLAYOUT_H_
