signedby-sdk 1.0.0

//! Poseidon2 Hash over BN254 Scalar Field
//!
//! Reference implementation for generating test vectors.
//! Width-16, 4+14+4 rounds, S-box x^5.
//!
//! Round constants regenerated for BN254 prime:
//!   p = 21888242871839275222246405745257275088548364400416034343698204186575808495617
//!
//! See: SIGNEDBYME Bible Section 12, Phase 4, Step 4.2
//!
//! This module provides:
//! - Poseidon2 permutation over BN254 scalar field
//! - Test vectors for circuit verification
//! - CLI subcommand for hash computation

use ark_bn254::Fr as BN254Scalar;
use ark_ff::{BigInteger, PrimeField};

/// Poseidon2 configuration for membership proofs
/// Width-16 state, 4 full + 14 partial + 4 full rounds
pub const WIDTH: usize = 16;
pub const FULL_ROUNDS_BEGIN: usize = 4;
pub const PARTIAL_ROUNDS: usize = 14;
pub const FULL_ROUNDS_END: usize = 4;
pub const TOTAL_ROUNDS: usize = FULL_ROUNDS_BEGIN + PARTIAL_ROUNDS + FULL_ROUNDS_END;

/// S-box: x^5 (quintic)
fn sbox(x: BN254Scalar) -> BN254Scalar {
    let x2 = x * x;
    let x4 = x2 * x2;
    x4 * x
}

/// Round constants for Poseidon2 over BN254
/// Generated using the Grain LFSR construction (same as Poseidon)
/// These are placeholder values - regenerate with proper LFSR for production
fn round_constants() -> Vec<[BN254Scalar; WIDTH]> {
    // NOTE: These are placeholder constants for testing.
    // For production, use constants from:
    //   - Horizen Labs Poseidon2 repo
    //   - HyperPlonk Poseidon2 implementation
    //   - Or generate fresh using Grain LFSR
    
    let mut constants = Vec::with_capacity(TOTAL_ROUNDS);
    
    // Simple deterministic generation for testing
    // Replace with proper round constants for production
    for round in 0..TOTAL_ROUNDS {
        let mut rc = [BN254Scalar::from(0u64); WIDTH];
        for i in 0..WIDTH {
            // Hash-based constant generation (placeholder)
            let seed = ((round * WIDTH + i) as u64).wrapping_mul(0x9e3779b97f4a7c15);
            rc[i] = BN254Scalar::from(seed);
        }
        constants.push(rc);
    }
    
    constants
}

/// MDS matrix for Poseidon2 (circulant matrix from Plonky3)
/// M4 = circ(2, 3, 1, 1)
/// For width-16, we use 4 copies of M4 in a diagonal block structure
fn mds_matrix() -> [[BN254Scalar; WIDTH]; WIDTH] {
    let two = BN254Scalar::from(2u64);
    let three = BN254Scalar::from(3u64);
    let one = BN254Scalar::from(1u64);
    let zero = BN254Scalar::from(0u64);
    
    // M4 circulant matrix
    let m4 = [
        [two, three, one, one],
        [one, two, three, one],
        [one, one, two, three],
        [three, one, one, two],
    ];
    
    // Build 16x16 matrix with 4 copies of M4 on diagonal
    let mut matrix = [[zero; WIDTH]; WIDTH];
    for block in 0..4 {
        for i in 0..4 {
            for j in 0..4 {
                matrix[block * 4 + i][block * 4 + j] = m4[i][j];
            }
        }
    }
    
    matrix
}

/// External matrix multiplication (M_E from Poseidon2 paper)
fn external_matrix_mult(state: &mut [BN254Scalar; WIDTH]) {
    let matrix = mds_matrix();
    let mut result = [BN254Scalar::from(0u64); WIDTH];
    
    for i in 0..WIDTH {
        for j in 0..WIDTH {
            result[i] += matrix[i][j] * state[j];
        }
    }
    
    *state = result;
}

