p3_dft/
traits.rs

Help
1use alloc::vec::Vec;
2
3use p3_field::TwoAdicField;
4use p3_matrix::bitrev::BitReversableMatrix;
5use p3_matrix::dense::RowMajorMatrix;
6use p3_matrix::util::swap_rows;
7use p3_matrix::Matrix;
8
9use crate::util::{coset_shift_cols, divide_by_height};
10
11pub trait TwoAdicSubgroupDft<F: TwoAdicField>: Clone + Default {
12    // Effectively this is either RowMajorMatrix or BitReversedMatrixView<RowMajorMatrix>.
13    // Always owned.
14    type Evaluations: BitReversableMatrix<F> + 'static;
15
16    /// Compute the discrete Fourier transform (DFT) `vec`.
17    fn dft(&self, vec: Vec<F>) -> Vec<F> {
18        self.dft_batch(RowMajorMatrix::new_col(vec))
19            .to_row_major_matrix()
20            .values
21    }
22
23    /// Compute the discrete Fourier transform (DFT) of each column in `mat`.
24    /// This is the only method an implementer needs to define, all other
25    /// methods can be derived from this one.
26    fn dft_batch(&self, mat: RowMajorMatrix<F>) -> Self::Evaluations;
27
28    /// Compute the "coset DFT" of `vec`. This can be viewed as interpolation onto a coset of a
29    /// multiplicative subgroup, rather than the subgroup itself.
30    fn coset_dft(&self, vec: Vec<F>, shift: F) -> Vec<F> {
31        self.coset_dft_batch(RowMajorMatrix::new_col(vec), shift)
32            .to_row_major_matrix()
33            .values
34    }
35
36    /// Compute the "coset DFT" of each column in `mat`. This can be viewed as interpolation onto a
37    /// coset of a multiplicative subgroup, rather than the subgroup itself.
38    fn coset_dft_batch(&self, mut mat: RowMajorMatrix<F>, shift: F) -> Self::Evaluations {
39        // Observe that
40        //     y_i = \sum_j c_j (s g^i)^j
41        //         = \sum_j (c_j s^j) (g^i)^j
42        // which has the structure of an ordinary DFT, except each coefficient c_j is first replaced
43        // by c_j s^j.
44        coset_shift_cols(&mut mat, shift);
45        self.dft_batch(mat)
46    }
47
48    /// Compute the inverse DFT of `vec`.
49    fn idft(&self, vec: Vec<F>) -> Vec<F> {
50        self.idft_batch(RowMajorMatrix::new(vec, 1)).values
51    }
52
53    /// Compute the inverse DFT of each column in `mat`.
54    fn idft_batch(&self, mat: RowMajorMatrix<F>) -> RowMajorMatrix<F> {
55        let mut dft = self.dft_batch(mat).to_row_major_matrix();
56        let h = dft.height();
57
58        divide_by_height(&mut dft);
59
60        for row in 1..h / 2 {
61            swap_rows(&mut dft, row, h - row);
62        }
63
64        dft
65    }
66
67    /// Compute the "coset iDFT" of `vec`. This can be viewed as an inverse operation of
68    /// "coset DFT", that interpolates over a coset of a multiplicative subgroup, rather than
69    /// subgroup itself.
70    fn coset_idft(&self, vec: Vec<F>, shift: F) -> Vec<F> {
71        self.coset_idft_batch(RowMajorMatrix::new(vec, 1), shift)
72            .values
73    }
74
75    /// Compute the "coset iDFT" of each column in `mat`. This can be viewed as an inverse operation
76    /// of "coset DFT", that interpolates over a coset of a multiplicative subgroup, rather than the
77    /// subgroup itself.
78    fn coset_idft_batch(&self, mut mat: RowMajorMatrix<F>, shift: F) -> RowMajorMatrix<F> {
79        mat = self.idft_batch(mat);
80        coset_shift_cols(&mut mat, shift.inverse());
81        mat
82    }
83
84    /// Compute the low-degree extension of `vec` onto a larger subgroup.
85    fn lde(&self, vec: Vec<F>, added_bits: usize) -> Vec<F> {
86        self.lde_batch(RowMajorMatrix::new(vec, 1), added_bits)
87            .to_row_major_matrix()
88            .values
89    }
90
91    /// Compute the low-degree extension of each column in `mat` onto a larger subgroup.
92    fn lde_batch(&self, mat: RowMajorMatrix<F>, added_bits: usize) -> Self::Evaluations {
93        let mut coeffs = self.idft_batch(mat);
94        coeffs
95            .values
96            .resize(coeffs.values.len() << added_bits, F::ZERO);
97        self.dft_batch(coeffs)
98    }
99
100    /// Compute the low-degree extension of each column in `mat` onto a coset of a larger subgroup.
101    fn coset_lde(&self, vec: Vec<F>, added_bits: usize, shift: F) -> Vec<F> {
102        self.coset_lde_batch(RowMajorMatrix::new(vec, 1), added_bits, shift)
103            .to_row_major_matrix()
104            .values
105    }
106
107    /// Compute the low-degree extension of each column in `mat` onto a coset of a larger subgroup.
108    fn coset_lde_batch(
109        &self,
110        mat: RowMajorMatrix<F>,
111        added_bits: usize,
112        shift: F,
113    ) -> Self::Evaluations {
114        let mut coeffs = self.idft_batch(mat);
115        // PANICS: possible panic if the new resized length overflows
116        coeffs.values.resize(
117            coeffs
118                .values
119                .len()
120                .checked_shl(added_bits.try_into().unwrap())
121                .unwrap(),
122            F::ZERO,
123        );
124        self.coset_dft_batch(coeffs, shift)
125    }
126}
p3_dft/traits.rs

p3_dft/
traits.rs
