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

p3_bn254/
poseidon2.rs
