p3_monty_31/aarch64_neon/
poseidon2.rs

1//! Eventually this will hold a Vectorized Neon implementation of Poseidon2 for MontyField31
2//! Currently this is essentially a placeholder to allow compilation on Neon devices.
3//!
4//! Converting the AVX2/AVX512 code across to Neon is on the TODO list.
5
6use alloc::vec::Vec;
7use core::marker::PhantomData;
8
9use p3_poseidon2::{
10    add_rc_and_sbox_generic, external_initial_permute_state, external_terminal_permute_state,
11    ExternalLayer, ExternalLayerConstants, ExternalLayerConstructor, InternalLayer,
12    InternalLayerConstructor, MDSMat4,
13};
14
15use crate::{
16    FieldParameters, InternalLayerBaseParameters, MontyField31, MontyParameters,
17    PackedMontyField31Neon,
18};
19
20/// The internal layers of the Poseidon2 permutation for Monty31 fields.
21///
22/// This is currently not optimized for the Neon architecture but this is on the TODO list.
23#[derive(Debug, Clone)]
24pub struct Poseidon2InternalLayerMonty31<
25    MP: MontyParameters,
26    const WIDTH: usize,
27    ILP: InternalLayerBaseParameters<MP, WIDTH>,
28> {
29    pub(crate) internal_constants: Vec<MontyField31<MP>>,
30    _phantom: PhantomData<ILP>,
31}
32
33/// The external layers of the Poseidon2 permutation for Monty31 fields.
34///
35/// This is currently not optimized for the Neon architecture but this is on the TODO list.
36#[derive(Debug, Clone)]
37pub struct Poseidon2ExternalLayerMonty31<MP: MontyParameters, const WIDTH: usize> {
38    pub(crate) external_constants: ExternalLayerConstants<MontyField31<MP>, WIDTH>,
39}
40
41impl<FP: FieldParameters, const WIDTH: usize, ILP: InternalLayerBaseParameters<FP, WIDTH>>
42    InternalLayerConstructor<PackedMontyField31Neon<FP>>
43    for Poseidon2InternalLayerMonty31<FP, WIDTH, ILP>
44{
45    fn new_from_constants(internal_constants: Vec<MontyField31<FP>>) -> Self {
46        Self {
47            internal_constants,
48            _phantom: PhantomData,
49        }
50    }
51}
52
53impl<FP: FieldParameters, const WIDTH: usize>
54    ExternalLayerConstructor<PackedMontyField31Neon<FP>, WIDTH>
55    for Poseidon2ExternalLayerMonty31<FP, WIDTH>
56{
57    fn new_from_constants(
58        external_constants: ExternalLayerConstants<MontyField31<FP>, WIDTH>,
59    ) -> Self {
60        Self { external_constants }
61    }
62}
63
64impl<FP, ILP, const WIDTH: usize, const D: u64> InternalLayer<PackedMontyField31Neon<FP>, WIDTH, D>
65    for Poseidon2InternalLayerMonty31<FP, WIDTH, ILP>
66where
67    FP: FieldParameters,
68    ILP: InternalLayerBaseParameters<FP, WIDTH>,
69{
70    /// Perform the internal layers of the Poseidon2 permutation on the given state.
71    fn permute_state(&self, state: &mut [PackedMontyField31Neon<FP>; WIDTH]) {
72        self.internal_constants.iter().for_each(|&rc| {
73            add_rc_and_sbox_generic::<_, D>(&mut state[0], rc);
74            ILP::generic_internal_linear_layer(state);
75        })
76    }
77}
78
79impl<FP, const D: u64, const WIDTH: usize> ExternalLayer<PackedMontyField31Neon<FP>, WIDTH, D>
80    for Poseidon2ExternalLayerMonty31<FP, WIDTH>
81where
82    FP: FieldParameters,
83{
84    /// Perform the initial external layers of the Poseidon2 permutation on the given state.
85    fn permute_state_initial(&self, state: &mut [PackedMontyField31Neon<FP>; WIDTH]) {
86        external_initial_permute_state(
87            state,
88            self.external_constants.get_initial_constants(),
89            add_rc_and_sbox_generic::<_, D>,
90            &MDSMat4,
91        );
92    }
93
94    /// Perform the terminal external layers of the Poseidon2 permutation on the given state.
95    fn permute_state_terminal(&self, state: &mut [PackedMontyField31Neon<FP>; WIDTH]) {
96        external_terminal_permute_state(
97            state,
98            self.external_constants.get_terminal_constants(),
99            add_rc_and_sbox_generic::<_, D>,
100            &MDSMat4,
101        );
102    }
103}