p3_poseidon1/

lib.rs

//! The Poseidon1 permutation.
//!
//! # Overview
//!
//! Poseidon1 is a cryptographic hash function designed for efficient verification inside
//! zero-knowledge proof systems (SNARKs, STARKs, etc). It operates natively over
//! prime fields, avoiding the overhead of bit-decomposition required by traditional hashes
//! like SHA-256 or Keccak.
//!
//! This crate provides [`Poseidon1`], an optimized implementation that uses the **sparse matrix
//! decomposition** from the original paper (Appendix B) to factor the dense MDS matrix into
//! sparse components. This reduces the cost of each partial round from O(t^2) to O(t), where
//! t = `WIDTH` is the state width.
//!
//! # Round Structure
//!
//! Poseidon1 alternates between two types of rounds:
//!
//! ```text
//!   ┌──────────────────┐   ┌───────────────────┐   ┌──────────────────┐
//!   │ RF/2 full rounds │ → │ RP partial rounds │ → │ RF/2 full rounds │
//!   │ (S-box on all)   │   │ (S-box on s[0])   │   │ (S-box on all)   │
//!   └──────────────────┘   └───────────────────┘   └──────────────────┘
//! ```
//!
//! - **Full rounds** (RF total, split equally): apply the S-box x^D to every state element,
//!   then multiply by the dense MDS matrix. These provide resistance against statistical
//!   attacks (differential, linear, rebound).
//!
//! - **Partial rounds** (RP total): apply the S-box only to `state[0]`, then multiply by a
//!   **sparse** matrix derived from the MDS factorization. These efficiently increase the
//!   algebraic degree to resist interpolation and Gröbner basis attacks.
//!
//! Each round follows the sequence: **AddRoundConstants → S-box → MixLayer**.
//!
//! # References
//!
//! - Grassi et al., "Poseidon1: A New Hash Function for Zero-Knowledge Proof Systems",
//!   USENIX Security 2021. <https://eprint.iacr.org/2019/458>
//! - HorizenLabs reference implementation: <https://github.com/HorizenLabs/poseidon2>

#![no_std]

extern crate alloc;

pub mod external;
pub mod generic;
pub mod internal;
pub mod utils;

use alloc::vec::Vec;
use core::marker::PhantomData;

pub use external::*;
pub use generic::*;
pub use internal::*;
use p3_field::{Algebra, InjectiveMonomial, PrimeField};
use p3_symmetric::{CryptographicPermutation, Permutation};
use rand::distr::{Distribution, StandardUniform};
use rand::{Rng, RngExt};

/// Raw Poseidon1 parameters before the sparse matrix optimization.
///
/// These are the "textbook" parameters as generated by the Grain LFSR or any
/// other parameter-generation script. They are transformed into the optimized
/// sparse form used at runtime.
///
/// # Round Constant Layout
///
/// Constants are stored in a flat array with three consecutive sections:
///
/// ```text
///   ┌──────────────┬──────────────────┬──────────────────┐
///   │ initial full  │  partial rounds  │  terminal full   │
///   │ (RF/2 items)  │  (RP items)      │  (RF/2 items)    │
///   └──────────────┴──────────────────┴──────────────────┘
/// ```
///
/// Each entry is a WIDTH-sized vector (one constant per state element per round).
#[derive(Clone, Debug)]
pub struct Poseidon1Constants<F, const WIDTH: usize> {
    /// Total number of full rounds (split equally between initial and terminal).
    pub rounds_f: usize,

    /// Number of partial rounds.
    pub rounds_p: usize,

    /// First column of the circulant MDS matrix, stored as signed integers.
    ///
    /// This matches the representation used by the MDS crate, so concrete
    /// fields can pass their verified constants directly without duplication.
    /// During initialization, the dense form is expanded once for the sparse
    /// decomposition, then discarded.
    pub mds_circ_col: [i64; WIDTH],

    /// Round constants, one WIDTH-sized vector per round.
    ///
    /// Total length = rounds_f + rounds_p.
    pub round_constants: Vec<[F; WIDTH]>,
}

