uhash-core 0.5.1

//! Core UniversalHash v4 Implementation (Spec-Compliant)
//!
//! The algorithm uses 4 parallel chains with 512KB scratchpads each.
//! Each chain performs rounds using raw compression functions (AES, SHA-256, BLAKE3)
//! in a sequential pattern that prevents GPU parallelism.
//!
//! This implementation follows the spec exactly:
//! - Seed generation: BLAKE3(header || (nonce ⊕ (c × golden_ratio)))
//! - Primitive rotation: (nonce + c) mod 3, then +1 before each round
//! - Write-back: Same address as read (not computed from new state)
//! - No cross-chain mixing (spec doesn't specify it)

#[cfg(not(feature = "std"))]
use alloc::vec;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;

use blake3::Hasher as Blake3;
use sha2::{Digest, Sha256};

use crate::params::*;
use crate::primitives::{aes_compress, blake3_compress, sha256_compress};

/// Mask for address calculation (BLOCKS_PER_SCRATCHPAD - 1)
/// Since BLOCKS_PER_SCRATCHPAD = 8192 = 2^13, this is 0x1FFF
const ADDRESS_MASK: usize = BLOCKS_PER_SCRATCHPAD - 1;

/// Golden ratio constant for seed generation (Fibonacci hashing constant)
const GOLDEN_RATIO: u64 = 0x9E3779B97F4A7C15;

#[inline(always)]
fn initial_primitive_index(nonce: u64, chain: usize) -> usize {
    // Keep primitive selection architecture-stable across 32-bit/64-bit targets.
    // (nonce + chain) mod 3 must use full u64 nonce, not truncated usize.
    ((nonce.wrapping_add(chain as u64)) % 3) as usize
}

/// UniversalHash v4 hasher
///
/// This struct maintains the scratchpads and chain states needed for hashing.
/// It can be reused for multiple hashes to avoid repeated allocations.
pub struct UniversalHash {
    /// 4 scratchpads, one per chain (512KB each)
    scratchpads: Vec<Vec<u8>>,
    /// Current state for each chain
    chain_states: [[u8; 32]; CHAINS],
    /// Effective nonce extracted from input (last 8 bytes)
    effective_nonce: u64,
}

impl UniversalHash {
    /// Create a new UniversalHash instance
    ///
    /// Allocates 2MB of memory for the scratchpads.
    pub fn new() -> Self {
        Self {
            scratchpads: vec![vec![0u8; SCRATCHPAD_SIZE]; CHAINS],
            chain_states: [[0u8; 32]; CHAINS],
            effective_nonce: 0,
        }
    }

    /// Compute the UniversalHash of input data
    ///
    /// The input should be formatted as:
    /// `epoch_seed || miner_address || timestamp || nonce`
    /// where nonce is the last 8 bytes.
    ///
    /// Returns a 32-byte hash.
    pub fn hash(&mut self, input: &[u8]) -> [u8; 32] {
        // Extract effective nonce from last 8 bytes of input (or hash if shorter)
        self.effective_nonce = extract_nonce(input);

        // Phase 1: Initialize scratchpads using input (spec-compliant seed generation)
        self.init_scratchpads(input);

        // Phase 2: Execute main mixing rounds (spec-compliant, no cross-chain mixing)
        self.execute_rounds();

        // Phase 3: Finalize and produce output
        self.finalize()
    }

    /// Initialize all scratchpads from input using expansion
    /// Spec: seed[c] = BLAKE3_256(header || (nonce ⊕ (c × golden_ratio)))
    fn init_scratchpads(&mut self, input: &[u8]) {
        let nonce = self.effective_nonce;
        let header_len = input.len().saturating_sub(8);

        for (chain, state) in self.chain_states.iter_mut().enumerate() {
            // Spec: nonce ⊕ (c × golden_ratio)
            let offset = (chain as u64).wrapping_mul(GOLDEN_RATIO);
            let modified_nonce = nonce ^ offset;

            // Spec: BLAKE3(header || modified_nonce)
            let mut hasher = Blake3::new();
            hasher.update(&input[..header_len]);
            hasher.update(&modified_nonce.to_le_bytes());
            let hash = hasher.finalize();

            let hash_bytes = hash.as_bytes();
            state.copy_from_slice(hash_bytes);

            // Fill scratchpad using AES-based expansion
            let mut seed_array = [0u8; 32];
            seed_array.copy_from_slice(hash_bytes);
            fill_scratchpad_aes(&mut self.scratchpads[chain], &seed_array);
        }
    }

    /// Execute the main mixing rounds (spec-compliant: no cross-chain mixing)
    fn execute_rounds(&mut self) {
        let nonce = self.effective_nonce;

        // Process each chain independently (spec-compliant: no cross-chain mixing)
        for chain in 0..CHAINS {
            // Spec: primitive = (nonce + c) mod 3
            let initial_primitive = initial_primitive_index(nonce, chain);

            // Execute all rounds for this chain
            for round in 0..ROUNDS {
                round_step_spec_compliant(
                    &mut self.scratchpads[chain],
                    &mut self.chain_states[chain],
                    initial_primitive,
                    round,
                );
            }
        }
    }

    /// Finalize and produce the 32-byte output hash per spec
    /// Spec: result = BLAKE3_256(SHA256_256(combined))
    fn finalize(&self) -> [u8; 32] {
        // XOR all chain states together
        let mut combined = [0u8; 32];
        for state in &self.chain_states {
            for i in 0..32 {
                combined[i] ^= state[i];
            }
        }

        // Double hash: SHA256 then BLAKE3 (per spec)
        let sha_hash = Sha256::digest(combined);
        let mut hasher = Blake3::new();
        hasher.update(&sha_hash);
        hasher.finalize().into()
    }
}

