// https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.202.pdf
// based on keccak-f[1600] permutation

import "utils/casts/u8_to_bits";
import "utils/casts/u8_from_bits";
import "utils/casts/u64_to_bits";
import "utils/casts/u64_from_bits";

const u32[24] RHO = [
    1,  3,  6,  10, 15, 21, 28, 36, 45, 55, 2,  14,
    27, 41, 56, 8,  25, 43, 62, 18, 39, 61, 20, 44
];

const u32[24] PI = [
    10, 7,  11, 17, 18, 3, 5,  16, 8,  21, 24, 4,
    15, 23, 19, 13, 12, 2, 20, 14, 22, 9,  6,  1
];

const u64[24] RC = [
    0x0000000000000001, 0x0000000000008082, 0x800000000000808a,
    0x8000000080008000, 0x000000000000808b, 0x0000000080000001,
    0x8000000080008081, 0x8000000000008009, 0x000000000000008a,
    0x0000000000000088, 0x0000000080008009, 0x000000008000000a,
    0x000000008000808b, 0x800000000000008b, 0x8000000000008089,
    0x8000000000008003, 0x8000000000008002, 0x8000000000000080,
    0x000000000000800a, 0x800000008000000a, 0x8000000080008081,
    0x8000000000008080, 0x0000000080000001, 0x8000000080008008
];

// left rotation
def rotl64(u64 x, u32 n) -> u64 {
    return ((x << n) | (x >> (64 - n)));
}

// change endianness
def swap_u64(u64 mut val) -> u64 {
    val = ((val << 8) & 0xFF00FF00FF00FF00) | ((val >> 8) & 0x00FF00FF00FF00FF);
    val = ((val << 16) & 0xFFFF0000FFFF0000) | ((val >> 16) & 0x0000FFFF0000FFFF);
    return (val << 32) | (val >> 32);
}

// compression function
def keccakf(u64[25] mut st) -> u64[25] {
    u64[5] mut bc = [0; 5];
    u64 mut t = 0;

    for u32 i in 0..25 {
        st[i] = swap_u64(st[i]);
    }

    for u32 r in 0..24 {
        // theta
        for u32 i in 0..5 {
            bc[i] = st[i] ^ st[i + 5] ^ st[i + 10] ^ st[i + 15] ^ st[i + 20];
        }

        for u32 i in 0..5 {
            t = bc[(i + 4) % 5] ^ rotl64(bc[(i + 1) % 5], 1);
            for u32 j in 0..5 {
                st[(j * 5) + i] = st[(j * 5) + i] ^ t;
            }
        }

        t = st[1];

        // rho pi
        for u32 i in 0..24 {
            u32 j = PI[i];
            bc[0] = st[j];
            st[j] = rotl64(t, RHO[i]);
            t = bc[0];
        }

        // chi
        for u32 i in 0..5 {
            for u32 j in 0..5 {
                bc[j] = st[(i * 5) + j];
            }
            for u32 j in 0..5 {
                u32 p = (i * 5) + j;
                st[p] = st[p] ^ (!bc[(j + 1) % 5] & bc[(j + 2) % 5]);
            }
        }

        // iota
        st[0] = st[0] ^ RC[r];
    }

    for u32 i in 0..25 {
        st[i] = swap_u64(st[i]);
    }

    return st;
}

def u64_from_u8_array(u8[8] input) -> u64 {
    bool[64] bits = [
        ...u8_to_bits(input[0]),
        ...u8_to_bits(input[1]),
        ...u8_to_bits(input[2]),
        ...u8_to_bits(input[3]),
        ...u8_to_bits(input[4]),
        ...u8_to_bits(input[5]),
        ...u8_to_bits(input[6]),
        ...u8_to_bits(input[7])
    ];
    return u64_from_bits(bits);
}

def u64_to_u8_array(u64 input) -> u8[8] {
    bool[64] bits = u64_to_bits(input);
    return [
        u8_from_bits(bits[0..8]),
        u8_from_bits(bits[8..16]),
        u8_from_bits(bits[16..24]),
        u8_from_bits(bits[24..32]),
        u8_from_bits(bits[32..40]),
        u8_from_bits(bits[40..48]),
        u8_from_bits(bits[48..56]),
        u8_from_bits(bits[56..64])
    ];
}

def to_bytes(u64[25] input) -> u8[200] {
    u8[200] mut output = [0; 200];
    for u32 i in 0..25 {
        u8[8] t = u64_to_u8_array(input[i]);
        for u32 j in 0..8 {
            output[i * 8 + j] = t[j];
        }
    }
    return output;
}

def from_bytes(u8[200] input) -> u64[25] {
    u64[25] mut output = [0; 25];
    for u32 i in 0..25 {
        output[i] = u64_from_u8_array(input[i*8..i*8+8]);
    }
    return output;
}

def main<N, W>(u8[N] mut input, u8 delim) -> u8[W] {
    u8[200] mut b = [0; 200];
    u32 rate = 200 - (2 * W);
    u32 mut pt = 0;

    // update
    for u32 i in 0..N {
        b[pt] = b[pt] ^ input[i];
        pt = (pt + 1) % rate;
        b = pt == 0 ? to_bytes(keccakf(from_bytes(b))) : b;
    }

    // finalize
    b[pt] = b[pt] ^ delim;
    b[rate - 1] = b[rate - 1] ^ 0x80;
    b = to_bytes(keccakf(from_bytes(b)));

    return b[..W];
}