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}