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::Matrix;
5use p3_matrix::dense::RowMajorMatrix;
6
7/// An additive NTT: evaluation of the novel polynomial basis on an `F_2`-linear subspace.
8///
9/// Columns hold coefficients in the novel polynomial basis `X_i = ∏_j Ŵ_j^{bit_j(i)}`, the
10/// products of normalised subspace polynomials; output row `i` is the evaluation at
11/// `shift + domain_point(i)`. Both sides are in natural order.
12///
13/// Heights are powers of two, and `ℓ = log2(height)` is at most the bit width `2^LOG_BITS` of
14/// `F`: `S_ℓ` is spanned by the first `ℓ` Cantor basis vectors, of which there are only that many.
15///
16/// No `Clone + Default` supertrait: an implementation carrying state — a precomputed twiddle
17/// table, a configurable grain, a non-Cantor subspace basis — should not be ruled out by the
18/// interface. [`AdditiveRsEncoder`](crate::AdditiveRsEncoder) requires those bounds itself.
19pub trait AdditiveNtt<F: TowerLevel> {
20    /// Evaluates each column on the coset `shift + S_ℓ`.
21    ///
22    /// # Panics
23    /// Panics if the height of `mat` is not a power of two, or if `ℓ` exceeds the bit width of
24    /// `F`, since `S_ℓ` then calls for a Cantor basis vector this level does not have.
25    fn shifted_ntt_batch(&self, mat: RowMajorMatrix<F>, shift: F) -> RowMajorMatrix<F>;
26
27    /// Inverse of [`shifted_ntt_batch`](Self::shifted_ntt_batch).
28    ///
29    /// # Panics
30    /// Panics under the same conditions as [`shifted_ntt_batch`](Self::shifted_ntt_batch).
31    fn shifted_intt_batch(&self, mat: RowMajorMatrix<F>, shift: F) -> RowMajorMatrix<F>;
32
33    /// Evaluates each column on `S_ℓ`.
34    fn ntt_batch(&self, mat: RowMajorMatrix<F>) -> RowMajorMatrix<F> {
35        self.shifted_ntt_batch(mat, F::ZERO)
36    }
37
38    /// Transforms a matrix whose coefficient prefix has been padded with zero rows.
39    ///
40    /// The caller supplies `2^log_inv_rate` times the original height and must leave
41    /// every entry after that original prefix zero. Implementations may ignore that tail.
42    /// The default evaluates the full matrix; optimized implementations can skip zero work.
43    ///
44    /// # Panics
45    /// Panics for an invalid transform height or if the padding exceeds that height.
46    fn ntt_batch_padded(&self, mat: RowMajorMatrix<F>, log_inv_rate: usize) -> RowMajorMatrix<F> {
47        let log_n = p3_util::log2_strict_usize(mat.height());
48        assert!(log_inv_rate <= log_n, "padding exceeds matrix height");
49        self.ntt_batch(mat)
50    }
51
52    /// Inverse of [`ntt_batch`](Self::ntt_batch).
53    fn intt_batch(&self, mat: RowMajorMatrix<F>) -> RowMajorMatrix<F> {
54        self.shifted_intt_batch(mat, F::ZERO)
55    }
56
57    /// Low-degree extension onto `S_{ℓ + added_bits} ⊃ S_ℓ`.
58    ///
59    /// Because `S_ℓ` is the index prefix of the larger domain, the input rows reappear as the
60    /// prefix of the output.
61    ///
62    /// # Panics
63    /// Panics unless `ℓ + added_bits` is at most the bit width of `F`, on top of the conditions
64    /// of [`shifted_ntt_batch`](Self::shifted_ntt_batch).
65    fn lde_batch(&self, mat: RowMajorMatrix<F>, added_bits: usize) -> RowMajorMatrix<F> {
66        self.shifted_lde_batch(mat, added_bits, F::ZERO)
67    }
68
69    /// [`lde_batch`](Self::lde_batch) on the coset `shift + S_{ℓ + added_bits}`.
70    fn shifted_lde_batch(
71        &self,
72        mat: RowMajorMatrix<F>,
73        added_bits: usize,
74        shift: F,
75    ) -> RowMajorMatrix<F> {
76        let log_n = p3_util::log2_strict_usize(mat.height());
77        assert!(log_n <= 1 << F::LOG_BITS, "domain exceeds field dimension");
78        if added_bits == 0 {
79            return mat;
80        }
81        let coeffs = self.shifted_intt_batch(mat, shift);
82        let width = coeffs.width;
83        let len = coeffs.values.len();
84        let padded_len = u32::try_from(added_bits)
85            .ok()
86            .and_then(|bits| len.checked_shl(bits))
87            // `checked_shl` only rejects a shift amount that is too wide; it does not detect
88            // the value itself overflowing, so recovering `len` from the shifted result is
89            // what actually proves no bits were lost.
90            .filter(|&padded| padded >> added_bits == len)
91            .expect("extended codeword length overflows usize");
92
93        // Zero-padding the coefficients keeps the same polynomial on the larger domain. `F::zero_vec`
94        // plus a `copy_from_slice` avoids `Vec::resize`'s reallocate-and-memcpy of the whole prefix.
95        let mut values = F::zero_vec(padded_len);
96        values[..len].copy_from_slice(&coeffs.values);
97        self.shifted_ntt_batch(RowMajorMatrix::new(values, width), shift)
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use alloc::vec;
104
105    use p3_binary_field::BinaryField8;
106    use p3_field::PrimeCharacteristicRing;
107    use p3_matrix::dense::RowMajorMatrix;
108
109    use super::AdditiveNtt;
110    use crate::LchNtt;
111
112    #[test]
113    fn identity_lde_preserves_input_allocation() {
114        let mat = RowMajorMatrix::new(vec![BinaryField8::ONE; 32], 4);
115        let ptr = mat.values.as_ptr();
116        let result = LchNtt::default().lde_batch(mat, 0);
117        assert_eq!(result.values.as_ptr(), ptr);
118        assert_eq!(result.values, vec![BinaryField8::ONE; 32]);
119    }
120
121    #[test]
122    #[should_panic]
123    fn identity_lde_rejects_invalid_height() {
124        let _ = LchNtt::default().lde_batch(RowMajorMatrix::new(vec![BinaryField8::ONE; 3], 1), 0);
125    }
126
127    #[test]
128    #[should_panic]
129    fn identity_lde_rejects_domain_above_field_dimension() {
130        let _ =
131            LchNtt::default().lde_batch(RowMajorMatrix::new(vec![BinaryField8::ONE; 512], 1), 0);
132    }
133}
p3_binary_dft/traits.rs

p3_binary_dft/
traits.rs
