// https://eprint.iacr.org/2019/458.pdf

from "./constants.zok" import POSEIDON_C, POSEIDON_M;

def ark<N>(field[N] mut state, field[497] c, u32 it) -> field[N] {
    for u32 i in 0..N {
        state[i] = state[i] + c[it + i];
    }
    return state;
}

def sbox<N>(field[N] mut state, u32 f, u32 p, u32 r) -> field[N] {
    state[0] = state[0]**5;
    for u32 i in 1..N {
        state[i] = ((r < f/2) || (r >= f/2 + p)) ? state[i]**5 : state[i];
    }
    return state;
}

def mix<N>(field[N] state, field[7][7] m) -> field[N] {
    field[N] mut out = [0; N];
    for u32 i in 0..N {
        field mut acc = 0;
        for u32 j in 0..N {
            acc = acc + (state[j] * m[i][j]);
        }
        out[i] = acc;
    }
    return out;
}

def main<N>(field[N] inputs) -> field {
    assert(N > 0 && N <= 6); // max 6 inputs

    u32 t = N + 1;
    u32[8] rounds_p = [56, 57, 56, 60, 60, 63, 64, 63];

    u32 f = 8;
    u32 p = rounds_p[(t - 2)];

    // Constants are padded with zeroes to the maximum value calculated by
    // t * (f + p) = 497, where `t` (number of inputs + 1) is a max of 7.
    // This is done to keep the function generic, as resulting array size depends on `t`
    // and we do not want callers passing down constants.
    // This should be revisited once compiler limitations are gone.

    field[497]  c = POSEIDON_C[t - 2];
    field[7][7] m = POSEIDON_M[t - 2];

    field[t] mut state = [0; t];
    for u32 i in 1..t {
        state[i] = inputs[i - 1];
    }

    for u32 r in 0..f+p {
        state = ark(state, c, r * t);
        state = sbox(state, f, p, r);
        state = mix(state, m);
    }

    return state[0];
}