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
16pub trait PartialRoundBaseParameters<MP: MontyParameters, const WIDTH: usize>:
28 Clone + Sync
29{
30 const USE_TEXTBOOK: bool = false;
35
36 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 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 for (s, &r) in state
89 .iter_mut()
90 .zip(self.internal_constants.textbook_residual.iter())
91 {
92 *s += r;
93 }
94 } else {
95 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
130pub 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}