Skip to main content

p3_baby_bear/
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, 3, 4) where multiplication is implemented using addition
7//! and inverse powers of 2 where it is possible to avoid monty reductions.
8//! Additionally, for technical reasons, having the first entry be -2 is useful.
9//!
10//! Optimized Diagonal for BabyBear16:
11//! [-2, 1, 2, 1/2, 3, 4, -1/2, -3, -4, 1/2^8, 1/4, 1/8, 1/2^27, -1/2^8, -1/16, -1/2^27].
12//! Optimized Diagonal for BabyBear24:
13//! [-2, 1, 2, 1/2, 3, 4, -1/2, -3, -4, 1/2^8, 1/4, 1/8, 1/16, 1/2^7, 1/2^9, 1/2^27, -1/2^8, -1/4, -1/8, -1/16, -1/32, -1/64, -1/2^7, -1/2^27]
14//! Optimized Diagonal for BabyBear32:
15//! [-2, 1, 2, 1/2, 3, 4, -1/2, -3, -4, 1/2^8, 1/4, 1/8, 1/16, 1/32, 1/64, 1/2^7, 1/2^9, 1/2^10, 1/2^12, 1/2^27, -1/2^8, -1/4, -1/8, -1/16, -1/32, -1/64, -1/2^7, -1/2^9, -1/2^10, -1/2^12, -1/2^14, -1/2^27]
16//! See poseidon2\src\diffusion.rs for information on how to double check these matrices in Sage.
17
18use p3_field::PrimeCharacteristicRing;
19use p3_monty_31::{
20    GenericPoseidon2LinearLayersMonty31, InternalLayerBaseParameters, InternalLayerParameters,
21    Poseidon2ExternalLayerMonty31, Poseidon2InternalLayerMonty31,
22};
23use p3_poseidon2::{ExternalLayerConstants, Poseidon2};
24
25use crate::{BABYBEAR_S_BOX_DEGREE, BabyBear, BabyBearParameters};
26
27pub type Poseidon2InternalLayerBabyBear<const WIDTH: usize> =
28    Poseidon2InternalLayerMonty31<BabyBearParameters, WIDTH, BabyBearInternalLayerParameters>;
29
30pub type Poseidon2ExternalLayerBabyBear<const WIDTH: usize> =
31    Poseidon2ExternalLayerMonty31<BabyBearParameters, WIDTH>;
32
33/// Number of full rounds per half for BabyBear 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 BABYBEAR_POSEIDON2_HALF_FULL_ROUNDS: usize = 4;
38
39/// Number of partial rounds for BabyBear 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.3562 · min{7.53, 15.5} = 11.682
45///
46/// With the +7.5% security margin: ⌈1.075 × 11.682⌉ = 13.
47pub const BABYBEAR_POSEIDON2_PARTIAL_ROUNDS_16: usize = 13;
48
49/// Number of partial rounds for BabyBear Poseidon2 (width 24).
50///
51/// Same Gröbner basis bound as width 16:
52///
53///   R_GB ≥ 17 + 0.3562 · min{5.12, 15.5} = 18.824
54///
55/// With the +7.5% security margin: ⌈1.075 × 18.824⌉ = 21.
56pub const BABYBEAR_POSEIDON2_PARTIAL_ROUNDS_24: usize = 21;
57
58/// An implementation of the Poseidon2 hash function specialised to run on the current architecture.
59///
60/// It acts on arrays of the form either `[BabyBear::Packing; WIDTH]` or `[BabyBear; WIDTH]`. For speed purposes,
61/// wherever possible, input arrays should of the form `[BabyBear::Packing; WIDTH]`.
62pub type Poseidon2BabyBear<const WIDTH: usize> = Poseidon2<
63    BabyBear,
64    Poseidon2ExternalLayerBabyBear<WIDTH>,
65    Poseidon2InternalLayerBabyBear<WIDTH>,
66    WIDTH,
67    BABYBEAR_S_BOX_DEGREE,
68>;
69
70/// An implementation of the matrix multiplications in the internal and external layers of Poseidon2.
71///
72/// This can act on `[A; WIDTH]` for any ring implementing `Algebra<BabyBear>`.
73/// If you have either `[BabyBear::Packing; WIDTH]` or `[BabyBear; WIDTH]` it will be much faster
74/// to use `Poseidon2BabyBear<WIDTH>` instead of building a Poseidon2 permutation using this.
75pub type GenericPoseidon2LinearLayersBabyBear =
76    GenericPoseidon2LinearLayersMonty31<BabyBearParameters, BabyBearInternalLayerParameters>;
77
78/// Round constants for width-16 Poseidon2 on BabyBear.
79///
80/// Generated by the Grain LFSR with parameters:
81///     field_type=1, alpha=7 (exp_flag=0), n=31, t=16, R_F=8, R_P=13
82///
83/// Generated by `poseidon2/generate_constants.py --field babybear --width 16`.
84///
85/// Layout: external_initial (4 rounds × 16 elements).
86pub const BABYBEAR_POSEIDON2_RC_16_EXTERNAL_INITIAL: [[BabyBear; 16]; 4] =
87    BabyBear::new_2d_array([
88        [
89            0x69cbb6af, 0x46ad93f9, 0x60a00f4e, 0x6b1297cd, 0x23189afe, 0x732e7bef, 0x72c246de,
90            0x2c941900, 0x0557eede, 0x1580496f, 0x3a3ea77b, 0x54f3f271, 0x0f49b029, 0x47872fe1,
91            0x221e2e36, 0x1ab7202e,
92        ],
93        [
94            0x487779a6, 0x3851c9d8, 0x38dc17c0, 0x209f8849, 0x268dcee8, 0x350c48da, 0x5b9ad32e,
95            0x0523272b, 0x3f89055b, 0x01e894b2, 0x13ddedde, 0x1b2ef334, 0x7507d8b4, 0x6ceeb94e,
96            0x52eb6ba2, 0x50642905,
97        ],
98        [
99            0x05453f3f, 0x06349efc, 0x6922787c, 0x04bfff9c, 0x768c714a, 0x3e9ff21a, 0x15737c9c,
100            0x2229c807, 0x0d47f88c, 0x097e0ecc, 0x27eadba0, 0x2d7d29e4, 0x3502aaa0, 0x0f475fd7,
101            0x29fbda49, 0x018afffd,
102        ],
103        [
104            0x0315b618, 0x6d4497d1, 0x1b171d9e, 0x52861abd, 0x2e5d0501, 0x3ec8646c, 0x6e5f250a,
105            0x148ae8e6, 0x17f5fa4a, 0x3e66d284, 0x0051aa3b, 0x483f7913, 0x2cfe5f15, 0x023427ca,
106            0x2cc78315, 0x1e36ea47,
107        ],
108    ]);
109
110/// Round constants for width-16 Poseidon2 on BabyBear.
111///
112/// Generated by the Grain LFSR with parameters:
113///     field_type=1, alpha=7 (exp_flag=0), n=31, t=16, R_F=8, R_P=13
114///
115/// Generated by `poseidon2/generate_constants.py --field babybear --width 16`.
116///
117/// Layout: external_final (4 rounds × 16 elements).
118pub const BABYBEAR_POSEIDON2_RC_16_EXTERNAL_FINAL: [[BabyBear; 16]; 4] = BabyBear::new_2d_array([
119    [
120        0x7290a80d, 0x6f7e5329, 0x598ec8a8, 0x76a859a0, 0x6559e868, 0x657b83af, 0x13271d3f,
121        0x1f876063, 0x0aeeae37, 0x706e9ca6, 0x46400cee, 0x72a05c26, 0x2c589c9e, 0x20bd37a7,
122        0x6a2d3d10, 0x20523767,
123    ],
124    [
125        0x5b8fe9c4, 0x2aa501d6, 0x1e01ac3e, 0x1448bc54, 0x5ce5ad1c, 0x4918a14d, 0x2c46a83f,
126        0x4fcf6876, 0x61d8d5c8, 0x6ddf4ff9, 0x11fda4d3, 0x02933a8f, 0x170eaf81, 0x5a9c314f,
127        0x49a12590, 0x35ec52a1,
128    ],
129    [
130        0x58eb1611, 0x5e481e65, 0x367125c9, 0x0eba33ba, 0x1fc28ded, 0x066399ad, 0x0cbec0ea,
131        0x75fd1af0, 0x50f5bf4e, 0x643d5f41, 0x6f4fe718, 0x5b3cbbde, 0x1e3afb3e, 0x296fb027,
132        0x45e1547b, 0x4a8db2ab,
133    ],
134    [
135        0x59986d19, 0x30bcdfa3, 0x1db63932, 0x1d7c2824, 0x53b33681, 0x0673b747, 0x038a98a3,
136        0x2c5bce60, 0x351979cd, 0x5008fb73, 0x547bca78, 0x711af481, 0x3f93bf64, 0x644d987b,
137        0x3c8bcd87, 0x608758b8,
138    ],
139]);
140
141/// Round constants for width-16 Poseidon2 on BabyBear.
142///
143/// Generated by the Grain LFSR with parameters:
144///     field_type=1, alpha=7 (exp_flag=0), n=31, t=16, R_F=8, R_P=13
145///
146/// Generated by `poseidon2/generate_constants.py --field babybear --width 16`.
147///
148/// Layout: internal (13 scalar constants).
149pub const BABYBEAR_POSEIDON2_RC_16_INTERNAL: [BabyBear; 13] = BabyBear::new_array([
150    0x5a8053c0, 0x693be639, 0x3858867d, 0x19334f6b, 0x128f0fd8, 0x4e2b1ccb, 0x61210ce0, 0x3c318939,
151    0x0b5b2f22, 0x2edb11d5, 0x213effdf, 0x0cac4606, 0x241af16d,
152]);
153
154/// Create a default width-16 Poseidon2 permutation for BabyBear.
155pub fn default_babybear_poseidon2_16() -> Poseidon2BabyBear<16> {
156    Poseidon2::new(
157        ExternalLayerConstants::new(
158            BABYBEAR_POSEIDON2_RC_16_EXTERNAL_INITIAL.to_vec(),
159            BABYBEAR_POSEIDON2_RC_16_EXTERNAL_FINAL.to_vec(),
160        ),
161        BABYBEAR_POSEIDON2_RC_16_INTERNAL.to_vec(),
162    )
163}
164
165/// Round constants for width-24 Poseidon2 on BabyBear.
166///
167/// Generated by the Grain LFSR with parameters:
168///     field_type=1, alpha=7 (exp_flag=0), n=31, t=24, R_F=8, R_P=21
169///
170/// Generated by `poseidon2/generate_constants.py --field babybear --width 24`.
171///
172/// Layout: external_initial (4 rounds × 24 elements).
173pub const BABYBEAR_POSEIDON2_RC_24_EXTERNAL_INITIAL: [[BabyBear; 24]; 4] =
174    BabyBear::new_2d_array([
175        [
176            0x0fa20c37, 0x0795bb97, 0x12c60b9c, 0x0eabd88e, 0x096485ca, 0x07093527, 0x1b1d4e50,
177            0x30a01ace, 0x3bd86f5a, 0x69af7c28, 0x3f94775f, 0x731560e8, 0x465a0ecd, 0x574ef807,
178            0x62fd4870, 0x52ccfe44, 0x14772b14, 0x4dedf371, 0x260acd7c, 0x1f51dc58, 0x75125532,
179            0x686a4d7b, 0x54bac179, 0x31947706,
180        ],
181        [
182            0x29799d3b, 0x6e01ae90, 0x203a7a64, 0x4f7e25be, 0x72503f77, 0x45bd3b69, 0x769bd6b4,
183            0x5a867f08, 0x4fdba082, 0x251c4318, 0x28f06201, 0x6788c43a, 0x4c6d6a99, 0x357784a8,
184            0x2abaf051, 0x770f7de6, 0x1794b784, 0x4796c57a, 0x724b7a10, 0x449989a7, 0x64935cf1,
185            0x59e14aac, 0x0e620bb8, 0x3af5a33b,
186        ],
187        [
188            0x4465cc0e, 0x019df68f, 0x4af8d068, 0x08784f82, 0x0cefdeae, 0x6337a467, 0x32fa7a16,
189            0x486f62d6, 0x386a7480, 0x20f17c4a, 0x54e50da8, 0x2012cf03, 0x5fe52950, 0x09afb6cd,
190            0x2523044e, 0x5c54d0ef, 0x71c01f3c, 0x60b2c4fb, 0x4050b379, 0x5e6a70a5, 0x418543f5,
191            0x71debe56, 0x1aad2994, 0x3368a483,
192        ],
193        [
194            0x07a86f3a, 0x5ea43ff1, 0x2443780e, 0x4ce444f7, 0x146f9882, 0x3132b089, 0x197ea856,
195            0x667030c3, 0x2317d5dc, 0x0c2c48a7, 0x56b2df66, 0x67bd81e9, 0x4fcdfb19, 0x4baaef32,
196            0x0328d30a, 0x6235760d, 0x12432912, 0x0a49e258, 0x030e1b70, 0x48caeb03, 0x49e4d9e9,
197            0x1051b5c6, 0x6a36dbbe, 0x4cff27a5,
198        ],
199    ]);
200
201/// Round constants for width-24 Poseidon2 on BabyBear.
202///
203/// Generated by the Grain LFSR with parameters:
204///     field_type=1, alpha=7 (exp_flag=0), n=31, t=24, R_F=8, R_P=21
205///
206/// Generated by `poseidon2/generate_constants.py --field babybear --width 24`.
207///
208/// Layout: external_final (4 rounds × 24 elements).
209pub const BABYBEAR_POSEIDON2_RC_24_EXTERNAL_FINAL: [[BabyBear; 24]; 4] = BabyBear::new_2d_array([
210    [
211        0x032959ad, 0x2b18af6a, 0x55d3dc8c, 0x43bd26c8, 0x0c41595f, 0x7048d2e2, 0x00db8983,
212        0x2af563d7, 0x6e84758f, 0x611d64e1, 0x1f9977e2, 0x64163a0a, 0x5c5fc27b, 0x02e22561,
213        0x3a2d75db, 0x1ba7b71a, 0x34343f64, 0x7406b35d, 0x19df8299, 0x6ff4480a, 0x514a81c8,
214        0x57ab52ce, 0x6ad69f52, 0x3e0c0e0d,
215    ],
216    [
217        0x48126114, 0x2a9d62cc, 0x17441f23, 0x485762bb, 0x2f218674, 0x06fdc64a, 0x0861b7f2,
218        0x3b36eee6, 0x70a11040, 0x04b31737, 0x3722a872, 0x2a351c63, 0x623560dc, 0x62584ab2,
219        0x382c7c04, 0x3bf9edc7, 0x0e38fe51, 0x376f3b10, 0x5381e178, 0x3afc61c7, 0x5c1bcb4d,
220        0x6643ce1f, 0x2d0af1c1, 0x08f583cc,
221    ],
222    [
223        0x5d6ff60f, 0x6324c1e5, 0x74412fb7, 0x70c0192e, 0x0b72f141, 0x4067a111, 0x57388c4f,
224        0x351009ec, 0x0974c159, 0x539a58b3, 0x038c0cff, 0x476c0392, 0x3f7bc15f, 0x4491dd2c,
225        0x4d1fef55, 0x04936ae3, 0x58214dd4, 0x683c6aad, 0x1b42f16b, 0x6dc79135, 0x2d4e71ec,
226        0x3e2946ea, 0x59dce8db, 0x6cee892a,
227    ],
228    [
229        0x47f07350, 0x7106ce93, 0x3bd4a7a9, 0x2bfe636a, 0x430011e9, 0x001cd66a, 0x307faf5b,
230        0x0d9ef3fe, 0x6d40043a, 0x2e8f470c, 0x1b6865e8, 0x0c0e6c01, 0x4d41981f, 0x423b9d3d,
231        0x410408cc, 0x263f0884, 0x5311bbd0, 0x4dae58d8, 0x30401cea, 0x09afa575, 0x4b3d5b42,
232        0x63ac0b37, 0x5fe5bb14, 0x5244e9d4,
233    ],
234]);
235
236/// Round constants for width-24 Poseidon2 on BabyBear.
237///
238/// Generated by the Grain LFSR with parameters:
239///     field_type=1, alpha=7 (exp_flag=0), n=31, t=24, R_F=8, R_P=21
240///
241/// Generated by `poseidon2/generate_constants.py --field babybear --width 24`.
242///
243/// Layout: internal (21 scalar constants).
244pub const BABYBEAR_POSEIDON2_RC_24_INTERNAL: [BabyBear; 21] = BabyBear::new_array([
245    0x1da78ec2, 0x730b0924, 0x3eb56cf3, 0x5bd93073, 0x37204c97, 0x51642d89, 0x66e943e8, 0x1a3e72de,
246    0x70beb1e9, 0x30ff3b3f, 0x4240d1c4, 0x12647b8d, 0x65d86965, 0x49ef4d7c, 0x47785697, 0x46b3969f,
247    0x5c7b7a0e, 0x7078fc60, 0x4f22d482, 0x482a9aee, 0x6beb839d,
248]);
249
250/// Create a default width-24 Poseidon2 permutation for BabyBear.
251pub fn default_babybear_poseidon2_24() -> Poseidon2BabyBear<24> {
252    Poseidon2::new(
253        ExternalLayerConstants::new(
254            BABYBEAR_POSEIDON2_RC_24_EXTERNAL_INITIAL.to_vec(),
255            BABYBEAR_POSEIDON2_RC_24_EXTERNAL_FINAL.to_vec(),
256        ),
257        BABYBEAR_POSEIDON2_RC_24_INTERNAL.to_vec(),
258    )
259}
260
261/// Contains data needed to define the internal layers of the Poseidon2 permutation.
262#[derive(Debug, Clone, Default)]
263pub struct BabyBearInternalLayerParameters;
264
265impl InternalLayerBaseParameters<BabyBearParameters, 16> for BabyBearInternalLayerParameters {
266    /// Perform the internal matrix multiplication: s -> (1 + Diag(V))s.
267    /// We ignore `state[0]` as it is handled separately.
268    fn internal_layer_mat_mul<R: PrimeCharacteristicRing>(state: &mut [R; 16], sum: R) {
269        // The diagonal matrix is defined by the vector:
270        // V = [-2, 1, 2, 1/2, 3, 4, -1/2, -3, -4, 1/2^8, 1/4, 1/8, 1/2^27, -1/2^8, -1/16, -1/2^27]
271        state[1] += sum.dup();
272        state[2] = state[2].double() + sum.dup();
273        state[3] = state[3].halve() + sum.dup();
274        state[4] = sum.dup() + state[4].double() + state[4].dup();
275        state[5] = sum.dup() + state[5].double().double();
276        state[6] = sum.dup() - state[6].halve();
277        state[7] = sum.dup() - (state[7].double() + state[7].dup());
278        state[8] = sum.dup() - state[8].double().double();
279        state[9] = state[9].div_2exp_u64(8) + sum.dup();
280        state[10] = state[10].div_2exp_u64(2) + sum.dup();
281        state[11] = state[11].div_2exp_u64(3) + sum.dup();
282        state[12] = state[12].div_2exp_u64(27) + sum.dup();
283        state[13] = sum.dup() - state[13].div_2exp_u64(8);
284        state[14] = sum.dup() - state[14].div_2exp_u64(4);
285        state[15] = sum - state[15].div_2exp_u64(27);
286    }
287}
288
289impl InternalLayerBaseParameters<BabyBearParameters, 24> for BabyBearInternalLayerParameters {
290    /// Perform the internal matrix multiplication: s -> (1 + Diag(V))s.
291    /// We ignore `state[0]` as it is handled separately.
292    fn internal_layer_mat_mul<R: PrimeCharacteristicRing>(state: &mut [R; 24], sum: R) {
293        // The diagonal matrix is defined by the vector:
294        // V = [-2, 1, 2, 1/2, 3, 4, -1/2, -3, -4, 1/2^8, 1/4, 1/8, 1/16, 1/2^7, 1/2^9, 1/2^27, -1/2^8, -1/4, -1/8, -1/16, -1/32, -1/64, -1/2^7, -1/2^27]
295        state[1] += sum.dup();
296        state[2] = state[2].double() + sum.dup();
297        state[3] = state[3].halve() + sum.dup();
298        state[4] = sum.dup() + state[4].double() + state[4].dup();
299        state[5] = sum.dup() + state[5].double().double();
300        state[6] = sum.dup() - state[6].halve();
301        state[7] = sum.dup() - (state[7].double() + state[7].dup());
302        state[8] = sum.dup() - state[8].double().double();
303        state[9] = state[9].div_2exp_u64(8) + sum.dup();
304        state[10] = state[10].div_2exp_u64(2) + sum.dup();
305        state[11] = state[11].div_2exp_u64(3) + sum.dup();
306        state[12] = state[12].div_2exp_u64(4) + sum.dup();
307        state[13] = state[13].div_2exp_u64(7) + sum.dup();
308        state[14] = state[14].div_2exp_u64(9) + sum.dup();
309        state[15] = state[15].div_2exp_u64(27) + sum.dup();
310        state[16] = sum.dup() - state[16].div_2exp_u64(8);
311        state[17] = sum.dup() - state[17].div_2exp_u64(2);
312        state[18] = sum.dup() - state[18].div_2exp_u64(3);
313        state[19] = sum.dup() - state[19].div_2exp_u64(4);
314        state[20] = sum.dup() - state[20].div_2exp_u64(5);
315        state[21] = sum.dup() - state[21].div_2exp_u64(6);
316        state[22] = sum.dup() - state[22].div_2exp_u64(7);
317        state[23] = sum - state[23].div_2exp_u64(27);
318    }
319}
320
321impl InternalLayerBaseParameters<BabyBearParameters, 32> for BabyBearInternalLayerParameters {
322    /// Perform the internal matrix multiplication: s -> (1 + Diag(V))s.
323    /// We ignore `state[0]` as it is handled separately.
324    fn internal_layer_mat_mul<R: PrimeCharacteristicRing>(state: &mut [R; 32], sum: R) {
325        // The diagonal matrix is defined by the vector:
326        // V = [-2, 1, 2, 1/2, 3, 4, -1/2, -3, -4,
327        //      1/2^8, 1/4, 1/8, 1/16, 1/32, 1/64, 1/2^7, 1/2^9, 1/2^10, 1/2^12, 1/2^27,
328        //      -1/2^8, -1/4, -1/8, -1/16, -1/32, -1/64, -1/2^7, -1/2^9, -1/2^10, -1/2^12, -1/2^14, -1/2^27]
329        state[1] += sum.dup();
330        state[2] = state[2].double() + sum.dup();
331        state[3] = state[3].halve() + sum.dup();
332        state[4] = sum.dup() + state[4].double() + state[4].dup();
333        state[5] = sum.dup() + state[5].double().double();
334        state[6] = sum.dup() - state[6].halve();
335        state[7] = sum.dup() - (state[7].double() + state[7].dup());
336        state[8] = sum.dup() - state[8].double().double();
337        state[9] = state[9].div_2exp_u64(8) + sum.dup();
338        state[10] = state[10].div_2exp_u64(2) + sum.dup();
339        state[11] = state[11].div_2exp_u64(3) + sum.dup();
340        state[12] = state[12].div_2exp_u64(4) + sum.dup();
341        state[13] = state[13].div_2exp_u64(5) + sum.dup();
342        state[14] = state[14].div_2exp_u64(6) + sum.dup();
343        state[15] = state[15].div_2exp_u64(7) + sum.dup();
344        state[16] = state[16].div_2exp_u64(9) + sum.dup();
345        state[17] = state[17].div_2exp_u64(10) + sum.dup();
346        state[18] = state[18].div_2exp_u64(12) + sum.dup();
347        state[19] = state[19].div_2exp_u64(27) + sum.dup();
348        state[20] = sum.dup() - state[20].div_2exp_u64(8);
349        state[21] = sum.dup() - state[21].div_2exp_u64(2);
350        state[22] = sum.dup() - state[22].div_2exp_u64(3);
351        state[23] = sum.dup() - state[23].div_2exp_u64(4);
352        state[24] = sum.dup() - state[24].div_2exp_u64(5);
353        state[25] = sum.dup() - state[25].div_2exp_u64(6);
354        state[26] = sum.dup() - state[26].div_2exp_u64(7);
355        state[27] = sum.dup() - state[27].div_2exp_u64(9);
356        state[28] = sum.dup() - state[28].div_2exp_u64(10);
357        state[29] = sum.dup() - state[29].div_2exp_u64(12);
358        state[30] = sum.dup() - state[30].div_2exp_u64(14);
359        state[31] = sum - state[31].div_2exp_u64(27);
360    }
361}
362
363impl InternalLayerParameters<BabyBearParameters, 16> for BabyBearInternalLayerParameters {}
364impl InternalLayerParameters<BabyBearParameters, 24> for BabyBearInternalLayerParameters {}
365impl InternalLayerParameters<BabyBearParameters, 32> for BabyBearInternalLayerParameters {}
366
367#[cfg(test)]
368mod tests {
369    use p3_symmetric::Permutation;
370    use rand::{RngExt, SeedableRng};
371    use rand_xoshiro::Xoroshiro128Plus;
372
373    use super::*;
374
375    type F = BabyBear;
376
377    /// Test on a roughly random input.
378    /// This random input is generated by the following sage code:
379    /// set_random_seed(16)
380    /// vector([BB.random_element() for t in range(16)]).
381    #[test]
382    fn test_poseidon2_width_16_random() {
383        let mut input: [F; 16] = BabyBear::new_array([
384            894848333, 1437655012, 1200606629, 1690012884, 71131202, 1749206695, 1717947831,
385            120589055, 19776022, 42382981, 1831865506, 724844064, 171220207, 1299207443, 227047920,
386            1783754913,
387        ]);
388
389        let expected: [F; 16] = BabyBear::new_array([
390            1255099308, 941729227, 93609187, 112406640, 492658670, 1824768948, 812517469,
391            1055381989, 670973674, 1407235524, 891397172, 1003245378, 1381303998, 1564172645,
392            1399931635, 1005462965,
393        ]);
394
395        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
396        let perm = Poseidon2BabyBear::new_from_rng_128(&mut rng);
397
398        perm.permute_mut(&mut input);
399        assert_eq!(input, expected);
400    }
401
402    /// Test on a roughly random input.
403    /// This random input is generated by the following sage code:
404    /// set_random_seed(24)
405    /// vector([BB.random_element() for t in range(24)]).
406    #[test]
407    fn test_poseidon2_width_24_random() {
408        let mut input: [F; 24] = BabyBear::new_array([
409            886409618, 1327899896, 1902407911, 591953491, 648428576, 1844789031, 1198336108,
410            355597330, 1799586834, 59617783, 790334801, 1968791836, 559272107, 31054313,
411            1042221543, 474748436, 135686258, 263665994, 1962340735, 1741539604, 449439011,
412            1131357108, 50869465, 1589724894,
413        ]);
414
415        let expected: [F; 24] = BabyBear::new_array([
416            249424342, 562262148, 757431114, 354243402, 57767055, 976981973, 1393169022,
417            1774550827, 1527742125, 1019514605, 1776327602, 266236737, 1412355182, 1070239213,
418            426390978, 1775539440, 1527732214, 1101406020, 1417710778, 1699632661, 413672313,
419            820348291, 1067197851, 1669055675,
420        ]);
421
422        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
423        let perm = Poseidon2BabyBear::new_from_rng_128(&mut rng);
424
425        perm.permute_mut(&mut input);
426
427        assert_eq!(input, expected);
428    }
429
430    /// Test the generic internal layer against the optimized internal layer
431    /// for a random input of width 16.
432    #[test]
433    fn test_generic_internal_linear_layer_16() {
434        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
435        let mut input1: [F; 16] = rng.random();
436        let mut input2 = input1;
437
438        let part_sum: F = input1[1..].iter().copied().sum();
439        let full_sum = part_sum + input1[0];
440
441        input1[0] = part_sum - input1[0];
442
443        BabyBearInternalLayerParameters::internal_layer_mat_mul(&mut input1, full_sum);
444        BabyBearInternalLayerParameters::generic_internal_linear_layer(&mut input2);
445
446        assert_eq!(input1, input2);
447    }
448
449    /// Test the generic internal layer against the optimized internal layer
450    /// for a random input of width 24.
451    #[test]
452    fn test_generic_internal_linear_layer_24() {
453        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
454        let mut input1: [F; 24] = rng.random();
455        let mut input2 = input1;
456
457        let part_sum: F = input1[1..].iter().copied().sum();
458        let full_sum = part_sum + input1[0];
459
460        input1[0] = part_sum - input1[0];
461
462        BabyBearInternalLayerParameters::internal_layer_mat_mul(&mut input1, full_sum);
463        BabyBearInternalLayerParameters::generic_internal_linear_layer(&mut input2);
464
465        assert_eq!(input1, input2);
466    }
467
468    #[test]
469    fn test_default_babybear_poseidon2_width_16() {
470        let mut input: [F; 16] = BabyBear::new_array([
471            894848333, 1437655012, 1200606629, 1690012884, 71131202, 1749206695, 1717947831,
472            120589055, 19776022, 42382981, 1831865506, 724844064, 171220207, 1299207443, 227047920,
473            1783754913,
474        ]);
475
476        let expected: [F; 16] = BabyBear::new_array([
477            516096821, 90309867, 1101817252, 1660784290, 360715097, 1789519026, 1788910906,
478            563338433, 319524748, 1741414159, 1650859320, 894311162, 1121347488, 1692793758,
479            1052633829, 1344246938,
480        ]);
481
482        let perm = default_babybear_poseidon2_16();
483        perm.permute_mut(&mut input);
484
485        assert_eq!(input, expected);
486    }
487
488    #[test]
489    fn test_default_babybear_poseidon2_width_24() {
490        let mut input: [F; 24] = BabyBear::new_array([
491            886409618, 1327899896, 1902407911, 591953491, 648428576, 1844789031, 1198336108,
492            355597330, 1799586834, 59617783, 790334801, 1968791836, 559272107, 31054313,
493            1042221543, 474748436, 135686258, 263665994, 1962340735, 1741539604, 2026927696,
494            449439011, 1131357108, 50869465,
495        ]);
496
497        let expected: [F; 24] = BabyBear::new_array([
498            882297297, 1264077610, 512812497, 782602970, 867738552, 1251075457, 309180082,
499            340784773, 524041877, 351272188, 404451680, 15001466, 322926653, 1773004150,
500            1718440818, 674682955, 1154713225, 1719133502, 324232301, 1005243141, 443371079,
501            268735940, 770060019, 718377682,
502        ]);
503
504        let perm = default_babybear_poseidon2_24();
505        perm.permute_mut(&mut input);
506
507        assert_eq!(input, expected);
508    }
509    /// Test on a roughly random input for width 32.
510    #[test]
511    fn test_poseidon2_width_32_random() {
512        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
513        let perm = Poseidon2BabyBear::<32>::new_from_rng_128(&mut rng);
514
515        let mut input: [F; 32] = rng.random();
516        let original = input;
517
518        perm.permute_mut(&mut input);
519        assert_ne!(input, original);
520    }
521
522    /// Test the generic internal layer against the optimized internal layer
523    /// for a random input of width 32.
524    #[test]
525    fn test_generic_internal_linear_layer_32() {
526        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
527        let mut input1: [F; 32] = rng.random();
528        let mut input2 = input1;
529
530        let part_sum: F = input1[1..].iter().copied().sum();
531        let full_sum = part_sum + input1[0];
532
533        input1[0] = part_sum - input1[0];
534
535        BabyBearInternalLayerParameters::internal_layer_mat_mul(&mut input1, full_sum);
536        BabyBearInternalLayerParameters::generic_internal_linear_layer(&mut input2);
537
538        assert_eq!(input1, input2);
539    }
540}