p3_koala_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 KoalaBear16:
11//! [-2, 1, 2, 1/2, 3, 4, -1/2, -3, -4, 1/2^8, 1/8, 1/2^24, -1/2^8, -1/8, -1/16, -1/2^24]
12//! Optimized Diagonal for KoalaBear24:
13//! [-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^24, -1/2^8, -1/8, -1/16, -1/32, -1/64, -1/2^7, -1/2^9, -1/2^24]
14//! See poseidon2\src\diffusion.rs for information on how to double check these matrices in Sage.
15
16use p3_field::PrimeCharacteristicRing;
17use p3_monty_31::{
18    GenericPoseidon2LinearLayersMonty31, InternalLayerBaseParameters, InternalLayerParameters,
19    Poseidon2ExternalLayerMonty31, Poseidon2InternalLayerMonty31,
20};
21use p3_poseidon2::{ExternalLayerConstants, Poseidon2};
22
23use crate::{KoalaBear, KoalaBearParameters};
24
25pub type Poseidon2InternalLayerKoalaBear<const WIDTH: usize> =
26    Poseidon2InternalLayerMonty31<KoalaBearParameters, WIDTH, KoalaBearInternalLayerParameters>;
27
28pub type Poseidon2ExternalLayerKoalaBear<const WIDTH: usize> =
29    Poseidon2ExternalLayerMonty31<KoalaBearParameters, WIDTH>;
30
31/// Degree of the chosen permutation polynomial for KoalaBear, used as the Poseidon2 S-Box.
32///
33/// As p - 1 = 127 * 2^{24} we have a lot of choice in degree D satisfying gcd(p - 1, D) = 1.
34/// Experimentation suggests that the optimal choice is the smallest available one, namely 3.
35const KOALABEAR_S_BOX_DEGREE: u64 = 3;
36
37/// An implementation of the Poseidon2 hash function specialised to run on the current architecture.
38///
39/// It acts on arrays of the form either `[KoalaBear::Packing; WIDTH]` or `[KoalaBear; WIDTH]`. For speed purposes,
40/// wherever possible, input arrays should of the form `[KoalaBear::Packing; WIDTH]`.
41pub type Poseidon2KoalaBear<const WIDTH: usize> = Poseidon2<
42    KoalaBear,
43    Poseidon2ExternalLayerKoalaBear<WIDTH>,
44    Poseidon2InternalLayerKoalaBear<WIDTH>,
45    WIDTH,
46    KOALABEAR_S_BOX_DEGREE,
47>;
48
49/// An implementation of the matrix multiplications in the internal and external layers of Poseidon2.
50///
51/// This can act on `[A; WIDTH]` for any ring implementing `Algebra<BabyBear>`.
52/// If you have either `[KoalaBear::Packing; WIDTH]` or `[KoalaBear; WIDTH]` it will be much faster
53/// to use `Poseidon2KoalaBear<WIDTH>` instead of building a Poseidon2 permutation using this.
54pub type GenericPoseidon2LinearLayersKoalaBear =
55    GenericPoseidon2LinearLayersMonty31<KoalaBearParameters, KoalaBearInternalLayerParameters>;
56
57/// Initial round constants for the 16-width Poseidon2 external layer on KoalaBear.
58///
59/// See Poseidon paper for more details: https://eprint.iacr.org/2019/458
60/// Sage script used to generate these constants: https://github.com/HorizenLabs/poseidon2/blob/main/poseidon2_rust_params.sage
61pub const KOALABEAR_RC16_EXTERNAL_INITIAL: [[KoalaBear; 16]; 4] = KoalaBear::new_2d_array([
62    [
63        2128964168, 288780357, 316938561, 2126233899, 426817493, 1714118888, 1045008582,
64        1738510837, 889721787, 8866516, 681576474, 419059826, 1596305521, 1583176088, 1584387047,
65        1529751136,
66    ],
67    [
68        1863858111, 1072044075, 517831365, 1464274176, 1138001621, 428001039, 245709561,
69        1641420379, 1365482496, 770454828, 693167409, 757905735, 136670447, 436275702, 525466355,
70        1559174242,
71    ],
72    [
73        1030087950, 869864998, 322787870, 267688717, 948964561, 740478015, 679816114, 113662466,
74        2066544572, 1744924186, 367094720, 1380455578, 1842483872, 416711434, 1342291586,
75        1692058446,
76    ],
77    [
78        1493348999, 1113949088, 210900530, 1071655077, 610242121, 1136339326, 2020858841,
79        1019840479, 678147278, 1678413261, 1361743414, 61132629, 1209546658, 64412292, 1936878279,
80        1980661727,
81    ],
82]);
83
84/// Final round constants for the 16-width Poseidon2's external layer on KoalaBear.
85///
86/// See Poseidon paper for more details: https://eprint.iacr.org/2019/458
87/// Sage script used to generate these constants: https://github.com/HorizenLabs/poseidon2/blob/main/poseidon2_rust_params.sage
88pub const KOALABEAR_RC16_EXTERNAL_FINAL: [[KoalaBear; 16]; 4] = KoalaBear::new_2d_array([
89    [
90        1423960925, 2101391318, 1915532054, 275400051, 1168624859, 1141248885, 356546469,
91        1165250474, 1320543726, 932505663, 1204226364, 1452576828, 1774936729, 926808140,
92        1184948056, 1186493834,
93    ],
94    [
95        843181003, 185193011, 452207447, 510054082, 1139268644, 630873441, 669538875, 462500858,
96        876500520, 1214043330, 383937013, 375087302, 636912601, 307200505, 390279673, 1999916485,
97    ],
98    [
99        1518476730, 1606686591, 1410677749, 1581191572, 1004269969, 143426723, 1747283099,
100        1016118214, 1749423722, 66331533, 1177761275, 1581069649, 1851371119, 852520128,
101        1499632627, 1820847538,
102    ],
103    [
104        150757557, 884787840, 619710451, 1651711087, 505263814, 212076987, 1482432120, 1458130652,
105        382871348, 417404007, 2066495280, 1996518884, 902934924, 582892981, 1337064375, 1199354861,
106    ],
107]);
108
109/// Round constants for the 16-width Poseidon2's internal layer on KoalaBear.
110///
111/// See Poseidon paper for more details: https://eprint.iacr.org/2019/458
112/// Sage script used to generate these constants: https://github.com/HorizenLabs/poseidon2/blob/main/poseidon2_rust_params.sage
113pub const KOALABEAR_RC16_INTERNAL: [KoalaBear; 20] = KoalaBear::new_array([
114    2102596038, 1533193853, 1436311464, 2012303432, 839997195, 1225781098, 2011967775, 575084315,
115    1309329169, 786393545, 995788880, 1702925345, 1444525226, 908073383, 1811535085, 1531002367,
116    1635653662, 1585100155, 867006515, 879151050,
117]);
118
119/// A default Poseidon2 for KoalaBear using the round constants from the original specification.
120///
121/// See Poseidon paper for more details: https://eprint.iacr.org/2019/458
122/// Sage script used to generate these constants: https://github.com/HorizenLabs/poseidon2/blob/main/poseidon2_rust_params.sage
123pub fn default_koalabear_poseidon2_16() -> Poseidon2KoalaBear<16> {
124    Poseidon2::new(
125        ExternalLayerConstants::new(
126            KOALABEAR_RC16_EXTERNAL_INITIAL.to_vec(),
127            KOALABEAR_RC16_EXTERNAL_FINAL.to_vec(),
128        ),
129        KOALABEAR_RC16_INTERNAL.to_vec(),
130    )
131}
132
133/// Initial round constants for the 24-width Poseidon2 external layer on KoalaBear.
134///
135/// See Poseidon paper for more details: https://eprint.iacr.org/2019/458
136/// Sage script used to generate these constants: https://github.com/HorizenLabs/poseidon2/blob/main/poseidon2_rust_params.sage
137pub const KOALABEAR_RC24_EXTERNAL_INITIAL: [[KoalaBear; 24]; 4] = KoalaBear::new_2d_array([
138    [
139        487143900, 1829048205, 1652578477, 646002781, 1044144830, 53279448, 1519499836, 22697702,
140        1768655004, 230479744, 1484895689, 705130286, 1429811285, 1695785093, 1417332623,
141        1115801016, 1048199020, 878062617, 738518649, 249004596, 1601837737, 24601614, 245692625,
142        364803730,
143    ],
144    [
145        1857019234, 1906668230, 1916890890, 835590867, 557228239, 352829675, 515301498, 973918075,
146        954515249, 1142063750, 1795549558, 608869266, 1850421928, 2028872854, 1197543771,
147        1027240055, 1976813168, 963257461, 652017844, 2113212249, 213459679, 90747280, 1540619478,
148        324138382,
149    ],
150    [
151        1377377119, 294744504, 512472871, 668081958, 907306515, 518526882, 1907091534, 1152942192,
152        1572881424, 720020214, 729527057, 1762035789, 86171731, 205890068, 453077400, 1201344594,
153        986483134, 125174298, 2050269685, 1895332113, 749706654, 40566555, 742540942, 1735551813,
154    ],
155    [
156        162985276, 1943496073, 1469312688, 703013107, 1979485151, 1278193166, 548674995,
157        2118718736, 749596440, 1476142294, 1293606474, 918523452, 890353212, 1691895663,
158        1932240646, 1180911992, 86098300, 1592168978, 895077289, 724819849, 1697986774, 1608418116,
159        1083269213, 691256798,
160    ],
161]);
162
163/// Final round constants for the 24-width Poseidon2's external layer on KoalaBear.
164///
165/// See Poseidon paper for more details: https://eprint.iacr.org/2019/458
166/// Sage script used to generate these constants: https://github.com/HorizenLabs/poseidon2/blob/main/poseidon2_rust_params.sage
167pub const KOALABEAR_RC24_EXTERNAL_FINAL: [[KoalaBear; 24]; 4] = KoalaBear::new_2d_array([
168    [
169        328586442, 1572520009, 1375479591, 322991001, 967600467, 1172861548, 1973891356,
170        1503625929, 1881993531, 40601941, 1155570620, 571547775, 1361622243, 1495024047,
171        1733254248, 964808915, 763558040, 1887228519, 994888261, 718330940, 213359415, 603124968,
172        1038411577, 2099454809,
173    ],
174    [
175        949846777, 630926956, 1168723439, 222917504, 1527025973, 1009157017, 2029957881, 805977836,
176        1347511739, 540019059, 589807745, 440771316, 1530063406, 761076336, 87974206, 1412686751,
177        1230318064, 514464425, 1469011754, 1770970737, 1510972858, 965357206, 209398053, 778802532,
178    ],
179    [
180        40567006, 1984217577, 1545851069, 879801839, 1611910970, 1215591048, 330802499, 1051639108,
181        321036, 511927202, 591603098, 1775897642, 115598532, 278200718, 233743176, 525096211,
182        1335507608, 830017835, 1380629279, 560028578, 598425701, 302162385, 567434115, 1859222575,
183    ],
184    [
185        958294793, 1582225556, 1781487858, 1570246000, 1067748446, 526608119, 1666453343,
186        1786918381, 348203640, 1860035017, 1489902626, 1904576699, 860033965, 1954077639,
187        1685771567, 971513929, 1877873770, 137113380, 520695829, 806829080, 1408699405, 1613277964,
188        793223662, 648443918,
189    ],
190]);
191
192/// Round constants for the 24-width Poseidon2's internal layer on KoalaBear.
193///
194/// See Poseidon paper for more details: https://eprint.iacr.org/2019/458
195/// Sage script used to generate these constants: https://github.com/HorizenLabs/poseidon2/blob/main/poseidon2_rust_params.sage
196pub const KOALABEAR_RC24_INTERNAL: [KoalaBear; 23] = KoalaBear::new_array([
197    893435011, 403879071, 1363789863, 1662900517, 2043370, 2109755796, 931751726, 2091644718,
198    606977583, 185050397, 946157136, 1350065230, 1625860064, 122045240, 880989921, 145137438,
199    1059782436, 1477755661, 335465138, 1640704282, 1757946479, 1551204074, 681266718,
200]);
201
202/// A default Poseidon2 for KoalaBear using the round constants from the original specification.
203///
204/// See Poseidon paper for more details: https://eprint.iacr.org/2019/458
205/// Sage script used to generate these constants: https://github.com/HorizenLabs/poseidon2/blob/main/poseidon2_rust_params.sage
206pub fn default_koalabear_poseidon2_24() -> Poseidon2KoalaBear<24> {
207    Poseidon2::new(
208        ExternalLayerConstants::new(
209            KOALABEAR_RC24_EXTERNAL_INITIAL.to_vec(),
210            KOALABEAR_RC24_EXTERNAL_FINAL.to_vec(),
211        ),
212        KOALABEAR_RC24_INTERNAL.to_vec(),
213    )
214}
215
216/// Contains data needed to define the internal layers of the Poseidon2 permutation.
217#[derive(Debug, Clone, Default)]
218pub struct KoalaBearInternalLayerParameters;
219
220impl InternalLayerBaseParameters<KoalaBearParameters, 16> for KoalaBearInternalLayerParameters {
221    /// Perform the internal matrix multiplication: s -> (1 + Diag(V))s.
222    /// We ignore `state[0]` as it is handled separately.
223    fn internal_layer_mat_mul<R: PrimeCharacteristicRing>(state: &mut [R; 16], sum: R) {
224        // The diagonal matrix is defined by the vector:
225        // V = [-2, 1, 2, 1/2, 3, 4, -1/2, -3, -4, 1/2^8, 1/8, 1/2^24, -1/2^8, -1/8, -1/16, -1/2^24]
226        state[1] += sum.clone();
227        state[2] = state[2].double() + sum.clone();
228        state[3] = state[3].halve() + sum.clone();
229        state[4] = sum.clone() + state[4].double() + state[4].clone();
230        state[5] = sum.clone() + state[5].double().double();
231        state[6] = sum.clone() - state[6].halve();
232        state[7] = sum.clone() - (state[7].double() + state[7].clone());
233        state[8] = sum.clone() - state[8].double().double();
234        state[9] = state[9].div_2exp_u64(8);
235        state[9] += sum.clone();
236        state[10] = state[10].div_2exp_u64(3);
237        state[10] += sum.clone();
238        state[11] = state[11].div_2exp_u64(24);
239        state[11] += sum.clone();
240        state[12] = state[12].div_2exp_u64(8);
241        state[12] = sum.clone() - state[12].clone();
242        state[13] = state[13].div_2exp_u64(3);
243        state[13] = sum.clone() - state[13].clone();
244        state[14] = state[14].div_2exp_u64(4);
245        state[14] = sum.clone() - state[14].clone();
246        state[15] = state[15].div_2exp_u64(24);
247        state[15] = sum - state[15].clone();
248    }
249}
250
251impl InternalLayerBaseParameters<KoalaBearParameters, 24> for KoalaBearInternalLayerParameters {
252    /// Perform the internal matrix multiplication: s -> (1 + Diag(V))s.
253    /// We ignore `state[0]` as it is handled separately.
254    fn internal_layer_mat_mul<R: PrimeCharacteristicRing>(state: &mut [R; 24], sum: R) {
255        // The diagonal matrix is defined by the vector:
256        // V = [-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^24, -1/2^8, -1/8, -1/16, -1/32, -1/64, -1/2^7, -1/2^9, -1/2^24]
257        state[1] += sum.clone();
258        state[2] = state[2].double() + sum.clone();
259        state[3] = state[3].halve() + sum.clone();
260        state[4] = sum.clone() + state[4].double() + state[4].clone();
261        state[5] = sum.clone() + state[5].double().double();
262        state[6] = sum.clone() - state[6].halve();
263        state[7] = sum.clone() - (state[7].double() + state[7].clone());
264        state[8] = sum.clone() - state[8].double().double();
265        state[9] = state[9].div_2exp_u64(8);
266        state[9] += sum.clone();
267        state[10] = state[10].div_2exp_u64(2);
268        state[10] += sum.clone();
269        state[11] = state[11].div_2exp_u64(3);
270        state[11] += sum.clone();
271        state[12] = state[12].div_2exp_u64(4);
272        state[12] += sum.clone();
273        state[13] = state[13].div_2exp_u64(5);
274        state[13] += sum.clone();
275        state[14] = state[14].div_2exp_u64(6);
276        state[14] += sum.clone();
277        state[15] = state[15].div_2exp_u64(24);
278        state[15] += sum.clone();
279        state[16] = state[16].div_2exp_u64(8);
280        state[16] = sum.clone() - state[16].clone();
281        state[17] = state[17].div_2exp_u64(3);
282        state[17] = sum.clone() - state[17].clone();
283        state[18] = state[18].div_2exp_u64(4);
284        state[18] = sum.clone() - state[18].clone();
285        state[19] = state[19].div_2exp_u64(5);
286        state[19] = sum.clone() - state[19].clone();
287        state[20] = state[20].div_2exp_u64(6);
288        state[20] = sum.clone() - state[20].clone();
289        state[21] = state[21].div_2exp_u64(7);
290        state[21] = sum.clone() - state[21].clone();
291        state[22] = state[22].div_2exp_u64(9);
292        state[22] = sum.clone() - state[22].clone();
293        state[23] = state[23].div_2exp_u64(24);
294        state[23] = sum - state[23].clone();
295    }
296}
297
298impl InternalLayerParameters<KoalaBearParameters, 16> for KoalaBearInternalLayerParameters {}
299impl InternalLayerParameters<KoalaBearParameters, 24> for KoalaBearInternalLayerParameters {}
300
301#[cfg(test)]
302mod tests {
303    use p3_symmetric::Permutation;
304    use rand::{Rng, SeedableRng};
305    use rand_xoshiro::Xoroshiro128Plus;
306
307    use super::*;
308
309    type F = KoalaBear;
310
311    // We need to make some round constants. We use Xoroshiro128Plus for this as we can easily match this PRNG in sage.
312    // See: https://github.com/0xPolygonZero/hash-constants for the sage code used to create all these tests.
313
314    /// Test on a roughly random input.
315    /// This random input is generated by the following sage code:
316    /// set_random_seed(16)
317    /// vector([KB.random_element() for t in range(16)]).
318    #[test]
319    fn test_poseidon2_width_16_random() {
320        let mut input: [F; 16] = KoalaBear::new_array([
321            894848333, 1437655012, 1200606629, 1690012884, 71131202, 1749206695, 1717947831,
322            120589055, 19776022, 42382981, 1831865506, 724844064, 171220207, 1299207443, 227047920,
323            1783754913,
324        ]);
325
326        let expected: [F; 16] = KoalaBear::new_array([
327            652590279, 1200629963, 1013089423, 1840372851, 19101828, 561050015, 1714865585,
328            994637181, 498949829, 729884572, 1957973925, 263012103, 535029297, 2121808603,
329            964663675, 1473622080,
330        ]);
331
332        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
333        let perm = Poseidon2KoalaBear::new_from_rng_128(&mut rng);
334
335        perm.permute_mut(&mut input);
336        assert_eq!(input, expected);
337    }
338
339    /// Test on a roughly random input.
340    /// This random input is generated by the following sage code:
341    /// set_random_seed(24)
342    /// vector([KB.random_element() for t in range(24)]).
343    #[test]
344    fn test_poseidon2_width_24_random() {
345        let mut input: [F; 24] = KoalaBear::new_array([
346            886409618, 1327899896, 1902407911, 591953491, 648428576, 1844789031, 1198336108,
347            355597330, 1799586834, 59617783, 790334801, 1968791836, 559272107, 31054313,
348            1042221543, 474748436, 135686258, 263665994, 1962340735, 1741539604, 2026927696,
349            449439011, 1131357108, 50869465,
350        ]);
351
352        let expected: [F; 24] = KoalaBear::new_array([
353            3825456, 486989921, 613714063, 282152282, 1027154688, 1171655681, 879344953,
354            1090688809, 1960721991, 1604199242, 1329947150, 1535171244, 781646521, 1156559780,
355            1875690339, 368140677, 457503063, 304208551, 1919757655, 835116474, 1293372648,
356            1254825008, 810923913, 1773631109,
357        ]);
358
359        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
360        let perm = Poseidon2KoalaBear::new_from_rng_128(&mut rng);
361
362        perm.permute_mut(&mut input);
363        assert_eq!(input, expected);
364    }
365
366    /// Test the generic internal layer against the optimized internal layer
367    /// for a random input of width 16.
368    #[test]
369    fn test_generic_internal_linear_layer_16() {
370        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
371        let mut input1: [F; 16] = rng.random();
372        let mut input2 = input1;
373
374        let part_sum: F = input1[1..].iter().copied().sum();
375        let full_sum = part_sum + input1[0];
376
377        input1[0] = part_sum - input1[0];
378
379        KoalaBearInternalLayerParameters::internal_layer_mat_mul(&mut input1, full_sum);
380        KoalaBearInternalLayerParameters::generic_internal_linear_layer(&mut input2);
381
382        assert_eq!(input1, input2);
383    }
384
385    /// Test the generic internal layer against the optimized internal layer
386    /// for a random input of width 16.
387    #[test]
388    fn test_generic_internal_linear_layer_24() {
389        let mut rng = Xoroshiro128Plus::seed_from_u64(1);
390        let mut input1: [F; 24] = rng.random();
391        let mut input2 = input1;
392
393        let part_sum: F = input1[1..].iter().copied().sum();
394        let full_sum = part_sum + input1[0];
395
396        input1[0] = part_sum - input1[0];
397
398        KoalaBearInternalLayerParameters::internal_layer_mat_mul(&mut input1, full_sum);
399        KoalaBearInternalLayerParameters::generic_internal_linear_layer(&mut input2);
400
401        assert_eq!(input1, input2);
402    }
403}