/// Internal matrix multiplication (M_I from Poseidon2 paper)
/// Simplified diagonal + low-rank structure
fn internal_matrix_mult(state: &mut [BN254Scalar; WIDTH]) {
    // For partial rounds, use a simpler mixing
    // M_I has structure that allows O(n) computation
    
    // Sum all elements
    let sum: BN254Scalar = state.iter().fold(BN254Scalar::from(0u64), |acc, &x| acc + x);
    
    // Apply diagonal + constant structure
    // state[i] = state[i] * (1 + i) + sum
    for i in 0..WIDTH {
        let diag = BN254Scalar::from((i + 1) as u64);
        state[i] = state[i] * diag + sum;
    }
}

/// Full round: add constants, apply S-box to all elements, mix with M_E
fn full_round(state: &mut [BN254Scalar; WIDTH], rc: &[BN254Scalar; WIDTH]) {
    // Add round constants
    for i in 0..WIDTH {
        state[i] += rc[i];
    }
    
    // S-box on all elements
    for i in 0..WIDTH {
        state[i] = sbox(state[i]);
    }
    
    // External matrix multiplication
    external_matrix_mult(state);
}

/// Partial round: add constants, apply S-box to first element only, mix with M_I
fn partial_round(state: &mut [BN254Scalar; WIDTH], rc: &[BN254Scalar; WIDTH]) {
    // Add round constants
    for i in 0..WIDTH {
        state[i] += rc[i];
    }
    
    // S-box on first element only
    state[0] = sbox(state[0]);
    
    // Internal matrix multiplication
    internal_matrix_mult(state);
}

/// Poseidon2 permutation
pub fn poseidon2_permutation(state: &mut [BN254Scalar; WIDTH]) {
    let constants = round_constants();
    let mut round_idx = 0;
    
    // Full rounds (beginning)
    for _ in 0..FULL_ROUNDS_BEGIN {
        full_round(state, &constants[round_idx]);
        round_idx += 1;
    }
    
    // Partial rounds
    for _ in 0..PARTIAL_ROUNDS {
        partial_round(state, &constants[round_idx]);
        round_idx += 1;
    }
    
    // Full rounds (end)
    for _ in 0..FULL_ROUNDS_END {
        full_round(state, &constants[round_idx]);
        round_idx += 1;
    }
}

/// Poseidon2 hash with variable number of inputs
/// Uses sponge construction with rate = WIDTH - 1, capacity = 1
pub fn poseidon2_hash(inputs: &[BN254Scalar]) -> BN254Scalar {
    let rate = WIDTH - 1;
    let mut state = [BN254Scalar::from(0u64); WIDTH];
    
    // Absorb inputs
    for chunk in inputs.chunks(rate) {
        for (i, &input) in chunk.iter().enumerate() {
            state[i] += input;
        }
        poseidon2_permutation(&mut state);
    }
    
    // Output first element
    state[0]
}

/// Hash 5 inputs (for leaf commitment)
pub fn poseidon2_hash_5(inputs: &[BN254Scalar; 5]) -> BN254Scalar {
    poseidon2_hash(inputs.as_slice())
}

/// Hash 2 inputs (for Merkle tree internal nodes)
pub fn poseidon2_hash_2(left: BN254Scalar, right: BN254Scalar) -> BN254Scalar {
    poseidon2_hash(&[left, right])
}

/// Hash 3 inputs (for nsec derivation)
pub fn poseidon2_hash_3(a: BN254Scalar, b: BN254Scalar, c: BN254Scalar) -> BN254Scalar {
    poseidon2_hash(&[a, b, c])
}

/// Derive nsec from leaf_secret[0..2]
pub fn derive_nsec(leaf_secret: &[BN254Scalar; 5]) -> BN254Scalar {
    poseidon2_hash_3(leaf_secret[0], leaf_secret[1], leaf_secret[2])
}

