/// Poseidon hash over BN254 (T=3, 2-input).

///

/// In Solidity this is generated as raw EVM bytecode by JavaScript.

/// In Fe it's readable code with typed field arithmetic.

///

/// All pure functions are const fn, enabling compile-time evaluation

/// via CTFE. The static_asserts below verify known-answer vectors at

/// compile time with zero gas cost.

pub use constants::{PRIME, C, S, M, P_MAT, N_ROUNDS_F, N_ROUNDS_P}

pub use fp::Fp

/// MDS matrix multiply: state = mat * state

const fn mds(s0: Fp, s1: Fp, s2: Fp, mat: [[u256; 3]; 3]) -> (Fp, Fp, Fp) {

(

Fp::new(val: mat[0][0]) * s0 + Fp::new(val: mat[1][0]) * s1 + Fp::new(val: mat[2][0]) * s2,

Fp::new(val: mat[0][1]) * s0 + Fp::new(val: mat[1][1]) * s1 + Fp::new(val: mat[2][1]) * s2,

Fp::new(val: mat[0][2]) * s0 + Fp::new(val: mat[1][2]) * s1 + Fp::new(val: mat[2][2]) * s2,

)

}

/// Full S-box round: pow5 + add round constants + MDS mix.

const fn full_round(s0: Fp, s1: Fp, s2: Fp, rc: usize, mat: [[u256; 3]; 3]) -> (Fp, Fp, Fp) {

let s0 = s0.pow5() + Fp::new(val: C[rc])

let s1 = s1.pow5() + Fp::new(val: C[rc + 1])

let s2 = s2.pow5() + Fp::new(val: C[rc + 2])

mds(s0, s1, s2, mat)

}

/// Poseidon hash of two field elements.

pub const fn hash(a: u256, b: u256) -> u256 {

let mut s0 = Fp::new(val: 0) + Fp::new(val: C[0])

let mut s1 = Fp::new(val: a) + Fp::new(val: C[1])

let mut s2 = Fp::new(val: b) + Fp::new(val: C[2])

// First half full rounds

let mut r: usize = 0

while r < N_ROUNDS_F / 2 - 1 {

let (r0, r1, r2) = full_round(s0, s1, s2, rc: (r + 1) * 3, mat: M)

s0 = r0

s1 = r1

s2 = r2

r += 1

}

// Pre-sparse round (uses P_MAT instead of M)

let (r0, r1, r2) = full_round(s0, s1, s2, rc: (N_ROUNDS_F / 2) * 3, mat: P_MAT)

s0 = r0

s1 = r1

s2 = r2

// Partial rounds (S-box on first element only, sparse matrix)

let mut r: usize = 0

while r < N_ROUNDS_P {

s0 = s0.pow5() + Fp::new(val: C[(N_ROUNDS_F / 2 + 1) * 3 + r])

let off: usize = 5 * r

let new_s0 = Fp::new(val: S[off]) * s0

+ Fp::new(val: S[off + 1]) * s1

+ Fp::new(val: S[off + 2]) * s2

s1 = s1 + s0 * Fp::new(val: S[off + 3])

s2 = s2 + s0 * Fp::new(val: S[off + 4])

s0 = new_s0

r += 1

}

// Second half full rounds

let mut r: usize = 0

while r < N_ROUNDS_F / 2 - 1 {

let (r0, r1, r2) = full_round(

s0,

s1,

s2,

rc: (N_ROUNDS_F / 2 + 1) * 3 + N_ROUNDS_P + r * 3,

mat: M,

)

s0 = r0

s1 = r1

s2 = r2

r += 1

}

// Final round (no round constants)

let s0 = s0.pow5()

let s1 = s1.pow5()

let s2 = s2.pow5()

let (out, _, _) = mds(s0, s1, s2, mat: M)

out.val

}

// Known-answer vectors from circomlibjs buildPoseidon(), verified at compile time.

static_assert(hash(a: 0, b: 0)

== 0x2098f5fb9e239eab3ceac3f27b81e481dc3124d55ffed523a839ee8446b64864)

static_assert(hash(a: 1, b: 2)

== 0x115cc0f5e7d690413df64c6b9662e9cf2a3617f2743245519e19607a4417189a)

static_assert(hash(a: 42, b: 17)

== 0x18ddf7be412cfc792364a46eba440316dbb2427346063b747899e6ee0c3aa214)