p3_monty_31/
poseidon2.rs

Help
1use core::marker::PhantomData;
2use core::ops::Mul;
3
4use p3_field::FieldAlgebra;
5use p3_poseidon2::{
6    add_rc_and_sbox_generic, external_initial_permute_state, external_terminal_permute_state,
7    ExternalLayer, GenericPoseidon2LinearLayers, InternalLayer, MDSMat4,
8};
9
10use crate::{
11    FieldParameters, MontyField31, MontyParameters, Poseidon2ExternalLayerMonty31,
12    Poseidon2InternalLayerMonty31,
13};
14
15/// Trait which handles the Poseidon2 internal layers.
16///
17/// Everything needed to compute multiplication by a `WIDTH x WIDTH` diffusion matrix whose monty form is `1 + Diag(vec)`.
18/// vec is assumed to be of the form `[-2, ...]` with all entries after the first being small powers of `2`.
19pub trait InternalLayerBaseParameters<MP: MontyParameters, const WIDTH: usize>:
20    Clone + Sync
21{
22    // Most of the time, ArrayLike will be `[u8; WIDTH - 1]`.
23    type ArrayLike: AsRef<[MontyField31<MP>]> + Sized;
24
25    // Long term INTERNAL_DIAG_MONTY will be removed.
26    // Currently it is needed for the Packed field implementations.
27    const INTERNAL_DIAG_MONTY: [MontyField31<MP>; WIDTH];
28
29    /// Perform the internal matrix multiplication: s -> (1 + Diag(V))s.
30    /// We ignore `state[0]` as it is handled separately.
31    fn internal_layer_mat_mul(state: &mut [MontyField31<MP>; WIDTH], sum: MontyField31<MP>);
32
33    /// Perform the internal matrix multiplication for any Abstract field
34    /// which implements multiplication by MontyField31 elements.
35    fn generic_internal_linear_layer<FA: FieldAlgebra + Mul<MontyField31<MP>, Output = FA>>(
36        state: &mut [FA; WIDTH],
37    );
38}
39
40#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
41pub trait InternalLayerParameters<FP: FieldParameters, const WIDTH: usize>:
42    InternalLayerBaseParameters<FP, WIDTH>
43{
44}
45#[cfg(all(
46    target_arch = "x86_64",
47    target_feature = "avx2",
48    not(all(feature = "nightly-features", target_feature = "avx512f"))
49))]
50pub trait InternalLayerParameters<FP: FieldParameters, const WIDTH: usize>:
51    InternalLayerBaseParameters<FP, WIDTH> + crate::InternalLayerParametersAVX2<FP, WIDTH>
52{
53}
54#[cfg(all(
55    feature = "nightly-features",
56    target_arch = "x86_64",
57    target_feature = "avx512f"
58))]
59pub trait InternalLayerParameters<FP: FieldParameters, const WIDTH: usize>:
60    InternalLayerBaseParameters<FP, WIDTH> + crate::InternalLayerParametersAVX512<FP, WIDTH>
61{
62}
63#[cfg(not(any(
64    all(target_arch = "aarch64", target_feature = "neon"),
65    all(
66        target_arch = "x86_64",
67        target_feature = "avx2",
68        not(all(feature = "nightly-features", target_feature = "avx512f"))
69    ),
70    all(
71        feature = "nightly-features",
72        target_arch = "x86_64",
73        target_feature = "avx512f"
74    ),
75)))]
76pub trait InternalLayerParameters<FP: FieldParameters, const WIDTH: usize>:
77    InternalLayerBaseParameters<FP, WIDTH>
78{
79}
80
81impl<FP, const WIDTH: usize, P2P, const D: u64> InternalLayer<MontyField31<FP>, WIDTH, D>
82    for Poseidon2InternalLayerMonty31<FP, WIDTH, P2P>
83where
84    FP: FieldParameters,
85    P2P: InternalLayerParameters<FP, WIDTH>,
86{
87    /// Perform the internal layers of the Poseidon2 permutation on the given state.
88    fn permute_state(&self, state: &mut [MontyField31<FP>; WIDTH]) {
89        self.internal_constants.iter().for_each(|rc| {
90            state[0] += *rc;
91            state[0] = state[0].exp_const_u64::<D>();
92            let part_sum: MontyField31<FP> = state[1..].iter().cloned().sum();
93            let full_sum = part_sum + state[0];
94            state[0] = part_sum - state[0];
95            P2P::internal_layer_mat_mul(state, full_sum);
96        })
97    }
98}
99
100impl<FP, const WIDTH: usize, const D: u64> ExternalLayer<MontyField31<FP>, WIDTH, D>
101    for Poseidon2ExternalLayerMonty31<FP, WIDTH>
102where
103    FP: FieldParameters,
104{
105    /// Perform the initial external layers of the Poseidon2 permutation on the given state.
106    fn permute_state_initial(&self, state: &mut [MontyField31<FP>; WIDTH]) {
107        external_initial_permute_state(
108            state,
109            self.external_constants.get_initial_constants(),
110            add_rc_and_sbox_generic::<_, D>,
111            &MDSMat4,
112        );
113    }
114
115    /// Perform the terminal external layers of the Poseidon2 permutation on the given state.
116    fn permute_state_terminal(&self, state: &mut [MontyField31<FP>; WIDTH]) {
117        external_terminal_permute_state(
118            state,
119            self.external_constants.get_terminal_constants(),
120            add_rc_and_sbox_generic::<_, D>,
121            &MDSMat4,
122        );
123    }
124}
125
126/// An implementation of the matrix multiplications in the internal and external layers of Poseidon2.
127///
128/// This can act on `[FA; WIDTH]` for any FieldAlgebra which implements multiplication by `Monty<31>` field elements.
129/// This will usually be slower than the Poseidon2 permutation built from `Poseidon2InternalLayerMonty31` and
130/// `Poseidon2ExternalLayerMonty31` but it does work in more cases.
131pub struct GenericPoseidon2LinearLayersMonty31<FP, ILBP> {
132    _phantom1: PhantomData<FP>,
133    _phantom2: PhantomData<ILBP>,
134}
135
136impl<FP, FA, ILBP, const WIDTH: usize> GenericPoseidon2LinearLayers<FA, WIDTH>
137    for GenericPoseidon2LinearLayersMonty31<FP, ILBP>
138where
139    FP: FieldParameters,
140    FA: FieldAlgebra + Mul<MontyField31<FP>, Output = FA>,
141    ILBP: InternalLayerBaseParameters<FP, WIDTH>,
142{
143    /// Perform the external matrix multiplication for any Abstract field
144    /// which implements multiplication by MontyField31 elements.
145    fn internal_linear_layer(state: &mut [FA; WIDTH]) {
146        ILBP::generic_internal_linear_layer(state);
147    }
148}
p3_monty_31/poseidon2.rs

p3_monty_31/
poseidon2.rs
