p3_bn254/
poseidon2.rs

1//! Diffusion matrix for Bn254
2//!
3//! Reference: https://github.com/HorizenLabs/poseidon2/blob/main/plain_implementations/src/poseidon2/poseidon2_instance_bn256.rs
4
5extern crate alloc;
6
7use alloc::vec::Vec;
8
9use p3_field::PrimeCharacteristicRing;
10use p3_poseidon2::{
11    ExternalLayer, ExternalLayerConstants, ExternalLayerConstructor, HLMDSMat4, InternalLayer,
12    InternalLayerConstructor, Poseidon2, add_rc_and_sbox_generic, external_initial_permute_state,
13    external_terminal_permute_state, internal_permute_state,
14};
15
16use crate::Bn254;
17
18/// Degree of the chosen permutation polynomial for BN254, used as the Poseidon2 S-Box.
19///
20/// As p - 1 is divisible by 2 and 3 the smallest choice for a degree D satisfying gcd(p - 1, D) = 1 is 5.
21const BN254_S_BOX_DEGREE: u64 = 5;
22
23/// An implementation of the Poseidon2 hash function for the Bn254Fr field.
24///
25/// It acts on arrays of the form `[Bn254Fr; WIDTH]`.
26pub type Poseidon2Bn254<const WIDTH: usize> = Poseidon2<
27    Bn254,
28    Poseidon2ExternalLayerBn254<WIDTH>,
29    Poseidon2InternalLayerBn254,
30    WIDTH,
31    BN254_S_BOX_DEGREE,
32>;
33
34/// Currently we only support a single width for Poseidon2 BN254.
35const BN254_WIDTH: usize = 3;
36
37#[derive(Debug, Clone, Default)]
38pub struct Poseidon2InternalLayerBn254 {
39    internal_constants: Vec<Bn254>,
40}
41
42impl InternalLayerConstructor<Bn254> for Poseidon2InternalLayerBn254 {
43    fn new_from_constants(internal_constants: Vec<Bn254>) -> Self {
44        Self { internal_constants }
45    }
46}
47
48/// A faster version of `matmul_internal` making use of the fact that
49/// the internal matrix is equal to:
50/// ```ignore
51///                             [2, 1, 1]
52///     1 + Diag([1, 1, 2]) =   [1, 2, 1]
53///                             [1, 1, 3]
54/// ```
55fn bn254_matmul_internal(state: &mut [Bn254; 3]) {
56    // We bracket in this way as the s-box is applied to state[0] so this lets us
57    // begin this computation before the s-box finishes.
58    let sum = state[0] + (state[1] + state[2]);
59
60    state[0] += sum;
61    state[1] += sum;
62    state[2] = state[2].double() + sum;
63}
64
65impl InternalLayer<Bn254, BN254_WIDTH, BN254_S_BOX_DEGREE> for Poseidon2InternalLayerBn254 {
66    /// Perform the internal layers of the Poseidon2 permutation on the given state.
67    fn permute_state(&self, state: &mut [Bn254; BN254_WIDTH]) {
68        internal_permute_state(state, bn254_matmul_internal, &self.internal_constants);
69    }
70}
71
72pub type Poseidon2ExternalLayerBn254<const WIDTH: usize> = ExternalLayerConstants<Bn254, WIDTH>;
73
74impl<const WIDTH: usize> ExternalLayerConstructor<Bn254, WIDTH>
75    for Poseidon2ExternalLayerBn254<WIDTH>
76{
77    fn new_from_constants(external_constants: Self) -> Self {
78        external_constants
79    }
80}
81
82impl<const WIDTH: usize> ExternalLayer<Bn254, WIDTH, BN254_S_BOX_DEGREE>
83    for Poseidon2ExternalLayerBn254<WIDTH>
84{
85    /// Perform the initial external layers of the Poseidon2 permutation on the given state.
86    fn permute_state_initial(&self, state: &mut [Bn254; WIDTH]) {
87        external_initial_permute_state(
88            state,
89            self.get_initial_constants(),
90            add_rc_and_sbox_generic,
91            &HLMDSMat4,
92        );
93    }
94
95    /// Perform the terminal external layers of the Poseidon2 permutation on the given state.
96    fn permute_state_terminal(&self, state: &mut [Bn254; WIDTH]) {
97        external_terminal_permute_state(
98            state,
99            self.get_terminal_constants(),
100            add_rc_and_sbox_generic,
101            &HLMDSMat4,
102        );
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use num_bigint::BigUint;
109    use p3_field::PrimeField;
110    use p3_poseidon2::ExternalLayerConstants;
111    use p3_symmetric::Permutation;
112    use rand::rngs::SmallRng;
113    use rand::{Rng, SeedableRng};
114    use zkhash::ark_ff::{BigInteger, PrimeField as ark_PrimeField};
115    use zkhash::fields::bn256::FpBN256 as ark_FpBN256;
116    use zkhash::poseidon2::poseidon2::Poseidon2 as Poseidon2Ref;
117    use zkhash::poseidon2::poseidon2_instance_bn256::{POSEIDON2_BN256_PARAMS, RC3};
118
119    use super::*;
120    use crate::BN254_MONTY_R_SQ;
121    use crate::helpers::monty_mul;
122
123    fn bn254_from_ark_ff(input: ark_FpBN256) -> Bn254 {
124        let mut full_bytes = [0; 32];
125        let bytes = input.into_bigint().to_bytes_le();
126        full_bytes[..bytes.len()].copy_from_slice(&bytes);
127
128        let value = Bn254::from_bytes_monty(&full_bytes);
129
130        value.map_or_else(
131            || panic!("Invalid field element"),
132            |field_elem| {
133                // From bytes does not convert into Monty form.
134                // Hence we need to do that ourselves.
135                let monty_form = monty_mul(BN254_MONTY_R_SQ, field_elem.value);
136                Bn254::new_monty(monty_form)
137            },
138        )
139    }
140
141    fn ark_ff_from_bn254(input: Bn254) -> ark_FpBN256 {
142        // We can't just use `input.into_bytes()` as we need to first convert out of MONTY form.
143        // Going via `BigUint` is a little unnecessary but is sufficient for our purposes.
144        let bigint = BigUint::from_bytes_le(&input.as_canonical_biguint().to_bytes_le());
145        ark_FpBN256::from(bigint)
146    }
147
148    #[test]
149    fn test_poseidon2_bn254() {
150        const WIDTH: usize = 3;
151        const ROUNDS_F: usize = 8;
152        const ROUNDS_P: usize = 56;
153
154        type F = Bn254;
155
156        let mut rng = SmallRng::seed_from_u64(1);
157
158        // Poiseidon2 reference implementation from zkhash repo.
159        let poseidon2_ref = Poseidon2Ref::new(&POSEIDON2_BN256_PARAMS);
160
161        // Copy over round constants from zkhash.
162        let mut round_constants: Vec<[F; WIDTH]> = RC3
163            .iter()
164            .map(|vec| {
165                vec.iter()
166                    .copied()
167                    .map(bn254_from_ark_ff)
168                    .collect::<Vec<_>>()
169                    .try_into()
170                    .unwrap()
171            })
172            .collect();
173
174        let internal_start = ROUNDS_F / 2;
175        let internal_end = (ROUNDS_F / 2) + ROUNDS_P;
176        let internal_round_constants = round_constants
177            .drain(internal_start..internal_end)
178            .map(|vec| vec[0])
179            .collect::<Vec<_>>();
180        let external_round_constants = ExternalLayerConstants::new(
181            round_constants[..(ROUNDS_F / 2)].to_vec(),
182            round_constants[(ROUNDS_F / 2)..].to_vec(),
183        );
184        // Our Poseidon2 implementation.
185        let poseidon2 = Poseidon2Bn254::new(external_round_constants, internal_round_constants);
186
187        // Generate random input and convert to both field formats.
188        let input = rng.random::<[F; WIDTH]>();
189        let input_ark_ff = input.map(ark_ff_from_bn254);
190
191        // Run reference implementation.
192        let output_ref: [ark_FpBN256; WIDTH] =
193            poseidon2_ref.permutation(&input_ark_ff).try_into().unwrap();
194        let expected: [F; WIDTH] = output_ref.map(bn254_from_ark_ff);
195
196        // Run our implementation.
197        let mut output = input;
198        poseidon2.permute_mut(&mut output);
199
200        assert_eq!(output, expected);
201    }
202}
p3_bn254/poseidon2.rs

p3_bn254/
poseidon2.rs
