trident/std/crypto/poseidon.tri

module std.crypto.poseidon

// Poseidon hash function for non-Triton targets.
//
// Poseidon is an algebraic hash function designed for ZK-friendly
// arithmetic circuits. It operates natively over field elements,
// making it efficient in SNARKs/STARKs that use different fields
// than Triton's Tip5.
//
// This module provides placeholder implementations. The actual
// Poseidon permutation requires:
//   - Round constants (specific to the field and security level)
//   - An S-box (typically x^alpha where alpha depends on the field)
//   - An MDS matrix for the linear mixing layer
//   - A specific number of full and partial rounds
//
// For the Goldilocks field (p = 2^64 - 2^32 + 1), a concrete
// Poseidon instantiation would use alpha = 7, with 8 full rounds
// and 22 partial rounds for a state width of 12.
//
// The stubs below use simple arithmetic combinations to show the
// API shape. They are NOT cryptographically secure and must be
// replaced with a proper Poseidon implementation for production use.
// ---------------------------------------------------------------------------
// Internal helpers: S-box and round simulation
// ---------------------------------------------------------------------------
// Placeholder S-box: x^5 over the field.
// The real Poseidon S-box for Goldilocks would be x^7,
// but x^5 illustrates the pattern with fewer multiplications.
fn sbox(x: Field) -> Field {
    let x2: Field = x * x
    let x4: Field = x2 * x2
    let x5: Field = x4 * x
    x5
}

// Placeholder linear mixing for a 2-element state.
// A real MDS matrix would be a 2x2 Cauchy matrix.
// This uses a simple invertible linear map instead.
fn mix2(a: Field, b: Field) -> (Field, Field) {
    // [ 2  1 ] [ a ]   =   [ 2a + b  ]
    // [ 1  3 ] [ b ]       [ a + 3b  ]
    let new_a: Field = a + a + b
    let new_b: Field = a + b + b + b
    (new_a, new_b)
}

// Placeholder linear mixing for a 4-element state.
// A real MDS matrix would be a 4x4 Cauchy matrix.
fn mix4(a: Field, b: Field, c: Field, d: Field) -> (Field, Field, Field, Field) {
    // Simple circulant-like mixing: each output depends on all inputs.
    let new_a: Field = a + a + b + c + d
    let new_b: Field = a + b + b + c + d
    let new_c: Field = a + b + c + c + d
    let new_d: Field = a + b + c + d + d
    (new_a, new_b, new_c, new_d)
}

// One full round for a 2-element state: S-box + mix.
fn round2(a: Field, b: Field, rc0: Field, rc1: Field) -> (Field, Field) {
    // Add round constants
    let a1: Field = a + rc0
    let b1: Field = b + rc1
    // Apply S-box
    let a2: Field = sbox(a1)
    let b2: Field = sbox(b1)
    // Linear mixing
    mix2(a2, b2)
}

// One full round for a 4-element state: S-box + mix.
fn round4(
    a: Field,
    b: Field,
    c: Field,
    d: Field,
    rc0: Field,
    rc1: Field,
    rc2: Field,
    rc3: Field
) -> (Field, Field, Field, Field) {
    // Add round constants
    let a1: Field = a + rc0
    let b1: Field = b + rc1
    let c1: Field = c + rc2
    let d1: Field = d + rc3
    // Apply S-box
    let a2: Field = sbox(a1)
    let b2: Field = sbox(b1)
    let c2: Field = sbox(c1)
    let d2: Field = sbox(d1)
    // Linear mixing
    mix4(a2, b2, c2, d2)
}

// ---------------------------------------------------------------------------
// Poseidon hash: 2 inputs -> 1 output
// ---------------------------------------------------------------------------
// Poseidon hash for 2 field elements -> 1 field element.
//
// PLACEHOLDER: Uses 4 rounds with dummy round constants.
// A real implementation needs the correct number of full/partial rounds
// and cryptographically derived round constants for the target field.
//
// The capacity element is initialized to 0 (not absorbed),
// and the output is taken from the first rate element.
pub fn hash2(a: Field, b: Field) -> Field {
    // Initial state: [a, b] (rate = 2, capacity is implicit)
    // In a real Poseidon with rate 2 and capacity 1, the state would be
    // [a, b, 0] and we'd run the permutation on all 3 elements.
    // For this 2-element stub we show the round structure.
    // Round 1: dummy round constants derived from small primes
    let (s0, s1) = round2(a, b, 3, 7)
    // Round 2
    let (s2, s3) = round2(s0, s1, 11, 13)
    // Round 3
    let (s4, s5) = round2(s2, s3, 17, 19)
    // Round 4
    let (s6, _s7) = round2(s4, s5, 23, 29)
    // Output: first element of final state
    s6
}

// ---------------------------------------------------------------------------
// Poseidon hash: 4 inputs -> 1 output
// ---------------------------------------------------------------------------
// Poseidon hash for 4 field elements -> 1 field element.
//
// PLACEHOLDER: Uses 4 rounds with dummy round constants.
// A real implementation would use rate = 4, capacity >= 1.
pub fn hash4(a: Field, b: Field, c: Field, d: Field) -> Field {
    // Round 1
    let (s0, s1, s2, s3) = round4(a, b, c, d, 3, 7, 11, 13)
    // Round 2
    let (s4, s5, s6, s7) = round4(s0, s1, s2, s3, 17, 19, 23, 29)
    // Round 3
    let (s8, s9, s10, s11) = round4(s4, s5, s6, s7, 31, 37, 41, 43)
    // Round 4
    let (s12, _s13, _s14, _s15) = round4(s8, s9, s10, s11, 47, 53, 59, 61)
    // Output: first element of final state
    s12
}

// ---------------------------------------------------------------------------
// Poseidon sponge: absorb-squeeze API
// ---------------------------------------------------------------------------
// Hash a single field element (domain separation for single-element inputs).
// Uses hash2 with the second input set to a domain tag.
pub fn hash1(a: Field) -> Field {
    // Domain tag: 1 indicates single-element input
    hash2(a, 1)
}

// Hash 3 field elements by chaining hash2 calls.
// Sponge-style: absorb 2, squeeze, absorb 1 more, squeeze.
pub fn hash3(a: Field, b: Field, c: Field) -> Field {
    let h1: Field = hash2(a, b)
    hash2(h1, c)
}