Skip to main content

p3_binary_dft/
traits.rs

1//! The additive NTT interface shared by the reference and fast transforms.
2
3use p3_binary_field::TowerLevel;
4use p3_matrix::dense::RowMajorMatrix;
5
6/// An additive NTT: evaluation of the novel polynomial basis on an `F_2`-linear subspace.
7///
8/// Columns hold coefficients in the novel polynomial basis `X_i = ∏_j Ŵ_j^{bit_j(i)}`, the
9/// products of normalised subspace polynomials; output row `i` is the evaluation at
10/// `shift + domain_point(i)`. Both sides are in natural order.
11///
12/// Heights are powers of two, and `ℓ = log2(height)` is at most the bit width `2^LOG_BITS` of
13/// `F`: `S_ℓ` is spanned by the first `ℓ` Cantor basis vectors, of which there are only that many.
14///
15/// No `Clone + Default` supertrait: an implementation carrying state — a precomputed twiddle
16/// table, a configurable grain, a non-Cantor subspace basis — should not be ruled out by the
17/// interface. [`AdditiveRsEncoder`](crate::AdditiveRsEncoder) requires those bounds itself.
18pub trait AdditiveNtt<F: TowerLevel> {
19    /// Evaluates each column on the coset `shift + S_ℓ`.
20    ///
21    /// # Panics
22    /// Panics if the height of `mat` is not a power of two, or if `ℓ` exceeds the bit width of
23    /// `F`, since `S_ℓ` then calls for a Cantor basis vector this level does not have.
24    fn shifted_ntt_batch(&self, mat: RowMajorMatrix<F>, shift: F) -> RowMajorMatrix<F>;
25
26    /// Inverse of [`shifted_ntt_batch`](Self::shifted_ntt_batch).
27    ///
28    /// # Panics
29    /// Panics under the same conditions as [`shifted_ntt_batch`](Self::shifted_ntt_batch).
30    fn shifted_intt_batch(&self, mat: RowMajorMatrix<F>, shift: F) -> RowMajorMatrix<F>;
31
32    /// Evaluates each column on `S_ℓ`.
33    fn ntt_batch(&self, mat: RowMajorMatrix<F>) -> RowMajorMatrix<F> {
34        self.shifted_ntt_batch(mat, F::ZERO)
35    }
36
37    /// Inverse of [`ntt_batch`](Self::ntt_batch).
38    fn intt_batch(&self, mat: RowMajorMatrix<F>) -> RowMajorMatrix<F> {
39        self.shifted_intt_batch(mat, F::ZERO)
40    }
41
42    /// Low-degree extension onto `S_{ℓ + added_bits} ⊃ S_ℓ`.
43    ///
44    /// Because `S_ℓ` is the index prefix of the larger domain, the input rows reappear as the
45    /// prefix of the output.
46    ///
47    /// # Panics
48    /// Panics unless `ℓ + added_bits` is at most the bit width of `F`, on top of the conditions
49    /// of [`shifted_ntt_batch`](Self::shifted_ntt_batch).
50    fn lde_batch(&self, mat: RowMajorMatrix<F>, added_bits: usize) -> RowMajorMatrix<F> {
51        self.shifted_lde_batch(mat, added_bits, F::ZERO)
52    }
53
54    /// [`lde_batch`](Self::lde_batch) on the coset `shift + S_{ℓ + added_bits}`.
55    fn shifted_lde_batch(
56        &self,
57        mat: RowMajorMatrix<F>,
58        added_bits: usize,
59        shift: F,
60    ) -> RowMajorMatrix<F> {
61        let coeffs = self.shifted_intt_batch(mat, shift);
62        let width = coeffs.width;
63        let len = coeffs.values.len();
64        let padded_len = u32::try_from(added_bits)
65            .ok()
66            .and_then(|bits| len.checked_shl(bits))
67            // `checked_shl` only rejects a shift amount that is too wide; it does not detect
68            // the value itself overflowing, so recovering `len` from the shifted result is
69            // what actually proves no bits were lost.
70            .filter(|&padded| padded >> added_bits == len)
71            .expect("extended codeword length overflows usize");
72
73        // Zero-padding the coefficients keeps the same polynomial on the larger domain. `F::zero_vec`
74        // plus a `copy_from_slice` avoids `Vec::resize`'s reallocate-and-memcpy of the whole prefix.
75        let mut values = F::zero_vec(padded_len);
76        values[..len].copy_from_slice(&coeffs.values);
77        self.shifted_ntt_batch(RowMajorMatrix::new(values, width), shift)
78    }
79}