p3_mds/
integrated_coset_mds.rs

1use alloc::vec::Vec;
2
3use p3_field::{Algebra, Field, Powers, TwoAdicField};
4use p3_symmetric::Permutation;
5use p3_util::{log2_strict_usize, reverse_slice_index_bits};
6
7use crate::MdsPermutation;
8use crate::butterflies::{bowers_g_layer, bowers_g_t_layer_integrated};
9
10/// Like `CosetMds`, with a few differences:
11/// - (Bit reversed, a la Bowers) DIF + DIT rather than DIT + DIF
12/// - We skip bit reversals of the inputs and outputs
13/// - We don't weight by `1/N`, since this doesn't affect the MDS property
14/// - We integrate the coset shifts into the DIF's twiddle factors
15#[derive(Clone, Debug)]
16pub struct IntegratedCosetMds<F, const N: usize> {
17    ifft_twiddles: Vec<F>,
18    fft_twiddles: Vec<Vec<F>>,
19}
20
21impl<F: TwoAdicField, const N: usize> Default for IntegratedCosetMds<F, N> {
22    fn default() -> Self {
23        let log_n = log2_strict_usize(N);
24        let root = F::two_adic_generator(log_n);
25        let root_inv = root.inverse();
26        let coset_shift = F::GENERATOR;
27
28        let mut ifft_twiddles = root_inv.powers().collect_n(N / 2);
29        reverse_slice_index_bits(&mut ifft_twiddles);
30
31        let fft_twiddles: Vec<Vec<F>> = (0..log_n)
32            .map(|layer| {
33                let shift_power = coset_shift.exp_power_of_2(layer);
34                let powers = Powers {
35                    base: root.exp_power_of_2(layer),
36                    current: shift_power,
37                };
38                let mut twiddles = powers.collect_n(N >> (layer + 1));
39                reverse_slice_index_bits(&mut twiddles);
40                twiddles
41            })
42            .collect();
43
44        Self {
45            ifft_twiddles,
46            fft_twiddles,
47        }
48    }
49}
50
51impl<F: Field, A: Algebra<F>, const N: usize> Permutation<[A; N]> for IntegratedCosetMds<F, N> {
52    fn permute(&self, mut input: [A; N]) -> [A; N] {
53        self.permute_mut(&mut input);
54        input
55    }
56
57    fn permute_mut(&self, values: &mut [A; N]) {
58        let log_n = log2_strict_usize(N);
59
60        // Bit-reversed DIF, aka Bowers G
61        for layer in 0..log_n {
62            bowers_g_layer(values, layer, &self.ifft_twiddles);
63        }
64
65        // Bit-reversed DIT, aka Bowers G^T
66        for layer in (0..log_n).rev() {
67            bowers_g_t_layer_integrated(values, layer, &self.fft_twiddles[layer]);
68        }
69    }
70}
71
72impl<F: Field, A: Algebra<F>, const N: usize> MdsPermutation<A, N> for IntegratedCosetMds<F, N> {}
73
74#[cfg(test)]
75mod tests {
76    use p3_baby_bear::BabyBear;
77    use p3_dft::{NaiveDft, TwoAdicSubgroupDft};
78    use p3_field::{Field, PrimeCharacteristicRing};
79    use p3_symmetric::Permutation;
80    use p3_util::reverse_slice_index_bits;
81    use rand::rngs::SmallRng;
82    use rand::{Rng, SeedableRng};
83
84    use crate::integrated_coset_mds::IntegratedCosetMds;
85
86    type F = BabyBear;
87    const N: usize = 16;
88
89    #[test]
90    fn matches_naive() {
91        let mut rng = SmallRng::seed_from_u64(1);
92        let mut arr: [F; N] = rng.random();
93
94        let mut arr_rev = arr.to_vec();
95        reverse_slice_index_bits(&mut arr_rev);
96
97        let shift = F::GENERATOR;
98        let mut coset_lde_naive = NaiveDft.coset_lde(arr_rev, 0, shift);
99        reverse_slice_index_bits(&mut coset_lde_naive);
100        coset_lde_naive
101            .iter_mut()
102            .for_each(|x| *x *= F::from_u8(N as u8));
103        IntegratedCosetMds::default().permute_mut(&mut arr);
104        assert_eq!(coset_lde_naive, arr);
105    }
106}
p3_mds/integrated_coset_mds.rs

p3_mds/
integrated_coset_mds.rs
