p3_binary_dft/
encoder.rs

1//! Reed–Solomon encoder over the additive NTT domain.
2
3use core::marker::PhantomData;
4
5use p3_binary_field::BinaryField128;
6use p3_commit::Encoder;
7use p3_field::PrimeCharacteristicRing;
8use p3_matrix::Matrix;
9use p3_matrix::dense::RowMajorMatrix;
10use p3_maybe_rayon::prelude::*;
11use p3_util::log2_strict_usize;
12
13use crate::domain::domain_point;
14use crate::poly::PolyBasisNtt;
15use crate::traits::AdditiveNtt;
16
17/// Reed–Solomon over the additive NTT domain.
18///
19/// The message holds the low-index novel-basis coefficients of each column, so the codeword is
20/// the evaluation of `f̂(Ŵ_0(x), …, Ŵ_{k−1}(x))` on `S_{k + log_inv_rate}`.
21///
22/// The alphabet is `BinaryField128`, where [`PolyBasisNtt`] is the faster transform and falls
23/// back to [`LchNtt`](crate::LchNtt) on a target without a carryless multiply, so it is the default.
24///
25/// `F` is phantom: [`Encoder`] is only implemented below for `F = BinaryField128`, and stays
26/// that way as long as the alphabet is fixed (D9), so the parameter carries no other instance.
27#[derive(Clone, Debug, Default)]
28pub struct AdditiveRsEncoder<F, Ntt = PolyBasisNtt> {
29    ntt: Ntt,
30    _marker: PhantomData<F>,
31}
32
33/// The alphabet is fixed at `BinaryField128` (D9), as [`Encoder`] requires of every impl outside
34/// `p3-commit`'s blanket one.
35impl<Ntt: AdditiveNtt<BinaryField128> + Sync> Encoder<BinaryField128>
36    for AdditiveRsEncoder<BinaryField128, Ntt>
37{
38    fn encode_batch(
39        &self,
40        message: RowMajorMatrix<BinaryField128>,
41        log_inv_rate: usize,
42    ) -> RowMajorMatrix<BinaryField128> {
43        if log_inv_rate == 0 {
44            return self.ntt.ntt_batch(message);
45        }
46
47        let width = message.width();
48        let len = message.values.len();
49        let padded_len = u32::try_from(log_inv_rate)
50            .ok()
51            .and_then(|rate| len.checked_shl(rate))
52            // `checked_shl` only rejects a shift amount that is too wide; it does not detect
53            // the value itself overflowing, so recovering `len` from the shifted result is
54            // what actually proves no bits were lost.
55            .filter(|&padded| padded >> log_inv_rate == len)
56            .expect("codeword length overflows usize");
57        let log_message_height = log2_strict_usize(message.height());
58
59        // Zero-padding to `S_{k+r}` and transforming the whole codeword leaves the appended
60        // `hi` coefficient half zero at every stage `j >= k`, so those butterflies only ever
61        // replicate `lo` — yet the carryless multiply against that zero still runs, `r` stages
62        // deep. The equivalent computation is `2^r` independent height-`2^k` transforms of the
63        // unpadded message, one per coset `c` of `S_k` in `S_{k+r}`, evaluated at
64        // `domain_point(c << k)`: every codeword row decomposes as `c*2^k + m` with
65        // `domain_point(c*2^k + m) = domain_point(c << k) + domain_point(m)`, and the padding
66        // rows never enter the computation at all. Each coset is independent, so they run in
67        // parallel; `shifted_ntt_batch` takes its message by value, so each needs its own copy.
68        let mut values = BinaryField128::zero_vec(padded_len);
69        values
70            .par_chunks_mut(len)
71            .enumerate()
72            .for_each(|(c, chunk)| {
73                let shift = domain_point::<BinaryField128>(c << log_message_height);
74                let coset = self
75                    .ntt
76                    .shifted_ntt_batch(RowMajorMatrix::new(message.values.clone(), width), shift);
77                chunk.copy_from_slice(&coset.values);
78            });
79
80        RowMajorMatrix::new(values, width)
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use alloc::vec;
87
88    use p3_binary_field::{BinaryField128, TowerLevel};
89    use p3_commit::Encoder;
90    use p3_field::PrimeCharacteristicRing;
91    use p3_matrix::Matrix;
92    use p3_matrix::dense::RowMajorMatrix;
93    use proptest::prelude::*;
94    use rand::SeedableRng;
95    use rand::rngs::SmallRng;
96
97    use super::AdditiveRsEncoder;
98    use crate::naive::NaiveAdditiveNtt;
99    use crate::traits::AdditiveNtt;
100
101    type F = BinaryField128;
102
103    /// Builds a matrix whose entries are distinct functions of the seed and the position.
104    fn matrix(log_n: usize, width: usize, seed: u64) -> RowMajorMatrix<F> {
105        RowMajorMatrix::new(
106            (0..(width << log_n))
107                .map(|i| {
108                    let bits = seed
109                        .wrapping_mul(0x9e37_79b9_7f4a_7c15)
110                        .wrapping_add(i as u64);
111                    F::from_le_byte_iter(bits.to_le_bytes().into_iter().cycle())
112                })
113                .collect(),
114            width,
115        )
116    }
117
118    /// Encoding is zero-padding the novel-basis coefficients and evaluating on the whole domain.
119    #[test]
120    fn encodes_by_padding_and_transforming() {
121        let mut rng = SmallRng::seed_from_u64(1);
122        let message = RowMajorMatrix::<F>::rand(&mut rng, 1 << 5, 3);
123
124        let mut padded = message.clone();
125        padded.values.resize(message.values.len() * 4, F::ZERO);
126        let expected = NaiveAdditiveNtt::<F>::default().ntt_batch(padded);
127
128        let encoded = AdditiveRsEncoder::<F>::default().encode_batch(message, 2);
129        assert_eq!(encoded.height(), 1 << 7);
130        assert_eq!(encoded, expected);
131    }
132
133    /// The codeword restricted to the message-sized prefix is the message's own transform: the
134    /// correspondence Phase 3 folds along.
135    #[test]
136    fn codeword_prefix_is_the_message_transform() {
137        let mut rng = SmallRng::seed_from_u64(2);
138        let message = RowMajorMatrix::<F>::rand(&mut rng, 1 << 5, 2);
139
140        let encoded = AdditiveRsEncoder::<F>::default().encode_batch(message.clone(), 1);
141        let direct = AdditiveRsEncoder::<F>::default().encode_batch(message, 0);
142        assert_eq!(&encoded.values[..direct.values.len()], &direct.values[..]);
143    }
144
145    #[test]
146    #[should_panic = "codeword length overflows usize"]
147    fn encode_batch_panics_when_the_codeword_length_overflows() {
148        let message = RowMajorMatrix::new(vec![F::ZERO; 2], 1);
149        let _ = AdditiveRsEncoder::<F>::default().encode_batch(message, usize::BITS as usize - 1);
150    }
151
152    proptest! {
153        #![proptest_config(ProptestConfig::with_cases(32))]
154
155        /// The coset-wise encoder agrees with zero-padding and transforming, across the sizes
156        /// D8's coset decomposition actually branches on: `log_inv_rate = 0` (no cosets),
157        /// `log_n = 0` (single-row cosets), and ordinary cases in between.
158        #[test]
159        fn encode_batch_matches_padding_and_transforming(
160            log_n in 0usize..=6,
161            log_inv_rate in 0usize..=3,
162            width in 1usize..=4,
163            seed in any::<u64>(),
164        ) {
165            let message = matrix(log_n, width, seed);
166
167            let mut padded = message.clone();
168            padded
169                .values
170                .resize(message.values.len() << log_inv_rate, F::ZERO);
171            let expected = NaiveAdditiveNtt::<F>::default().ntt_batch(padded);
172
173            let encoded = AdditiveRsEncoder::<F>::default().encode_batch(message, log_inv_rate);
174            prop_assert_eq!(encoded, expected);
175        }
176    }
177}
p3_binary_dft/encoder.rs

p3_binary_dft/
encoder.rs