/// Extract nonce from input (last 8 bytes, or hash if shorter)
#[inline(always)]
fn extract_nonce(input: &[u8]) -> u64 {
    if input.len() >= 8 {
        // Use last 8 bytes as nonce
        let nonce_bytes: [u8; 8] = input[input.len() - 8..].try_into().unwrap();
        u64::from_le_bytes(nonce_bytes)
    } else {
        // For short inputs, hash to get a nonce
        let hash = blake3::hash(input);
        let bytes: [u8; 8] = hash.as_bytes()[..8].try_into().unwrap();
        u64::from_le_bytes(bytes)
    }
}

/// Fill a scratchpad using AES-based expansion per spec
/// Spec (§5.3.2 + reference implementation §11):
///   key = AES_KeyExpand(seed[0:16])
///   state = seed[16:32]
///   For i = 0 to NUM_BLOCKS - 1:
///     4 sequential AES_4Rounds expansions fill each 64-byte block
///     state carries forward from the 4th expansion
#[inline(always)]
fn fill_scratchpad_aes(scratchpad: &mut [u8], seed: &[u8; 32]) {
    use crate::primitives::{aes128_key_expand, aes_expand_block};

    let raw_key: [u8; 16] = seed[0..16].try_into().unwrap();
    let round_keys = aes128_key_expand(&raw_key);
    let mut state: [u8; 16] = seed[16..32].try_into().unwrap();

    for i in 0..BLOCKS_PER_SCRATCHPAD {
        let offset = i * BLOCK_SIZE;

        // 4 sequential AES expansions — each 16-byte chunk is unique
        state = aes_expand_block(&state, &round_keys);
        scratchpad[offset..offset + 16].copy_from_slice(&state);

        state = aes_expand_block(&state, &round_keys);
        scratchpad[offset + 16..offset + 32].copy_from_slice(&state);

        state = aes_expand_block(&state, &round_keys);
        scratchpad[offset + 32..offset + 48].copy_from_slice(&state);

        state = aes_expand_block(&state, &round_keys);
        scratchpad[offset + 48..offset + 64].copy_from_slice(&state);
    }
}

/// Single round step for one chain (spec-compliant version)
///
/// Spec:
/// - Address: computed from current state
/// - Primitive: (initial_primitive + round + 1) mod 3  (increment BEFORE use)
/// - Write-back: SAME address as read (not new address)
#[inline(always)]
fn round_step_spec_compliant(
    scratchpad: &mut [u8],
    state: &mut [u8; 32],
    initial_primitive: usize,
    round: usize,
) {
    // Compute memory address from state per spec formula
    let addr = compute_address(state, round);

    // Read block from scratchpad
    // SAFETY: addr is always within bounds due to ADDRESS_MASK
    let block: [u8; BLOCK_SIZE] =
        unsafe { core::ptr::read(scratchpad.as_ptr().add(addr) as *const [u8; BLOCK_SIZE]) };

    // Spec: primitive = (primitive + 1) mod 3 BEFORE applying
    // Where primitive starts at (nonce + chain) mod 3
    // So at round r: primitive = (initial_primitive + r + 1) mod 3
    let primitive = (initial_primitive + round + 1) % 3;

    // Apply raw compression function based on primitive
    let new_state = match primitive {
        0 => aes_compress(state, &block),
        1 => sha256_compress(state, &block),
        _ => blake3_compress(state, &block),
    };

    // Spec: Write back to SAME address as read (not computed from new_state!)
    // SAFETY: addr is always within bounds due to ADDRESS_MASK
    unsafe {
        core::ptr::copy_nonoverlapping(new_state.as_ptr(), scratchpad.as_mut_ptr().add(addr), 32);
    }

    // Update chain state
    *state = new_state;
}

/// Compute scratchpad address from state per spec
/// Spec: mixed = state[0:8] ⊕ state[8:16] ⊕ rotl64(round, 13) ⊕ (round × 0x517cc1b727220a95)
///       addr = (mixed mod NUM_BLOCKS) × BLOCK_SIZE
#[inline(always)]
fn compute_address(state: &[u8; 32], round: usize) -> usize {
    const MIXING_CONSTANT: u64 = 0x517cc1b727220a95;

    // Read u64s directly using pointer reads (faster than try_into)
    // SAFETY: state is 32 bytes, reading at offsets 0 and 8 is safe
    let state_lo = unsafe { core::ptr::read_unaligned(state.as_ptr() as *const u64) };
    let state_hi = unsafe { core::ptr::read_unaligned(state.as_ptr().add(8) as *const u64) };
    let round_u64 = round as u64;

    // Spec formula for unpredictable address
    let mixed =
        state_lo ^ state_hi ^ round_u64.rotate_left(13) ^ round_u64.wrapping_mul(MIXING_CONSTANT);

    // Use bitwise AND instead of modulo (NUM_BLOCKS is power of 2)
    ((mixed as usize) & ADDRESS_MASK) * BLOCK_SIZE
}

impl Default for UniversalHash {
    fn default() -> Self {
        Self::new()
    }
}

/// Convenience function for single-shot hashing
///
/// Creates a new hasher, computes the hash, and returns it.
/// For multiple hashes, prefer creating a `UniversalHash` instance
/// and reusing it to avoid repeated memory allocation.
pub fn hash(input: &[u8]) -> [u8; 32] {
    let mut hasher = UniversalHash::new();
    hasher.hash(input)
}