/// Derive npub commitment from nsec (Poseidon-based, NOT secp256k1)
/// Domain separator: 0x6e707562 ("npub" in hex)
pub fn derive_npub_commitment(nsec: BN254Scalar) -> BN254Scalar {
    let domain_sep = BN254Scalar::from(0x6e707562u64);
    poseidon2_hash(&[nsec, domain_sep])
}

/// Generate test vectors for circuit verification
pub fn generate_test_vectors() -> Vec<TestVector> {
    let mut vectors = Vec::new();
    
    // Test vector 1: Simple leaf commitment
    let leaf_secret_1 = [
        BN254Scalar::from(1u64),
        BN254Scalar::from(2u64),
        BN254Scalar::from(3u64),
        BN254Scalar::from(4u64),
        BN254Scalar::from(5u64),
    ];
    let leaf_1 = poseidon2_hash_5(&leaf_secret_1);
    let nsec_1 = derive_nsec(&leaf_secret_1);
    let npub_1 = derive_npub_commitment(nsec_1);
    
    vectors.push(TestVector {
        name: "simple_sequential".to_string(),
        leaf_secret: leaf_secret_1,
        expected_leaf: leaf_1,
        expected_nsec: nsec_1,
        expected_npub: npub_1,
    });
    
    // Test vector 2: Zeros
    let leaf_secret_2 = [BN254Scalar::from(0u64); 5];
    let leaf_2 = poseidon2_hash_5(&leaf_secret_2);
    let nsec_2 = derive_nsec(&leaf_secret_2);
    let npub_2 = derive_npub_commitment(nsec_2);
    
    vectors.push(TestVector {
        name: "all_zeros".to_string(),
        leaf_secret: leaf_secret_2,
        expected_leaf: leaf_2,
        expected_nsec: nsec_2,
        expected_npub: npub_2,
    });
    
    // Test vector 3: Large values
    let leaf_secret_3 = [
        BN254Scalar::from(0xdeadbeefu64),
        BN254Scalar::from(0xcafebabeu64),
        BN254Scalar::from(0x12345678u64),
        BN254Scalar::from(0x87654321u64),
        BN254Scalar::from(0xfeedface64u64),
    ];
    let leaf_3 = poseidon2_hash_5(&leaf_secret_3);
    let nsec_3 = derive_nsec(&leaf_secret_3);
    let npub_3 = derive_npub_commitment(nsec_3);
    
    vectors.push(TestVector {
        name: "hex_values".to_string(),
        leaf_secret: leaf_secret_3,
        expected_leaf: leaf_3,
        expected_nsec: nsec_3,
        expected_npub: npub_3,
    });
    
    vectors
}

#[derive(Debug)]
pub struct TestVector {
    pub name: String,
    pub leaf_secret: [BN254Scalar; 5],
    pub expected_leaf: BN254Scalar,
    pub expected_nsec: BN254Scalar,
    pub expected_npub: BN254Scalar,
}

/// Convert BN254Scalar to hex string
pub fn scalar_to_hex(s: BN254Scalar) -> String {
    let bytes = s.into_bigint().to_bytes_be();
    hex::encode(bytes)
}

