// https://tools.ietf.org/html/rfc7693

import "EMBED/u32_to_bits" as to_bits
import "EMBED/u32_from_bits" as from_bits

def rotr32<N>(u32 x) -> u32:
    return (x >> N) | (x << (32 - N))

def blake2s_iv() -> (u32[8]):
    return [
        0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A,
        0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19
    ]

def blake2s_sigma() -> (u32[10][16]):
    return [
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
        [14, 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3],
        [11, 8, 12, 0, 5, 2, 15, 13, 10, 14, 3, 6, 7, 1, 9, 4],
        [7, 9, 3, 1, 13, 12, 11, 14, 2, 6, 5, 10, 4, 0, 15, 8],
        [9, 0, 5, 7, 2, 4, 10, 15, 14, 1, 11, 12, 6, 8, 3, 13],
        [2, 12, 6, 10, 0, 11, 8, 3, 4, 13, 7, 5, 15, 14, 1, 9],
        [12, 5, 1, 15, 14, 13, 4, 10, 0, 7, 6, 3, 9, 2, 8, 11],
        [13, 11, 7, 14, 12, 1, 3, 9, 5, 0, 15, 4, 8, 6, 2, 10],
        [6, 15, 14, 9, 11, 3, 0, 8, 12, 2, 13, 7, 1, 4, 10, 5],
        [10, 2, 8, 4, 7, 6, 1, 5, 15, 11, 9, 14, 3, 12, 13, 0]
    ]

def mixing_g(u32[16] v, u32 a, u32 b, u32 c, u32 d, u32 x, u32 y) -> (u32[16]):
    v[a] = (v[a] + v[b] + x)
    v[d] = rotr32::<16>(v[d] ^ v[a])
    v[c] = (v[c] + v[d])
    v[b] = rotr32::<12>(v[b] ^ v[c])
    v[a] = (v[a] + v[b] + y)
    v[d] = rotr32::<8>(v[d] ^ v[a])
    v[c] = (v[c] + v[d])
    v[b] = rotr32::<7>(v[b] ^ v[c])
    return v

def blake2s_compression(u32[8] h, u32[16] m, u32[2] t, bool last) -> (u32[8]):
    u32[16] v = [...h, ...blake2s_iv()]

    v[12] = v[12] ^ t[0]
    v[13] = v[13] ^ t[1]
    v[14] = if last then v[14] ^ 0xFFFFFFFF else v[14] fi

    u32[10][16] sigma = blake2s_sigma()

    for u32 i in 0..10 do
        u32[16] s = sigma[i]
        v = mixing_g(v, 0, 4,  8, 12, m[s[0]],  m[s[1]])
        v = mixing_g(v, 1, 5,  9, 13, m[s[2]],  m[s[3]])
        v = mixing_g(v, 2, 6, 10, 14, m[s[4]],  m[s[5]])
        v = mixing_g(v, 3, 7, 11, 15, m[s[6]],  m[s[7]])
        v = mixing_g(v, 0, 5, 10, 15, m[s[8]],  m[s[9]])
        v = mixing_g(v, 1, 6, 11, 12, m[s[10]], m[s[11]])
        v = mixing_g(v, 2, 7,  8, 13, m[s[12]], m[s[13]])
        v = mixing_g(v, 3, 4,  9, 14, m[s[14]], m[s[15]])
    endfor

    for u32 i in 0..8 do
        h[i] = h[i] ^ v[i] ^ v[i + 8]
    endfor

    return h

def blake2s_init(u32[2] p) -> (u32[8]):
    u32[8] iv = blake2s_iv()
    u32[8] h = [
        iv[0] ^ 0x01010000 ^ 0x00000020,
        iv[1],
        iv[2],
        iv[3],
        iv[4],
        iv[5],
        iv[6] ^ p[0],
        iv[7] ^ p[1]
    ]
    return h

def main<K>(u32[K][16] input, u32[2] p) -> (u32[8]):
    u32[8] h = blake2s_init(p)

    u32 t0 = 0
    u32 t1 = 0

    for u32 i in 0..K-1 do
        t0 = (i + 1) * 64
        t1 = if t0 == 0 then t1 + 1 else t1 fi
        h = blake2s_compression(h, input[i], [t0, t1], false)
    endfor

    t0 = t0 + 64
    t1 = if t0 == 0 then t1 + 1 else t1 fi

    h = blake2s_compression(h, input[K - 1], [t0, t1], true)
    return h