p3_dft/
butterflies.rs

1use p3_field::{Field, PackedField, PackedValue};
2
3pub trait Butterfly<F: Field>: Copy + Send + Sync {
4    fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF);
5
6    #[inline]
7    fn apply_in_place<PF: PackedField<Scalar = F>>(&self, x_1: &mut PF, x_2: &mut PF) {
8        (*x_1, *x_2) = self.apply(*x_1, *x_2);
9    }
10
11    #[inline]
12    fn apply_to_rows(&self, row_1: &mut [F], row_2: &mut [F]) {
13        let (shorts_1, suffix_1) = F::Packing::pack_slice_with_suffix_mut(row_1);
14        let (shorts_2, suffix_2) = F::Packing::pack_slice_with_suffix_mut(row_2);
15        assert_eq!(shorts_1.len(), shorts_2.len());
16        assert_eq!(suffix_1.len(), suffix_2.len());
17        for (x_1, x_2) in shorts_1.iter_mut().zip(shorts_2) {
18            self.apply_in_place(x_1, x_2);
19        }
20        for (x_1, x_2) in suffix_1.iter_mut().zip(suffix_2) {
21            self.apply_in_place(x_1, x_2);
22        }
23    }
24}
25
26#[derive(Copy, Clone)]
27pub struct DifButterfly<F>(pub F);
28impl<F: Field> Butterfly<F> for DifButterfly<F> {
29    #[inline]
30    fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF) {
31        (x_1 + x_2, (x_1 - x_2) * self.0)
32    }
33}
34
35#[derive(Copy, Clone)]
36pub struct DitButterfly<F>(pub F);
37impl<F: Field> Butterfly<F> for DitButterfly<F> {
38    #[inline]
39    fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF) {
40        let x_2_twiddle = x_2 * self.0;
41        (x_1 + x_2_twiddle, x_1 - x_2_twiddle)
42    }
43}
44
45/// Butterfly with twiddle factor 1 (works in either DIT or DIF).
46#[derive(Copy, Clone)]
47pub struct TwiddleFreeButterfly;
48impl<F: Field> Butterfly<F> for TwiddleFreeButterfly {
49    #[inline]
50    fn apply<PF: PackedField<Scalar = F>>(&self, x_1: PF, x_2: PF) -> (PF, PF) {
51        (x_1 + x_2, x_1 - x_2)
52    }
53}
p3_dft/butterflies.rs

p3_dft/
butterflies.rs