impl<F: PrimeField, const WIDTH: usize> Poseidon1Constants<F, WIDTH> {
    /// Compute the optimized sparse-form constants from these raw parameters.
    ///
    /// This performs two transformations:
    ///
    /// 1. **Sparse matrix decomposition**: factors the dense MDS matrix into one
    ///    dense transition matrix and several sparse matrices. Each sparse matrix
    ///    is parameterized by two vectors of length WIDTH-1.
    ///
    /// 2. **Round constant compression**: via backward substitution through the
    ///    inverse MDS matrix, reduces each partial round's full constant vector
    ///    to a single scalar, except for the first partial round.
    pub fn to_optimized(
        &self,
    ) -> (
        FullRoundConstants<F, WIDTH>,
        PartialRoundConstants<F, WIDTH>,
    ) {
        let half_f = self.rounds_f / 2;
        let rounds_p = self.rounds_p;

        // Split the flat round-constant array into three sections.
        //
        // Layout: [initial_full (RF/2) | partial (RP) | terminal_full (RF/2)]
        let initial_rc: Vec<[F; WIDTH]> = self.round_constants[..half_f].to_vec();
        let partial_rc: Vec<[F; WIDTH]> = self.round_constants[half_f..half_f + rounds_p].to_vec();
        let terminal_rc: Vec<[F; WIDTH]> = self.round_constants[half_f + rounds_p..].to_vec();

        // Expand circulant first column into dense MDS matrix for the sparse decomposition.
        let mds: [[F; WIDTH]; WIDTH] = utils::circulant_to_dense(&self.mds_circ_col);

        // Compress round constants and factor MDS into sparse matrices.
        let (first_round_constants, opt_partial_rc, m_i, sparse_v, sparse_first_row) =
            utils::compute_optimized_constants::<F, WIDTH>(&mds, rounds_p, &partial_rc);

        // Compute textbook path constants (forward constant substitution).
        let (textbook_scalar_constants, textbook_residual) =
            utils::forward_constant_substitution::<F, WIDTH>(&mds, &partial_rc);

        let full_constants = FullRoundConstants {
            initial: initial_rc,
            terminal: terminal_rc,
            dense_mds: mds,
        };

        let partial_constants = PartialRoundConstants {
            first_round_constants,
            m_i,
            sparse_first_row,
            v: sparse_v,
            round_constants: opt_partial_rc,
            textbook_scalar_constants,
            textbook_residual,
        };

        (full_constants, partial_constants)
    }
}

/// The optimized Poseidon1 permutation.
///
/// Holds the pre-computed full and partial round layers and applies them
/// in sequence: initial full rounds, then partial rounds, then terminal
/// full rounds.
#[derive(Clone, Debug)]
pub struct Poseidon1<F, FullRoundPerm, PartialRoundPerm, const WIDTH: usize, const D: u64> {
    /// The initial and terminal full round layers.
    ///
    /// Full rounds apply the S-box to every state element, then multiply
    /// by the dense MDS matrix.
    full_round_layer: FullRoundPerm,

    /// The partial round layer using sparse matrix decomposition.
    ///
    /// Partial rounds apply the S-box only to the first state element,
    /// then multiply by a sparse matrix in O(t) operations.
    partial_round_layer: PartialRoundPerm,

    _phantom: PhantomData<F>,
}

impl<F, FullRoundPerm, PartialRoundPerm, const WIDTH: usize, const D: u64>
    Poseidon1<F, FullRoundPerm, PartialRoundPerm, WIDTH, D>
