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

jf_poseidon2/
lib.rs
