Skip to main content

p3_monty_31/
poseidon1.rs

1use core::marker::PhantomData;
2
3use p3_field::{Field, InjectiveMonomial};
4use p3_poseidon1::external::{
5    FullRoundLayer, full_round_initial_permute_state, full_round_terminal_permute_state,
6};
7use p3_poseidon1::generic::GenericPoseidon1LinearLayers;
8use p3_poseidon1::internal::{PartialRoundLayer, partial_permute_state};
9use p3_symmetric::Permutation;
10
11use crate::{
12    FieldParameters, MDSUtils, MdsMatrixMontyField31, MontyField31, MontyParameters,
13    Poseidon1ExternalLayerMonty31, Poseidon1InternalLayerMonty31, RelativelyPrimePower,
14};
15
16/// Trait for Poseidon1 partial round scalar operations.
17///
18/// Provides compile-time dispatch between two partial round strategies:
19///
20/// - **Sparse decomposition** (`USE_TEXTBOOK = false`, default): Uses the sparse matrix
21///   factorization from Appendix B of the Poseidon1 paper. Best for most field/width combos.
22///
23/// - **Textbook with scalar constants** (`USE_TEXTBOOK = true`): Keeps the fast MDS
24///   permutation (e.g., Karatsuba convolution) per round, but folds `state[1..WIDTH]`
25///   constants forward so only a scalar is added to `state[0]` each round. Best when
26///   the MDS is very fast (e.g., BabyBear width-16 with power-of-2 Karatsuba).
27pub trait PartialRoundBaseParameters<MP: MontyParameters, const WIDTH: usize>:
28    Clone + Sync
29{
30    /// Whether to use the textbook (MDS-per-round) path for partial rounds.
31    ///
32    /// - When `true`, the Karatsuba MDS is applied per round with scalar constants.
33    /// - When `false` (default), the sparse matrix decomposition is used.
34    const USE_TEXTBOOK: bool = false;
35
36    /// Apply the MDS permutation. Only called when `USE_TEXTBOOK` is `true`.
37    fn mds_permute(_state: &mut [MontyField31<MP>; WIDTH]) {}
38}
39
40#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
41pub trait PartialRoundParameters<FP: FieldParameters, const WIDTH: usize>:
42    PartialRoundBaseParameters<FP, WIDTH> + crate::PartialRoundParametersNeon<FP, WIDTH>
43{
44}
45#[cfg(all(
46    target_arch = "x86_64",
47    target_feature = "avx2",
48    not(target_feature = "avx512f")
49))]
50pub trait PartialRoundParameters<FP: FieldParameters, const WIDTH: usize>:
51    PartialRoundBaseParameters<FP, WIDTH> + crate::PartialRoundParametersAVX2<FP, WIDTH>
52{
53}
54#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]
55pub trait PartialRoundParameters<FP: FieldParameters, const WIDTH: usize>:
56    PartialRoundBaseParameters<FP, WIDTH> + crate::PartialRoundParametersAVX512<FP, WIDTH>
57{
58}
59#[cfg(not(any(
60    all(target_arch = "aarch64", target_feature = "neon"),
61    all(
62        target_arch = "x86_64",
63        target_feature = "avx2",
64        not(target_feature = "avx512f")
65    ),
66    all(target_arch = "x86_64", target_feature = "avx512f"),
67)))]
68pub trait PartialRoundParameters<FP: FieldParameters, const WIDTH: usize>:
69    PartialRoundBaseParameters<FP, WIDTH>
70{
71}
72
73impl<FP, const WIDTH: usize, P1P, const D: u64> PartialRoundLayer<MontyField31<FP>, WIDTH, D>
74    for Poseidon1InternalLayerMonty31<FP, WIDTH, P1P>
75where
76    FP: FieldParameters + RelativelyPrimePower<D>,
77    P1P: PartialRoundParameters<FP, WIDTH>,
78{
79    fn permute_state(&self, state: &mut [MontyField31<FP>; WIDTH]) {
80        if P1P::USE_TEXTBOOK {
81            // Textbook: scalar constant + S-box + Karatsuba MDS per round.
82            for &c in &self.internal_constants.textbook_scalar_constants {
83                state[0] += c;
84                state[0] = InjectiveMonomial::<D>::injective_exp_n(&state[0]);
85                P1P::mds_permute(state);
86            }
87            // Add residual after all partial rounds.
88            for (s, &r) in state
89                .iter_mut()
90                .zip(self.internal_constants.textbook_residual.iter())
91            {
92                *s += r;
93            }
94        } else {
95            // Sparse decomposition (default).
96            partial_permute_state::<MontyField31<FP>, MontyField31<FP>, WIDTH, D>(
97                state,
98                &self.internal_constants,
99            );
100        }
101    }
102}
103
104impl<FP, MU, const WIDTH: usize, const D: u64> FullRoundLayer<MontyField31<FP>, WIDTH, D>
105    for Poseidon1ExternalLayerMonty31<FP, MU, WIDTH>
106where
107    FP: FieldParameters + RelativelyPrimePower<D>,
108    MU: MDSUtils + Default,
109    MdsMatrixMontyField31<MU>: Permutation<[MontyField31<FP>; WIDTH]>,
110{
111    fn permute_state_initial(&self, state: &mut [MontyField31<FP>; WIDTH]) {
112        let mds = MdsMatrixMontyField31::<MU>::default();
113        full_round_initial_permute_state::<MontyField31<FP>, MontyField31<FP>, _, WIDTH, D>(
114            state,
115            &self.external_constants,
116            &mds,
117        );
118    }
119
120    fn permute_state_terminal(&self, state: &mut [MontyField31<FP>; WIDTH]) {
121        let mds = MdsMatrixMontyField31::<MU>::default();
122        full_round_terminal_permute_state::<MontyField31<FP>, MontyField31<FP>, _, WIDTH, D>(
123            state,
124            &self.external_constants,
125            &mds,
126        );
127    }
128}
129
130/// Generic Poseidon1 linear layers for MontyField31.
131pub struct GenericPoseidon1LinearLayersMonty31<FP, PRBP> {
132    _phantom1: PhantomData<FP>,
133    _phantom2: PhantomData<PRBP>,
134}
135
136impl<FP, PRBP, F, const WIDTH: usize> GenericPoseidon1LinearLayers<F, WIDTH>
137    for GenericPoseidon1LinearLayersMonty31<FP, PRBP>
138where
139    FP: FieldParameters,
140    PRBP: PartialRoundBaseParameters<FP, WIDTH>,
141    F: Field,
142{
143}