where
    F: PrimeField,
    FullRoundPerm: FullRoundLayerConstructor<F, WIDTH>,
    PartialRoundPerm: PartialRoundLayerConstructor<F, WIDTH>,
{
    /// Create a new optimized Poseidon1 from raw constants.
    ///
    /// Internally computes the sparse matrix decomposition and the
    /// optimized round constants. This is the typical entry point.
    pub fn new(raw: &Poseidon1Constants<F, WIDTH>) -> Self {
        // Compile-time structural checks from the Poseidon specification.
        const {
            // Section 2.3: POSEIDONπ takes inputs of t ≥ 2 words.
            assert!(
                WIDTH >= 2,
                "Poseidon1 requires WIDTH >= 2 (paper Section 2.3: t >= 2)"
            );
            // Section 2.3: S-box(x) = x^α where α ≥ 3.
            assert!(
                D >= 3,
                "Poseidon1 requires D >= 3 (paper Section 2.3: alpha >= 3)"
            );
        }

        // Section 5.5.1 / Eq.(2): RF ≥ 6 for statistical attack resistance.
        assert!(
            raw.rounds_f >= 6,
            "Poseidon1 requires rounds_f >= 6 (paper Section 5.5.1, Eq.(2): statistical attacks)"
        );
        // Round constants layout: [initial_full (RF/2) | partial (RP) | terminal_full (RF/2)].
        assert_eq!(
            raw.round_constants.len(),
            raw.rounds_f + raw.rounds_p,
            "Poseidon1 round_constants length must equal rounds_f + rounds_p"
        );

        let (full_constants, partial_constants) = raw.to_optimized();
        Self::new_from_precomputed(full_constants, partial_constants)
    }

    /// Create a new optimized Poseidon1 from pre-computed constants.
    ///
    /// Use this when the sparse matrix decomposition has already been performed
    /// (e.g., constants loaded from a file or hardcoded).
    fn new_from_precomputed(
        full_constants: FullRoundConstants<F, WIDTH>,
        partial_constants: PartialRoundConstants<F, WIDTH>,
    ) -> Self {
        Self {
            full_round_layer: FullRoundPerm::new_from_constants(full_constants),
            partial_round_layer: PartialRoundPerm::new_from_constants(partial_constants),
            _phantom: PhantomData,
        }
    }

    /// Create a new Poseidon1 with random round constants and a given MDS permutation.
    ///
    /// Builds the dense MDS matrix by applying the permutation to unit vectors,
    /// generates random round constants, and computes the sparse matrix decomposition.
    ///
    /// Primarily useful for testing.
    pub fn new_from_rng(
        half_num_full_rounds: usize,
        num_partial_rounds: usize,
        mds: &impl Permutation<[F; WIDTH]>,
        rng: &mut impl Rng,
    ) -> Self
    where
        StandardUniform: Distribution<F>,
    {
        // Compile-time structural checks from the Poseidon specification.
        const {
            // Section 2.3: POSEIDONπ takes inputs of t ≥ 2 words.
            assert!(
                WIDTH >= 2,
                "Poseidon1 requires WIDTH >= 2 (paper Section 2.3: t >= 2)"
            );
            // Section 2.3: S-box(x) = x^α where α ≥ 3.
            assert!(
                D >= 3,
                "Poseidon1 requires D >= 3 (paper Section 2.3: alpha >= 3)"
            );
        }

        let rounds_f = 2 * half_num_full_rounds;

        // Section 5.5.1 / Eq.(2): RF ≥ 6 for statistical attack resistance.
        assert!(
            rounds_f >= 6,
            "Poseidon1 requires rounds_f >= 6 (paper Section 5.5.1, Eq.(2): statistical attacks)"
        );

        let num_rounds = rounds_f + num_partial_rounds;

        // Generate random round constants.
        let round_constants: Vec<[F; WIDTH]> = (0..num_rounds)
            .map(|_| core::array::from_fn(|_| rng.sample(StandardUniform)))
            .collect();

        // Build dense MDS by applying the permutation to unit vectors.
        let columns: [[F; WIDTH]; WIDTH] = core::array::from_fn(|j| {
            let mut e = [F::ZERO; WIDTH];
            e[j] = F::ONE;
            mds.permute(e)
        });
        // Transpose to row-major format.
        let dense_mds: [[F; WIDTH]; WIDTH] =
            core::array::from_fn(|i| core::array::from_fn(|j| columns[j][i]));

        // Split round constants into three sections.
        let half_f = half_num_full_rounds;
        let initial_rc = round_constants[..half_f].to_vec();
        let partial_rc = round_constants[half_f..half_f + num_partial_rounds].to_vec();
        let terminal_rc = round_constants[half_f + num_partial_rounds..].to_vec();

        // Compute optimized sparse constants.
        let (first_round_constants, opt_partial_rc, m_i, sparse_v, sparse_first_row) =
            utils::compute_optimized_constants::<F, WIDTH>(
                &dense_mds,
                num_partial_rounds,
                &partial_rc,
            );

        // Compute textbook path constants (forward constant substitution).
        let (textbook_scalar_constants, textbook_residual) =
            utils::forward_constant_substitution::<F, WIDTH>(&dense_mds, &partial_rc);

        let full_constants = FullRoundConstants {
            initial: initial_rc,
            terminal: terminal_rc,
            dense_mds,
        };

        let partial_constants = PartialRoundConstants {
            first_round_constants,
            m_i,
            sparse_first_row,
            v: sparse_v,
            round_constants: opt_partial_rc,
            textbook_scalar_constants,
            textbook_residual,
        };

        Self::new_from_precomputed(full_constants, partial_constants)
    }
}

impl<F, A, FullRoundPerm, PartialRoundPerm, const WIDTH: usize, const D: u64>
    Permutation<[A; WIDTH]> for Poseidon1<F, FullRoundPerm, PartialRoundPerm, WIDTH, D>
where
    F: PrimeField + InjectiveMonomial<D>,
    A: Algebra<F> + Sync + InjectiveMonomial<D>,
    FullRoundPerm: FullRoundLayer<A, WIDTH, D>,
    PartialRoundPerm: PartialRoundLayer<A, WIDTH, D>,
{
    fn permute_mut(&self, state: &mut [A; WIDTH]) {
        // RF/2 full rounds (S-box on all elements + dense MDS).
        self.full_round_layer.permute_state_initial(state);
        // RP partial rounds (S-box on state[0] only + sparse matrix).
        self.partial_round_layer.permute_state(state);
        // RF/2 full rounds (S-box on all elements + dense MDS).
        self.full_round_layer.permute_state_terminal(state);
    }
}

impl<F, A, FullRoundPerm, PartialRoundPerm, const WIDTH: usize, const D: u64>
    CryptographicPermutation<[A; WIDTH]> for Poseidon1<F, FullRoundPerm, PartialRoundPerm, WIDTH, D>
where
    F: PrimeField + InjectiveMonomial<D>,
    A: Algebra<F> + Sync + InjectiveMonomial<D>,
    FullRoundPerm: FullRoundLayer<A, WIDTH, D>,
    PartialRoundPerm: PartialRoundLayer<A, WIDTH, D>,
{
}