#include "mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir::quant {

// Permutes `values` with `permutation`. Returns the permuted values. Sizes of
// `values` and `permutation` must be equal, and the elements of `permutation`
// should be less than `values.size()`.
template <typename T,
          typename = std::enable_if_t<std::is_default_constructible_v<T>, void>>

  // Operation data format after folding `permutation`.
  StringRef target_data_format = [&]() -> StringRef {
    if (op->getDataFormat() == "NHWC" && permutation.equals(kNchwToNhwc)) {
      return "NCHW";  // cancel NCHW->NHWC operand permutation
    } else if (op->getDataFormat() == "NCHW" &&
               permutation.equals(kNhwcToNchw)) {
      return "NHWC";  // cancel NHWC->NCHW operand permutation
    } else {
      return "";
    }

Type ShuffleRankedTensorType(Type type, ArrayRef<int64_t> permutation) {
  if (auto ranked_type = mlir::dyn_cast<RankedTensorType>(type)) {
    ArrayRef<int64_t> shape = ranked_type.getShape();
    assert(permutation.size() == shape.size());

    SmallVector<int64_t, 4> new_shape(permutation.size());
    for (size_t i = 0; i < permutation.size(); ++i)
      new_shape[i] = shape[permutation[i]];

        permutation_(permutation) {}

  // Transposes `values` with the permutation. Returns the transposed values.
  SmallVector<float> TransposeValues(const ArrayRef<float> values) const {
    SmallVector<float> transposed_values(values.size());
    SmallVector<int64_t> current_indices = {};
    TransposeRecursively(values, transposed_values, current_indices);

    return transposed_values;

// Returns a ShapedType for a permutation and the shape of input after
// applying the permutation to the given shape through a transpose.
class GetTransposedType<string perm> : NativeCodeCall<
  "GetTransposedType($0, " # perm # ")">;

// Function to map final permutation to initial permutation
// initial -> permutation1 -> permutation2 -> final
def RemapPermutation: NativeCodeCall<"RemapPermutation($0, $1)">;

var feistelBox [8][64]uint32

var feistelBoxOnce sync.Once

// general purpose function to perform DES block permutations.
func permuteBlock(src uint64, permutation []uint8) (block uint64) {
	for position, n := range permutation {
		bit := (src >> n) & 1
		block |= bit << uint((len(permutation)-1)-position)
	}
	return
}

func initFeistelBox() {
	for s := range sBoxes {
		for i := 0; i < 4; i++ {

    def "resolve pair #permutation"() {
        given:
        def candidates = permutation.candidates
        def expectedSingle = permutation.expectedSingle
        def expectedMulti = permutation.expectedMulti

        expect:
        checkScenarioResolution(expectedSingle, expectedMulti, candidates)

        where:
        permutation << VersionRangeResolveTestScenarios.SCENARIOS_TWO_DEPENDENCIES_BATCH1

    def "resolve prefer pair #permutation"() {
        given:
        def candidates = permutation.candidates
        def expectedSingle = permutation.expectedSingle
        def expectedMulti = permutation.expectedMulti

        expect:
        checkScenarioResolution(expectedSingle, expectedMulti, candidates)

        where:
        permutation << VersionRangeResolveTestScenarios.SCENARIOS_PREFER_BATCH2
    }

namespace odml {

struct PermutationAndShape {
  DenseIntElementsAttr permutation;
  ShapedType shape;
};

// Check that `arr` is an R1 iota with integer element type starting from `0`
// with `size` number of values.
bool IsIotaAttr(ArrayRef<int64_t> arr, int64_t size);

// Returns a DenseIntElementsAttr for a permutation and the shape after
// applying the permutation to a given shape through a transpose.

    def "resolve reject pair #permutation"() {
        given:
        def candidates = permutation.candidates
        def expectedSingle = permutation.expectedSingle
        def expectedMulti = permutation.expectedMulti

        expect:
        checkScenarioResolution(expectedSingle, expectedMulti, candidates)

        where:
        permutation << VersionRangeResolveTestScenarios.SCENARIOS_DEPENDENCY_WITH_REJECT