/* blake3_device.cuh — BLAKE3 __device__ primitives.
 *
 * Derived from Blaze-3/BLAKE3-gpu cuda/blaze3.cuh + cuda/blaze3_cpu.cuh
 * (MIT, Rehan Vipin 2021 — see LICENSE in this directory).  Stripped
 * to the bare compression primitives needed for a single-input,
 * single-thread chunk-tree walk.  Tree-mode merkle / dynamic-parallelism
 * removed — see NOTICE.
 *
 * Constants and round function track the BLAKE3 spec
 * (https://github.com/BLAKE3-team/BLAKE3/blob/master/reference_impl/
 * reference_impl.rs) byte-for-byte.
 *
 * AGPLv3 on the surrounding wrapper, MIT on this header — see NOTICE.
 */
#ifndef BLAKE3_DEVICE_CUH
#define BLAKE3_DEVICE_CUH

#include <stdint.h>

namespace blake3_gpu {

typedef uint8_t  u8;
typedef uint32_t u32;
typedef uint64_t u64;

constexpr u32 OUT_LEN   = 32;
constexpr u32 KEY_LEN   = 32;
constexpr u32 BLOCK_LEN = 64;
constexpr u32 CHUNK_LEN = 1024;

constexpr u32 CHUNK_START         = 1U << 0;
constexpr u32 CHUNK_END           = 1U << 1;
constexpr u32 PARENT              = 1U << 2;
constexpr u32 ROOT                = 1U << 3;

__constant__ static const u32 IV[8] = {
    0x6A09E667U, 0xBB67AE85U, 0x3C6EF372U, 0xA54FF53AU,
    0x510E527FU, 0x9B05688CU, 0x1F83D9ABU, 0x5BE0CD19U,
};

__constant__ static const int MSG_PERMUTATION[16] = {
    2, 6, 3, 10, 7, 0, 4, 13,
    1, 11, 12, 5, 9, 14, 15, 8,
};

__device__ __forceinline__ u32 rotr32(u32 v, int n) {
    return (v >> n) | (v << (32 - n));
}

__device__ __forceinline__ void g(u32 *st, int a, int b, int c, int d,
                                  u32 mx, u32 my) {
    st[a] = st[a] + st[b] + mx;
    st[d] = rotr32(st[d] ^ st[a], 16);
    st[c] = st[c] + st[d];
    st[b] = rotr32(st[b] ^ st[c], 12);
    st[a] = st[a] + st[b] + my;
    st[d] = rotr32(st[d] ^ st[a], 8);
    st[c] = st[c] + st[d];
    st[b] = rotr32(st[b] ^ st[c], 7);
}

__device__ __forceinline__ void round_fn(u32 *st, const u32 *m) {
    /* columns */
    g(st, 0, 4,  8, 12, m[0],  m[1]);
    g(st, 1, 5,  9, 13, m[2],  m[3]);
    g(st, 2, 6, 10, 14, m[4],  m[5]);
    g(st, 3, 7, 11, 15, m[6],  m[7]);
    /* diagonals */
    g(st, 0, 5, 10, 15, m[8],  m[9]);
    g(st, 1, 6, 11, 12, m[10], m[11]);
    g(st, 2, 7,  8, 13, m[12], m[13]);
    g(st, 3, 4,  9, 14, m[14], m[15]);
}

__device__ __forceinline__ void permute_msg(u32 *m) {
    u32 p[16];
    for (int i = 0; i < 16; i++) p[i] = m[MSG_PERMUTATION[i]];
    for (int i = 0; i < 16; i++) m[i] = p[i];
}

/* compress: writes a 16-word state. After 7 rounds and the spec's XOR
 * fold, state[0..8] is the chaining value (also the first 32 B of root
 * output when the ROOT flag is set). state[8..16] forms the second
 * 32 B of root output. Counter / block_len / flags follow the BLAKE3
 * compression-function contract. */
__device__ __forceinline__ void compress(const u32 *cv,
                                         const u32 *block_words,
                                         u64 counter,
                                         u32 block_len,
                                         u32 flags,
                                         u32 *out16)
{
    u32 st[16];
    st[0]  = cv[0]; st[1]  = cv[1]; st[2]  = cv[2]; st[3]  = cv[3];
    st[4]  = cv[4]; st[5]  = cv[5]; st[6]  = cv[6]; st[7]  = cv[7];
    st[8]  = IV[0]; st[9]  = IV[1]; st[10] = IV[2]; st[11] = IV[3];
    st[12] = (u32)counter;
    st[13] = (u32)(counter >> 32);
    st[14] = block_len;
    st[15] = flags;

    u32 m[16];
    for (int i = 0; i < 16; i++) m[i] = block_words[i];

    round_fn(st, m); permute_msg(m);
    round_fn(st, m); permute_msg(m);
    round_fn(st, m); permute_msg(m);
    round_fn(st, m); permute_msg(m);
    round_fn(st, m); permute_msg(m);
    round_fn(st, m); permute_msg(m);
    round_fn(st, m); /* round 7, no permute after */

    for (int i = 0; i < 8; i++) {
        st[i]     ^= st[i + 8];
        st[i + 8] ^= cv[i];
    }
    for (int i = 0; i < 16; i++) out16[i] = st[i];
}

/* Load up to BLOCK_LEN bytes little-endian into 16 u32 words; tail
 * positions zero. */
__device__ __forceinline__ void words_from_bytes_le(const u8 *bytes,
                                                    u32 len,
                                                    u32 *words)
{
    for (int i = 0; i < 16; i++) words[i] = 0U;
    for (u32 i = 0; i < len; i++) {
        words[i >> 2] |= ((u32)bytes[i]) << (8 * (i & 3));
    }
}

} /* namespace blake3_gpu */

#endif /* BLAKE3_DEVICE_CUH */
