/*

 * Copyright (c) 2023, the SerenityOS developers

 *

 * SPDX-License-Identifier: BSD-2-Clause

 */

#include <AK/ByteReader.h>

#include <LibCrypto/Hash/BLAKE2b.h>

namespace Crypto::Hash {

constexpr static auto ROTRIGHT(u64 a, size_t b) { return (a >> b) | (a << (64 - b)); }

void BLAKE2b::update(u8 const* in, size_t inlen)

{

    if (inlen > 0) {

        size_t left = m_internal_state.buffer_length;

        size_t fill = BLAKE2bConstants::blockbytes - left;

        if (inlen > fill) {

            m_internal_state.buffer_length = 0;

            // Fill the buffer.

            __builtin_memcpy(m_internal_state.buffer + left, in, fill);

            increment_counter_by(BLAKE2bConstants::blockbytes);

            transform(m_internal_state.buffer);

            in += fill;

            inlen -= fill;

            while (inlen > BLAKE2bConstants::blockbytes) {

                increment_counter_by(BLAKE2bConstants::blockbytes);

                transform(in);

                in += BLAKE2bConstants::blockbytes;

                inlen -= BLAKE2bConstants::blockbytes;

            }

        }

        __builtin_memcpy(m_internal_state.buffer + m_internal_state.buffer_length, in, inlen);

        m_internal_state.buffer_length += inlen;

    }

}

BLAKE2b::DigestType BLAKE2b::peek()

{

    DigestType digest;

    increment_counter_by(m_internal_state.buffer_length);

    // Set this as the last block

    m_internal_state.is_at_last_block = UINT64_MAX;

    // Pad the buffer with zeros

    __builtin_memset(m_internal_state.buffer + m_internal_state.buffer_length, 0, BLAKE2bConstants::blockbytes - m_internal_state.buffer_length);

    transform(m_internal_state.buffer);

    for (size_t i = 0; i < 8; ++i)

        __builtin_memcpy(&digest.data[0] + sizeof(m_internal_state.hash_state[i]) * i, &m_internal_state.hash_state[i], sizeof(m_internal_state.hash_state[i]));

    return digest;

}

BLAKE2b::DigestType BLAKE2b::digest()

{

    auto digest = peek();

    reset();

    return digest;

}

void BLAKE2b::increment_counter_by(u64 const amount)

{

    m_internal_state.message_byte_offset[0] += amount;

    m_internal_state.message_byte_offset[1] += (m_internal_state.message_byte_offset[0] < amount);

}

void BLAKE2b::mix(u64* work_array, u64 a, u64 b, u64 c, u64 d, u64 x, u64 y)

{

    constexpr auto rotation_constant_1 = 32;

    constexpr auto rotation_constant_2 = 24;

    constexpr auto rotation_constant_3 = 16;

    constexpr auto rotation_constant_4 = 63;

    work_array[a] = work_array[a] + work_array[b] + x;

    work_array[d] = ROTRIGHT(work_array[d] ^ work_array[a], rotation_constant_1);

    work_array[c] = work_array[c] + work_array[d];

    work_array[b] = ROTRIGHT(work_array[b] ^ work_array[c], rotation_constant_2);

    work_array[a] = work_array[a] + work_array[b] + y;

    work_array[d] = ROTRIGHT(work_array[d] ^ work_array[a], rotation_constant_3);

    work_array[c] = work_array[c] + work_array[d];

    work_array[b] = ROTRIGHT(work_array[b] ^ work_array[c], rotation_constant_4);

}

void BLAKE2b::transform(u8 const* block)

{

    u64 m[16];

    u64 v[16];

    for (size_t i = 0; i < 16; ++i)

        m[i] = ByteReader::load64(block + i * sizeof(m[i]));

    for (size_t i = 0; i < 8; ++i)

        v[i] = m_internal_state.hash_state[i];

    v[8] = SHA512Constants::InitializationHashes[0];

    v[9] = SHA512Constants::InitializationHashes[1];

    v[10] = SHA512Constants::InitializationHashes[2];

    v[11] = SHA512Constants::InitializationHashes[3];

    v[12] = SHA512Constants::InitializationHashes[4] ^ m_internal_state.message_byte_offset[0];

    v[13] = SHA512Constants::InitializationHashes[5] ^ m_internal_state.message_byte_offset[1];

    v[14] = SHA512Constants::InitializationHashes[6] ^ m_internal_state.is_at_last_block;

    v[15] = SHA512Constants::InitializationHashes[7];

    for (size_t i = 0; i < 12; ++i) {

        u64 sigma_selection[16];

        for (size_t j = 0; j < 16; ++j)

            sigma_selection[j] = BLAKE2bSigma[i % 10][j];

        mix(v, 0, 4, 8, 12, m[sigma_selection[0]], m[sigma_selection[1]]);

        mix(v, 1, 5, 9, 13, m[sigma_selection[2]], m[sigma_selection[3]]);

        mix(v, 2, 6, 10, 14, m[sigma_selection[4]], m[sigma_selection[5]]);

        mix(v, 3, 7, 11, 15, m[sigma_selection[6]], m[sigma_selection[7]]);

        mix(v, 0, 5, 10, 15, m[sigma_selection[8]], m[sigma_selection[9]]);

        mix(v, 1, 6, 11, 12, m[sigma_selection[10]], m[sigma_selection[11]]);

        mix(v, 2, 7, 8, 13, m[sigma_selection[12]], m[sigma_selection[13]]);

        mix(v, 3, 4, 9, 14, m[sigma_selection[14]], m[sigma_selection[15]]);

    }

    for (size_t i = 0; i < 8; ++i)

        m_internal_state.hash_state[i] = m_internal_state.hash_state[i] ^ v[i] ^ v[i + 8];

}

}