jf_poseidon2/
lib.rs

1// Copyright (c) 2024 Espresso Systems (espressosys.com)
2// This file is part of the Jellyfish library.
3
4// You should have received a copy of the MIT License
5// along with the Jellyfish library. If not, see <https://mit-license.org/>.
6
7//! The Poseidon2 permutation.
8//!
9//! # Available instances
10//! We have implemented Poseidon2 instances: (field, state_size)
11//! - (Bls12_381::Fr, 2), (Bls12_381::Fr, 3)
12//! - (Bn254::Fr, 3)
13//!
14//! This implementation was based upon the following resources:
15//! - https://github.com/HorizenLabs/poseidon2/blob/main/plain_implementations/src/poseidon2/poseidon2.rs
16//! - https://eprint.iacr.org/2023/323.pdf
17//! - https://github.com/Plonky3/Plonky3/blob/main/poseidon2/
18
19#![no_std]
20#![deny(missing_docs)]
21
22use ark_ff::PrimeField;
23use ark_std::{borrow::ToOwned, marker::PhantomData, string::String};
24use displaydoc::Display;
25
26pub mod constants;
27pub mod crhf;
28mod external;
29mod internal;
30pub mod sponge;
31
32/// Parameters required for a Poseidon2 permutation instance.
33///
34/// # Generic parameters
35/// - `F`: field choice
36/// - `T`: state size = rate + capacity, `T` is made generic for easy trait
37///   bound on `permute<F,T>(input: [F; N])` and type safety on `external_rc()`
38///   return type.
39pub trait Poseidon2Params<F: PrimeField, const T: usize>: Clone {
40    /// t: state size = rate + capacity
41    const T: usize = T;
42    /// d: sbox degree
43    const D: usize;
44    /// round_F: number of external rounds (incl. initial and terminal)
45    /// round_F = 2 * round_f
46    const EXT_ROUNDS: usize;
47    /// round_P: number of internal rounds
48    const INT_ROUNDS: usize;
49
50    /// round constants for all external rounds
51    fn external_rc() -> &'static [[F; T]];
52    /// round constants for internal rounds
53    fn internal_rc() -> &'static [F];
54    /// diffusion (diagonal) matrix minus one used in internal rounds
55    fn internal_mat_diag_m_1() -> &'static [F; T];
56
57    /// A default sanity check on the parameters and constant getters
58    ///
59    /// State size only supports: 2, 3, 4, 8, 12, 16, 20, 24 for now
60    /// S-box size only supports: 3, 5, 7, 11
61    ///
62    /// # Round constants
63    /// Rust doesn't permit generic param to be used in const operations, thus
64    /// leveraging type system to ensure sanity such as `const INT_RC: &'static
65    /// [F; Self::INT_ROUNDS]` is not allowed.
66    fn sanity_check() -> bool {
67        let ext_rc = Self::external_rc();
68        let int_rc = Self::internal_rc();
69
70        // TODO: consider adding more security-related check, incl. number of internal
71        // rounds in terms of field size to achieve 128-bit security. see
72        // `poseidon2_round_numbers_128` in plonky3, we skip for now as GCD is not
73        // implemented in arkworks, and params are trusted.
74        [2, 3, 4, 8, 12, 16, 20, 24].contains(&T)
75            && [3, 5, 7, 11].contains(&Self::D)
76            && ext_rc.len() == Self::EXT_ROUNDS
77            && int_rc.len() == Self::INT_ROUNDS
78            && Self::EXT_ROUNDS % 2 == 0
79    }
80}
81
82/// A Poseidon2 permutation family <https://eprint.iacr.org/2023/323>
83pub struct Poseidon2<F: PrimeField>(PhantomData<F>);
84
85impl<F: PrimeField> Poseidon2<F> {
86    /// Apply Poseidon2 permutation on `input` and return the permuted result
87    pub fn permute<P, const T: usize>(input: &[F; T]) -> [F; T]
88    where
89        P: Poseidon2Params<F, T>,
90    {
91        let mut input = input.to_owned();
92        Self::permute_mut::<P, T>(&mut input);
93        input
94    }
95
96    /// Apply Poseidon2 permutation on `input` in place
97    pub fn permute_mut<P, const T: usize>(input: &mut [F; T])
98    where
99        P: Poseidon2Params<F, T>,
100    {
101        assert!(P::sanity_check(), "Unexpected: Invalid Poseidon2 param!");
102        // M_e * x
103        external::matmul_external(input);
104
105        // Initial external rounds (first EXT_ROUNDS/2 rounds)
106        let ext_rc = P::external_rc();
107        for rc in ext_rc.iter().take(P::EXT_ROUNDS / 2) {
108            external::permute_state(input, rc, P::D);
109        }
110
111        // Internal rounds
112        let int_rc = P::internal_rc();
113        let mat_diag_minus_1 = P::internal_mat_diag_m_1();
114        for rc in int_rc.iter() {
115            internal::permute_state(input, *rc, P::D, mat_diag_minus_1);
116        }
117
118        // Terminal external rounds (second EXT_ROUNDS/2 rounds)
119        for rc in ext_rc.iter().take(P::EXT_ROUNDS).skip(P::EXT_ROUNDS / 2) {
120            external::permute_state(input, rc, P::D);
121        }
122    }
123}
124
125// @credit: `add_rc_and_sbox_generic()` in plonky3
126/// add RCs to the entire state
127#[inline(always)]
128pub(crate) fn add_rcs<F: PrimeField, const T: usize>(state: &mut [F; T], rc: &[F; T]) {
129    for i in 0..T {
130        state[i] += rc[i];
131    }
132}
133
134/// `s -> s^d`
135#[inline(always)]
136pub(crate) fn s_box<F: PrimeField>(val: &mut F, d: usize) {
137    if d == 5 {
138        // Perform unrolled computation for val^5, faster
139        let original = *val;
140        val.square_in_place();
141        val.square_in_place();
142        *val *= &original;
143    } else {
144        *val = val.pow([d as u64]);
145    }
146}
147
148/// Poseidon2 Error type
149#[derive(Debug, Clone, Display, Eq, PartialEq)]
150pub enum Poseidon2Error {
151    /// Bad parameter: {0}
152    ParamErr(String),
153}
154impl ark_std::error::Error for Poseidon2Error {}
jf_poseidon2/lib.rs

jf_poseidon2/
lib.rs
