Skip to main content

p3_mersenne_31/
poseidon2.rs

1//! Implementation of Poseidon2, see: `<https://eprint.iacr.org/2023/323>`
2//!
3//! For the diffusion matrix, 1 + Diag(V), we perform a search to find an optimized
4//! vector V composed of elements with efficient multiplication algorithms in AVX2/AVX512/NEON.
5//!
6//! This leads to using small values (e.g. 1, 2) where multiplication is implemented using addition
7//! and powers of 2 where multiplication is implemented using shifts.
8//! Additionally, for technical reasons, having the first entry be -2 is useful.
9//!
10//! Optimized Diagonal for Mersenne31 width 16:
11//! [-2, 2^0, 2, 4, 8, 16, 32, 64, 2^7, 2^8, 2^10, 2^12, 2^13, 2^14, 2^15, 2^16]
12//! Optimized Diagonal for Mersenne31 width 24:
13//! [-2, 2^0, 2, 4, 8, 16, 32, 64, 2^7, 2^8, 2^9, 2^10, 2^11, 2^12, 2^13, 2^14, 2^15, 2^16, 2^17, 2^18, 2^19, 2^20, 2^21, 2^22]
14//! See poseidon2\src\diffusion.rs for information on how to double check these matrices in Sage.
15
16use p3_field::PrimeCharacteristicRing;
17use p3_poseidon2::{
18    ExternalLayer, ExternalLayerConstants, GenericPoseidon2LinearLayers, InternalLayer, MDSMat4,
19    Poseidon2, add_rc_and_sbox_generic, external_initial_permute_state,
20    external_terminal_permute_state, internal_permute_state,
21};
22
23use crate::{
24    Mersenne31, Poseidon2ExternalLayerMersenne31, Poseidon2InternalLayerMersenne31, from_u62,
25};
26
27/// S-box degree for Mersenne31 Poseidon2.
28///
29/// Since `p - 1 = 2 × 3^2 × 7 × 11 × 31 × 151 × 331`, both 3 and 4 share factors
30/// with `p - 1`. The smallest valid exponent satisfying `gcd(α, p - 1) = 1` is 5.
31pub const MERSENNE31_S_BOX_DEGREE: u64 = 5;
32
33/// Number of full rounds per half for Mersenne31 Poseidon2 (`RF / 2`).
34///
35/// The total number of full rounds is `RF = 8` (4 beginning + 4 ending).
36/// Follows the Poseidon2 paper's security analysis with a +2 RF margin.
37pub const MERSENNE31_POSEIDON2_HALF_FULL_ROUNDS: usize = 4;
38
39/// Number of partial rounds for Mersenne31 Poseidon2 (width 16).
40///
41/// Derived from the Gröbner basis bound in the Poseidon2 paper (Eq. 1, R_GB term 3):
42///
43///   R_GB ≥ t − 7 + log_α(2) · min{κ/(t+1), log_2(p)/2}
44///        = 9 + 0.4307 · min{7.53, 15.5} = 12.243
45///
46/// With the +7.5% security margin: ⌈1.075 × 12.243⌉ = 14.
47pub const MERSENNE31_POSEIDON2_PARTIAL_ROUNDS_16: usize = 14;
48
49/// Number of partial rounds for Mersenne31 Poseidon2 (width 24).
50///
51/// Same Gröbner basis bound as width 16:
52///
53///   R_GB ≥ 17 + 0.4307 · min{5.12, 15.5} = 19.205
54///
55/// With the +7.5% security margin: ⌈1.075 × 19.205⌉ = 21.
56///
57/// The official round number script yields R_P = 22 for this configuration
58/// (matching the Grain LFSR parameters used to generate the round constants).
59pub const MERSENNE31_POSEIDON2_PARTIAL_ROUNDS_24: usize = 22;
60
61/// An implementation of the Poseidon2 hash function specialised to run on the current architecture.
62///
63/// It acts on arrays of the form either `[Mersenne31::Packing; WIDTH]` or `[Mersenne31; WIDTH]`. For speed purposes,
64/// wherever possible, input arrays should of the form `[Mersenne31::Packing; WIDTH]`.
65pub type Poseidon2Mersenne31<const WIDTH: usize> = Poseidon2<
66    Mersenne31,
67    Poseidon2ExternalLayerMersenne31<WIDTH>,
68    Poseidon2InternalLayerMersenne31,
69    WIDTH,
70    MERSENNE31_S_BOX_DEGREE,
71>;
72
73/// Round constants for width-16 Poseidon2 on Mersenne31.
74///
75/// Generated by the Grain LFSR with parameters:
76///     field_type=1, alpha=5 (exp_flag=0), n=31, t=16, R_F=8, R_P=14
77///
78/// Generated by `poseidon2/generate_constants.py --field mersenne31 --width 16`.
79///
80/// Layout: external_initial (4 rounds × 16 elements).
81pub const MERSENNE31_POSEIDON2_RC_16_EXTERNAL_INITIAL: [[Mersenne31; 16]; 4] = [
82    Mersenne31::new_array([
83        0x768bab52, 0x70e0ab7d, 0x3d266c8a, 0x6da42045, 0x600fef22, 0x41dace6b, 0x64f9bdd4,
84        0x5d42d4fe, 0x76b1516d, 0x6fc9a717, 0x70ac4fb6, 0x00194ef6, 0x22b644e2, 0x1f7916d5,
85        0x47581be2, 0x2710a123,
86    ]),
87    Mersenne31::new_array([
88        0x6284e867, 0x018d3afe, 0x5df99ef3, 0x4c1e467b, 0x566f6abc, 0x2994e427, 0x538a6d42,
89        0x5d7bf2cf, 0x7fda2dab, 0x0fd854c4, 0x46922fca, 0x3d7763a1, 0x19fd05ca, 0x0a4bbb43,
90        0x15075851, 0x3d903d76,
91    ]),
92    Mersenne31::new_array([
93        0x2d290ff7, 0x40809fa0, 0x59dac6ec, 0x127927a2, 0x6bbf0ea0, 0x0294140f, 0x24742976,
94        0x6e84c081, 0x22484f4a, 0x354cae59, 0x0453ffe1, 0x3f47a3cc, 0x0088204e, 0x6066e109,
95        0x3b7c4b80, 0x6b55665d,
96    ]),
97    Mersenne31::new_array([
98        0x3bc4b897, 0x735bf378, 0x508daf42, 0x1884fc2b, 0x7214f24c, 0x7498be0a, 0x1a60e640,
99        0x3303f928, 0x29b46376, 0x5c96bb68, 0x65d097a5, 0x1d358e9f, 0x4a9a9017, 0x4724cf76,
100        0x347af70f, 0x1e77e59a,
101    ]),
102];
103
104/// Round constants for width-16 Poseidon2 on Mersenne31.
105///
106/// Generated by the Grain LFSR with parameters:
107///     field_type=1, alpha=5 (exp_flag=0), n=31, t=16, R_F=8, R_P=14
108///
109/// Generated by `poseidon2/generate_constants.py --field mersenne31 --width 16`.
110///
111/// Layout: external_final (4 rounds × 16 elements).
112pub const MERSENNE31_POSEIDON2_RC_16_EXTERNAL_FINAL: [[Mersenne31; 16]; 4] = [
113    Mersenne31::new_array([
114        0x57090613, 0x1fa42108, 0x17bbef50, 0x1ff7e11c, 0x047b24ca, 0x4e140275, 0x4fa086f5,
115        0x079b309c, 0x1159bd47, 0x6d37e4e5, 0x075d8dce, 0x12121ca0, 0x7f6a7c40, 0x68e182ba,
116        0x5493201b, 0x0444a80e,
117    ]),
118    Mersenne31::new_array([
119        0x0064f4c6, 0x6467abe6, 0x66975762, 0x2af68f9b, 0x345b33be, 0x1b70d47f, 0x053db717,
120        0x381189cb, 0x43b915f8, 0x20df3694, 0x0f459d26, 0x77a0e97b, 0x2f73e739, 0x1876c2f9,
121        0x65a0e29a, 0x4cabefbe,
122    ]),
123    Mersenne31::new_array([
124        0x5abd1268, 0x4d34a760, 0x12771799, 0x69a0c9ac, 0x39091e55, 0x7f611cd0, 0x3af055da,
125        0x7ac0bbdf, 0x6e0f3a24, 0x41e3b6f7, 0x49b3756d, 0x568bc538, 0x20c079d8, 0x1701c72c,
126        0x7670dc6c, 0x5a439035,
127    ]),
128    Mersenne31::new_array([
129        0x7c93e00e, 0x561fbb4d, 0x1178907b, 0x02737406, 0x32fb24f1, 0x6323b60a, 0x6ab12418,
130        0x42c99cea, 0x155a0b97, 0x53d1c6aa, 0x2bd20347, 0x279b3d73, 0x4f5f3c70, 0x0245af6c,
131        0x238359d3, 0x49966a59,
132    ]),
133];
134
135/// Round constants for width-16 Poseidon2 on Mersenne31.
136///
137/// Generated by the Grain LFSR with parameters:
138///     field_type=1, alpha=5 (exp_flag=0), n=31, t=16, R_F=8, R_P=14
139///
140/// Generated by `poseidon2/generate_constants.py --field mersenne31 --width 16`.
141///
142/// Layout: internal (14 scalar constants).
143pub const MERSENNE31_POSEIDON2_RC_16_INTERNAL: [Mersenne31; 14] = Mersenne31::new_array([
144    0x7f7ec4bf, 0x0421926f, 0x5198e669, 0x34db3148, 0x4368bafd, 0x66685c7f, 0x78d3249a, 0x60187881,
145    0x76dad67a, 0x0690b437, 0x1ea95311, 0x40e5369a, 0x38f103fc, 0x1d226a21,
146]);
147
148/// Round constants for width-24 Poseidon2 on Mersenne31.
149///
150/// Generated by the Grain LFSR with parameters:
151///     field_type=1, alpha=5 (exp_flag=0), n=31, t=24, R_F=8, R_P=22
152///
153/// Generated by `poseidon2/generate_constants.py --field mersenne31 --width 24`.
154///
155/// Layout: external_initial (4 rounds × 24 elements).
156pub const MERSENNE31_POSEIDON2_RC_24_EXTERNAL_INITIAL: [[Mersenne31; 24]; 4] = [
157    Mersenne31::new_array([
158        0x1feaba61, 0x53224454, 0x6bceb9e2, 0x5019f9b4, 0x48726592, 0x2b22d0a8, 0x6151bbf9,
159        0x2f474b21, 0x2eb5f337, 0x3b645d87, 0x0942cef0, 0x65228c52, 0x78ffb30f, 0x4d2837c8,
160        0x0e17ac4f, 0x05546686, 0x046c06cc, 0x0b51c3b6, 0x568db763, 0x38b334e4, 0x57f5acf0,
161        0x19d32611, 0x77d02f4b, 0x6c82e9b8,
162    ]),
163    Mersenne31::new_array([
164        0x7148c1b6, 0x08067c75, 0x46d1e8c9, 0x30973b07, 0x20614f3b, 0x5c3ff851, 0x30503329,
165        0x4972e7cc, 0x02d1d8bc, 0x09d5bfa6, 0x097104c0, 0x7ba49a34, 0x4a07c2fc, 0x24c1ee69,
166        0x28a6ab41, 0x5d9108a0, 0x3a7851c7, 0x1dd495f9, 0x12b49ff4, 0x7bad5760, 0x5fed64c2,
167        0x66f5c96c, 0x7eafbd02, 0x39b3593b,
168    ]),
169    Mersenne31::new_array([
170        0x4a653b49, 0x75091dc1, 0x56e488e0, 0x1704a355, 0x745e4ff3, 0x392ef16e, 0x31e33fdf,
171        0x02c28c66, 0x36c3083a, 0x3104d1fa, 0x5b03cda3, 0x6641e1af, 0x37754b56, 0x396f5af9,
172        0x1a1a461a, 0x688e26f2, 0x6f829784, 0x1bb91d69, 0x5b788016, 0x704aa5c5, 0x0181869c,
173        0x41211e56, 0x0ce803a0, 0x23bff3a0,
174    ]),
175    Mersenne31::new_array([
176        0x17fb7064, 0x47317220, 0x76914b53, 0x219c1905, 0x16655528, 0x4df35544, 0x60808465,
177        0x3350f833, 0x03bccdc7, 0x0a87180a, 0x017a99f5, 0x6e945726, 0x15445504, 0x780533b1,
178        0x3b91bf38, 0x3fc77eb1, 0x4b4d960e, 0x3cd93d2e, 0x0ea4e976, 0x1d5306cc, 0x3a7ac284,
179        0x0ec22934, 0x4d979713, 0x51a41c65,
180    ]),
181];
182
183/// Round constants for width-24 Poseidon2 on Mersenne31.
184///
185/// Generated by the Grain LFSR with parameters:
186///     field_type=1, alpha=5 (exp_flag=0), n=31, t=24, R_F=8, R_P=22
187///
188/// Generated by `poseidon2/generate_constants.py --field mersenne31 --width 24`.
189///
190/// Layout: external_final (4 rounds × 24 elements).
191pub const MERSENNE31_POSEIDON2_RC_24_EXTERNAL_FINAL: [[Mersenne31; 24]; 4] = [
192    Mersenne31::new_array([
193        0x1c662299, 0x057c955a, 0x7ab6c0f2, 0x25a6ad0a, 0x75850b58, 0x48fd3793, 0x0b4366b1,
194        0x0fdd0d49, 0x7db419f9, 0x49b9cc0f, 0x48949716, 0x29c35890, 0x76445485, 0x1c27d30c,
195        0x10aa7a3b, 0x30f34fb6, 0x6fe06435, 0x02135ecd, 0x6caaba96, 0x3eb290d0, 0x22fd8d3b,
196        0x768b1525, 0x5be95814, 0x523d7fe9,
197    ]),
198    Mersenne31::new_array([
199        0x55e94cec, 0x47c42e1f, 0x1aa53b5e, 0x2fd1fe7e, 0x59230e91, 0x7472da66, 0x6443f2df,
200        0x2d9de19d, 0x6f7f6a84, 0x77800430, 0x0f014bc8, 0x7bf3d095, 0x26afd318, 0x582561f7,
201        0x5ee3198c, 0x6acc0000, 0x2f315e26, 0x27cac040, 0x2595081e, 0x5963b7da, 0x7e073565,
202        0x6cf3f5f1, 0x09f8a3a4, 0x0da8ccfe,
203    ]),
204    Mersenne31::new_array([
205        0x60be2365, 0x7ed742f5, 0x668b8031, 0x4bb03494, 0x59019333, 0x700e2878, 0x1cc45856,
206        0x1d1617f7, 0x7b988da6, 0x4eb4936c, 0x78c9f87e, 0x63ce3e94, 0x7178341b, 0x45bc2f86,
207        0x05b775bc, 0x704b0244, 0x29eed278, 0x47f43032, 0x2127b2e5, 0x1997903f, 0x24b3ce03,
208        0x0c32298c, 0x7d2b6f3a, 0x17fcaa81,
209    ]),
210    Mersenne31::new_array([
211        0x72f37fef, 0x3028e7a9, 0x5edd4d96, 0x1f96583b, 0x4cd6918a, 0x14880f0e, 0x69170359,
212        0x173cbd33, 0x0969e7f4, 0x6e7f23ab, 0x6182ea87, 0x4dcb1f5c, 0x585fa113, 0x729cb3b6,
213        0x01b3a27a, 0x1ba173e7, 0x4b33bcea, 0x63d93bbb, 0x6b3fbf99, 0x6f17e9d1, 0x0c3dd8ba,
214        0x0bc1f9a8, 0x64d3f370, 0x465a6a18,
215    ]),
216];
217
218/// Round constants for width-24 Poseidon2 on Mersenne31.
219///
220/// Generated by the Grain LFSR with parameters:
221///     field_type=1, alpha=5 (exp_flag=0), n=31, t=24, R_F=8, R_P=22
222///
223/// Generated by `poseidon2/generate_constants.py --field mersenne31 --width 24`.
224///
225/// Layout: internal (22 scalar constants).
226pub const MERSENNE31_POSEIDON2_RC_24_INTERNAL: [Mersenne31; 22] = Mersenne31::new_array([
227    0x22776a11, 0x5fa34268, 0x1415528d, 0x563fbd14, 0x34f45244, 0x120ea1b6, 0x261368a5, 0x27665ec1,
228    0x36be2805, 0x345c4784, 0x17efdcc1, 0x393e6530, 0x6da0b4b8, 0x31e5ded3, 0x675b27ac, 0x0ae88c30,
229    0x577841cc, 0x5fe06dec, 0x56b0691a, 0x7242de1f, 0x3c377529, 0x339b7523,
230]);
231
232/// Create a default width-16 Poseidon2 permutation for Mersenne31.
233pub fn default_mersenne31_poseidon2_16() -> Poseidon2Mersenne31<16> {
234    Poseidon2::new(
235        ExternalLayerConstants::new(
236            MERSENNE31_POSEIDON2_RC_16_EXTERNAL_INITIAL.to_vec(),
237            MERSENNE31_POSEIDON2_RC_16_EXTERNAL_FINAL.to_vec(),
238        ),
239        MERSENNE31_POSEIDON2_RC_16_INTERNAL.to_vec(),
240    )
241}
242
243/// Create a default width-24 Poseidon2 permutation for Mersenne31.
244pub fn default_mersenne31_poseidon2_24() -> Poseidon2Mersenne31<24> {
245    Poseidon2::new(
246        ExternalLayerConstants::new(
247            MERSENNE31_POSEIDON2_RC_24_EXTERNAL_INITIAL.to_vec(),
248            MERSENNE31_POSEIDON2_RC_24_EXTERNAL_FINAL.to_vec(),
249        ),
250        MERSENNE31_POSEIDON2_RC_24_INTERNAL.to_vec(),
251    )
252}
253
254/// An implementation of the matrix multiplications in the internal and external layers of Poseidon2.
255///
256/// This can act on `[A; WIDTH]` for any ring implementing `Algebra<Mersenne31>`.
257/// If you have either `[Mersenne31::Packing; WIDTH]` or `[Mersenne31; WIDTH]` it will be much faster
258/// to use `Poseidon2Mersenne31<WIDTH>` instead of building a Poseidon2 permutation using this.
259pub struct GenericPoseidon2LinearLayersMersenne31 {}
260
261const POSEIDON2_INTERNAL_MATRIX_DIAG_16_SHIFTS: [u8; 15] =
262    [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 13, 14, 15, 16];
263
264const POSEIDON2_INTERNAL_MATRIX_DIAG_24_SHIFTS: [u8; 23] = [
265    0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22,
266];
267
268const POSEIDON2_INTERNAL_MATRIX_DIAG_32_SHIFTS: [u8; 31] = [
269    0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25,
270    26, 27, 28, 29, 30,
271];
272
273/// Multiply state by the matrix (1 + Diag(V))
274///
275/// Here V is the vector [-2] + 1 << shifts. This used delayed reduction to be slightly faster.
276fn permute_mut<const N: usize>(state: &mut [Mersenne31; N], shifts: &[u8]) {
277    debug_assert_eq!(shifts.len() + 1, N);
278    let part_sum: u64 = state[1..].iter().map(|x| x.value as u64).sum();
279    let full_sum = part_sum + (state[0].value as u64);
280    let s0 = part_sum + (-state[0]).value as u64;
281    state[0] = from_u62(s0);
282    for i in 1..N {
283        let si = full_sum + ((state[i].value as u64) << shifts[i - 1]);
284        state[i] = from_u62(si);
285    }
286}
287
288impl InternalLayer<Mersenne31, 16, MERSENNE31_S_BOX_DEGREE> for Poseidon2InternalLayerMersenne31 {
289    /// Perform the internal layers of the Poseidon2 permutation on the given state.
290    fn permute_state(&self, state: &mut [Mersenne31; 16]) {
291        internal_permute_state(
292            state,
293            |x| permute_mut(x, &POSEIDON2_INTERNAL_MATRIX_DIAG_16_SHIFTS),
294            &self.internal_constants,
295        );
296    }
297}
298
299impl InternalLayer<Mersenne31, 24, MERSENNE31_S_BOX_DEGREE> for Poseidon2InternalLayerMersenne31 {
300    /// Perform the internal layers of the Poseidon2 permutation on the given state.
301    fn permute_state(&self, state: &mut [Mersenne31; 24]) {
302        internal_permute_state(
303            state,
304            |x| permute_mut(x, &POSEIDON2_INTERNAL_MATRIX_DIAG_24_SHIFTS),
305            &self.internal_constants,
306        );
307    }
308}
309
310impl InternalLayer<Mersenne31, 32, MERSENNE31_S_BOX_DEGREE> for Poseidon2InternalLayerMersenne31 {
311    /// Perform the internal layers of the Poseidon2 permutation on the given state.
312    fn permute_state(&self, state: &mut [Mersenne31; 32]) {
313        internal_permute_state(
314            state,
315            |x| permute_mut(x, &POSEIDON2_INTERNAL_MATRIX_DIAG_32_SHIFTS),
316            &self.internal_constants,
317        );
318    }
319}
320
321impl<const WIDTH: usize> ExternalLayer<Mersenne31, WIDTH, MERSENNE31_S_BOX_DEGREE>
322    for Poseidon2ExternalLayerMersenne31<WIDTH>
323{
324    /// Perform the initial external layers of the Poseidon2 permutation on the given state.
325    fn permute_state_initial(&self, state: &mut [Mersenne31; WIDTH]) {
326        external_initial_permute_state(
327            state,
328            self.external_constants.get_initial_constants(),
329            add_rc_and_sbox_generic,
330            &MDSMat4,
331        );
332    }
333
334    /// Perform the terminal external layers of the Poseidon2 permutation on the given state.
335    fn permute_state_terminal(&self, state: &mut [Mersenne31; WIDTH]) {
336        external_terminal_permute_state(
337            state,
338            self.external_constants.get_terminal_constants(),
339            add_rc_and_sbox_generic,
340            &MDSMat4,
341        );
342    }
343}
344
345impl GenericPoseidon2LinearLayers<16> for GenericPoseidon2LinearLayersMersenne31 {
346    fn internal_linear_layer<R: PrimeCharacteristicRing>(state: &mut [R; 16]) {
347        let part_sum: R = state[1..].iter().map(|r| r.dup()).sum();
348        let full_sum = part_sum.dup() + state[0].dup();
349
350        // The first three diagonal elements are -2, 1, 2 so we do something custom.
351        state[0] = part_sum - state[0].dup();
352        state[1] = full_sum.dup() + state[1].dup();
353        state[2] = full_sum.dup() + state[2].double();
354
355        // For the remaining elements we use the mul_2exp_u64 method.
356        // We need state[1..] as POSEIDON2_INTERNAL_MATRIX_DIAG_16_SHIFTS
357        // doesn't include the shift for the 0'th element as it is -2.
358        state[1..]
359            .iter_mut()
360            .zip(POSEIDON2_INTERNAL_MATRIX_DIAG_16_SHIFTS)
361            .skip(2)
362            .for_each(|(val, diag_shift)| {
363                *val = full_sum.dup() + val.dup().mul_2exp_u64(diag_shift as u64);
364            });
365    }
366}
367
368impl GenericPoseidon2LinearLayers<24> for GenericPoseidon2LinearLayersMersenne31 {
369    fn internal_linear_layer<R: PrimeCharacteristicRing>(state: &mut [R; 24]) {
370        let part_sum: R = state[1..].iter().map(|r| r.dup()).sum();
371        let full_sum = part_sum.dup() + state[0].dup();
372
373        // The first three diagonal elements are -2, 1, 2 so we do something custom.
374        state[0] = part_sum - state[0].dup();
375        state[1] = full_sum.dup() + state[1].dup();
376        state[2] = full_sum.dup() + state[2].double();
377
378        // For the remaining elements we use the mul_2exp_u64 method.
379        // We need state[1..] as POSEIDON2_INTERNAL_MATRIX_DIAG_24_SHIFTS
380        // doesn't include the shift for the 0'th element as it is -2.
381        state[1..]
382            .iter_mut()
383            .zip(POSEIDON2_INTERNAL_MATRIX_DIAG_24_SHIFTS)
384            .skip(2)
385            .for_each(|(val, diag_shift)| {
386                *val = full_sum.dup() + val.dup().mul_2exp_u64(diag_shift as u64);
387            });
388    }
389}
390
391impl GenericPoseidon2LinearLayers<32> for GenericPoseidon2LinearLayersMersenne31 {
392    fn internal_linear_layer<R: PrimeCharacteristicRing>(state: &mut [R; 32]) {
393        let part_sum: R = state[1..].iter().map(|r| r.dup()).sum();
394        let full_sum = part_sum.dup() + state[0].dup();
395
396        // The first three diagonal elements are -2, 1, 2 so we do something custom.
397        state[0] = part_sum - state[0].dup();
398        state[1] = full_sum.dup() + state[1].dup();
399        state[2] = full_sum.dup() + state[2].double();
400
401        // For the remaining elements we use the mul_2exp_u64 method.
402        state[1..]
403            .iter_mut()
404            .zip(POSEIDON2_INTERNAL_MATRIX_DIAG_32_SHIFTS)
405            .skip(2)
406            .for_each(|(val, diag_shift)| {
407                *val = full_sum.dup() + val.dup().mul_2exp_u64(diag_shift as u64);
408            });
409    }
410}
411
412#[cfg(test)]
413mod tests {
414    use p3_symmetric::Permutation;
415    use rand::SeedableRng;
416    use rand_xoshiro::Xoroshiro128Plus;
417
418    use super::*;
419
420    type F = Mersenne31;
421
422    // We need to make some round constants. We use Xoroshiro128Plus for this as we can easily match this PRNG in sage.
423    // See: https://github.com/0xPolygonZero/hash-constants for the sage code used to create all these tests.
424
425    /// Test on a roughly random input.
426    /// This random input is generated by the following sage code:
427    /// set_random_seed(16)
428    /// vector([M31.random_element() for t in range(16)]).
429    #[test]
430    fn test_poseidon2_width_16_random() {
431        let mut input: [F; 16] = Mersenne31::new_array([
432            894848333, 1437655012, 1200606629, 1690012884, 71131202, 1749206695, 1717947831,
433            120589055, 19776022, 42382981, 1831865506, 724844064, 171220207, 1299207443, 227047920,
434            1783754913,
435        ]);
436
437        let expected: [F; 16] = Mersenne31::new_array([
438            1124552602, 2127602268, 1834113265, 1207687593, 1891161485, 245915620, 981277919,
439            627265710, 1534924153, 1580826924, 887997842, 1526280482, 547791593, 1028672510,
440            1803086471, 323071277,
441        ]);
442
443        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
444        let perm = Poseidon2Mersenne31::new_from_rng_128(&mut rng);
445
446        perm.permute_mut(&mut input);
447        assert_eq!(input, expected);
448    }
449
450    /// Test on a roughly random input.
451    /// This random input is generated by the following sage code:
452    /// set_random_seed(24)
453    /// vector([M31.random_element() for t in range(24)]).
454    #[test]
455    fn test_poseidon2_width_24_random() {
456        let mut input: [F; 24] = Mersenne31::new_array([
457            886409618, 1327899896, 1902407911, 591953491, 648428576, 1844789031, 1198336108,
458            355597330, 1799586834, 59617783, 790334801, 1968791836, 559272107, 31054313,
459            1042221543, 474748436, 135686258, 263665994, 1962340735, 1741539604, 2026927696,
460            449439011, 1131357108, 50869465,
461        ]);
462
463        let expected: [F; 24] = Mersenne31::new_array([
464            87189408, 212775836, 954807335, 1424761838, 1222521810, 1264950009, 1891204592,
465            710452896, 957091834, 1776630156, 1091081383, 786687731, 1101902149, 1281649821,
466            436070674, 313565599, 1961711763, 2002894460, 2040173120, 854107426, 25198245,
467            1967213543, 604802266, 2086190331,
468        ]);
469
470        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
471        let perm = Poseidon2Mersenne31::new_from_rng_128(&mut rng);
472
473        perm.permute_mut(&mut input);
474        assert_eq!(input, expected);
475    }
476
477    #[test]
478    fn test_default_mersenne31_poseidon2_width_16() {
479        let mut input: [F; 16] =
480            Mersenne31::new_array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]);
481
482        let expected: [F; 16] = Mersenne31::new_array([
483            0x0b2c803a, 0x5b1ee4d1, 0x49c6b1e3, 0x2cdc280c, 0x310a60c8, 0x530a729e, 0x4e61bcb4,
484            0x2e84d3c3, 0x58709c08, 0x7e82ac42, 0x2162bcef, 0x6d153ab6, 0x742cf0e3, 0x2f21632d,
485            0x61adce1e, 0x1973d6f1,
486        ]);
487
488        let perm = default_mersenne31_poseidon2_16();
489        perm.permute_mut(&mut input);
490
491        assert_eq!(input, expected);
492    }
493
494    #[test]
495    fn test_default_mersenne31_poseidon2_width_24() {
496        let mut input: [F; 24] = Mersenne31::new_array([
497            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
498        ]);
499
500        let expected: [F; 24] = Mersenne31::new_array([
501            0x2040f051, 0x7261dbfa, 0x4fbd519e, 0x2320ecaf, 0x039ef27c, 0x48d60ad5, 0x73ca17ff,
502            0x6023111a, 0x6c5e31e7, 0x373cd90d, 0x75a3ae11, 0x00ecc878, 0x33a7c097, 0x244c2171,
503            0x7552a38e, 0x58d20817, 0x00feecb7, 0x47c43c88, 0x30d3001c, 0x24d09ba6, 0x71f241d9,
504            0x1c72ab2e, 0x4749f79d, 0x61ff7579,
505        ]);
506
507        let perm = default_mersenne31_poseidon2_24();
508        perm.permute_mut(&mut input);
509
510        assert_eq!(input, expected);
511    }
512
513    /// Test on a roughly random input for width 32.
514    #[test]
515    fn test_poseidon2_width_32_random() {
516        let mut input: [F; 32] = Mersenne31::new_array([
517            886409618, 1327899896, 1902407911, 591953491, 648428576, 1844789031, 1198336108,
518            355597330, 1799586834, 59617783, 790334801, 1968791836, 559272107, 31054313,
519            1042221543, 474748436, 135686258, 263665994, 1962340735, 1741539604, 2026927696,
520            449439011, 1131357108, 50869465, 894848333, 1437655012, 1200606629, 1690012884,
521            71131202, 1749206695, 1717947831, 120589055,
522        ]);
523
524        let original = input;
525
526        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
527        let perm = Poseidon2Mersenne31::new_from_rng_128(&mut rng);
528
529        perm.permute_mut(&mut input);
530        assert_ne!(input, original);
531    }
532}