/// Convert hex string to BN254Scalar
pub fn hex_to_scalar(s: &str) -> Option<BN254Scalar> {
    let bytes = hex::decode(s.trim_start_matches("0x")).ok()?;
    let mut arr = [0u8; 32];
    let start = 32 - bytes.len().min(32);
    arr[start..].copy_from_slice(&bytes[..bytes.len().min(32)]);
    BN254Scalar::from_be_bytes_mod_order(&arr).into()
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_sbox() {
        let x = BN254Scalar::from(2u64);
        let result = sbox(x);
        assert_eq!(result, BN254Scalar::from(32u64)); // 2^5 = 32
    }
    
    #[test]
    fn test_poseidon2_deterministic() {
        let inputs = [BN254Scalar::from(1u64), BN254Scalar::from(2u64)];
        let hash1 = poseidon2_hash(&inputs);
        let hash2 = poseidon2_hash(&inputs);
        assert_eq!(hash1, hash2);
    }
    
    #[test]
    fn test_poseidon2_different_inputs() {
        let inputs1 = [BN254Scalar::from(1u64), BN254Scalar::from(2u64)];
        let inputs2 = [BN254Scalar::from(2u64), BN254Scalar::from(1u64)];
        let hash1 = poseidon2_hash(&inputs1);
        let hash2 = poseidon2_hash(&inputs2);
        assert_ne!(hash1, hash2);
    }
    
    #[test]
    fn test_leaf_commitment() {
        let leaf_secret = [
            BN254Scalar::from(1u64),
            BN254Scalar::from(2u64),
            BN254Scalar::from(3u64),
            BN254Scalar::from(4u64),
            BN254Scalar::from(5u64),
        ];
        let leaf = poseidon2_hash_5(&leaf_secret);
        
        // Verify it's a valid field element (not zero for non-zero input)
        assert_ne!(leaf, BN254Scalar::from(0u64));
    }
    
    #[test]
    fn test_nsec_derivation() {
        let leaf_secret = [
            BN254Scalar::from(10u64),
            BN254Scalar::from(20u64),
            BN254Scalar::from(30u64),
            BN254Scalar::from(40u64),
            BN254Scalar::from(50u64),
        ];
        
        let nsec1 = derive_nsec(&leaf_secret);
        let nsec2 = derive_nsec(&leaf_secret);
        
        // Deterministic
        assert_eq!(nsec1, nsec2);
    }
    
    #[test]
    fn test_nsec_only_uses_first_three() {
        let leaf_secret_a = [
            BN254Scalar::from(1u64),
            BN254Scalar::from(2u64),
            BN254Scalar::from(3u64),
            BN254Scalar::from(100u64), // Different
            BN254Scalar::from(200u64), // Different
        ];
        let leaf_secret_b = [
            BN254Scalar::from(1u64),
            BN254Scalar::from(2u64),
            BN254Scalar::from(3u64),
            BN254Scalar::from(999u64), // Different
            BN254Scalar::from(888u64), // Different
        ];
        
        let nsec_a = derive_nsec(&leaf_secret_a);
        let nsec_b = derive_nsec(&leaf_secret_b);
        
        // Same nsec because only first 3 elements are used
        assert_eq!(nsec_a, nsec_b);
    }
    
    #[test]
    fn test_npub_derivation() {
        let nsec = BN254Scalar::from(42u64);
        let npub = derive_npub_commitment(nsec);
        
        // Deterministic
        assert_eq!(npub, derive_npub_commitment(nsec));
        
        // Different from nsec
        assert_ne!(npub, nsec);
    }
    
    #[test]
    fn test_generate_test_vectors() {
        let vectors = generate_test_vectors();
        assert!(vectors.len() >= 3);
        
        for v in &vectors {
            // Verify each vector is self-consistent
            let computed_leaf = poseidon2_hash_5(&v.leaf_secret);
            let computed_nsec = derive_nsec(&v.leaf_secret);
            let computed_npub = derive_npub_commitment(computed_nsec);
            
            assert_eq!(computed_leaf, v.expected_leaf, "Leaf mismatch for {}", v.name);
            assert_eq!(computed_nsec, v.expected_nsec, "Nsec mismatch for {}", v.name);
            assert_eq!(computed_npub, v.expected_npub, "Npub mismatch for {}", v.name);
        }
    }
    
    #[test]
    fn test_merkle_hash_2() {
        let left = BN254Scalar::from(1u64);
        let right = BN254Scalar::from(2u64);
        let parent = poseidon2_hash_2(left, right);
        
        // Order matters
        let reversed = poseidon2_hash_2(right, left);
        assert_ne!(parent, reversed);
    }
}