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_mut(&self, values: &mut [A; N]) {
53        let log_n = log2_strict_usize(N);
54
55        // Bit-reversed DIF, aka Bowers G
56        for layer in 0..log_n {
57            bowers_g_layer(values, layer, &self.ifft_twiddles);
58        }
59
60        // Bit-reversed DIT, aka Bowers G^T
61        for layer in (0..log_n).rev() {
62            bowers_g_t_layer_integrated(values, layer, &self.fft_twiddles[layer]);
63        }
64    }
65}
66
67impl<F: Field, A: Algebra<F>, const N: usize> MdsPermutation<A, N> for IntegratedCosetMds<F, N> {}
68
69#[cfg(test)]
70mod tests {
71    use p3_baby_bear::BabyBear;
72    use p3_dft::{NaiveDft, TwoAdicSubgroupDft};
73    use p3_field::{Field, PrimeCharacteristicRing};
74    use p3_symmetric::Permutation;
75    use p3_util::reverse_slice_index_bits;
76    use rand::rngs::SmallRng;
77    use rand::{Rng, SeedableRng};
78
79    use crate::integrated_coset_mds::IntegratedCosetMds;
80
81    type F = BabyBear;
82    const N: usize = 16;
83
84    #[test]
85    fn matches_naive() {
86        let mut rng = SmallRng::seed_from_u64(1);
87        let mut arr: [F; N] = rng.random();
88
89        let mut arr_rev = arr.to_vec();
90        reverse_slice_index_bits(&mut arr_rev);
91
92        let shift = F::GENERATOR;
93        let mut coset_lde_naive = NaiveDft.coset_lde(arr_rev, 0, shift);
94        reverse_slice_index_bits(&mut coset_lde_naive);
95        coset_lde_naive
96            .iter_mut()
97            .for_each(|x| *x *= F::from_u8(N as u8));
98        IntegratedCosetMds::default().permute_mut(&mut arr);
99        assert_eq!(coset_lde_naive, arr);
100    }
101}
p3_mds/integrated_coset_mds.rs

p3_mds/
integrated_coset_mds.rs
