zip_plus/code/
iprs.rs

1mod pntt;
2
3use crate::{ZipError, code::LinearCode, pcs::structs::ZipTypes};
4use crypto_primitives::{
5    BaseFieldConfig, ProjectElementWithConfig, ProjectPrimitiveIntegersWithConfig,
6};
7use num_traits::{CheckedAdd, CheckedMul};
8use pntt::radix8::params::Config as PnttConfig;
9pub use pntt::radix8::params::{PnttConfigF65537, PnttInt, Radix8PnttParams};
10use std::{
11    fmt::Debug,
12    iter::Sum,
13    marker::PhantomData,
14    ops::{Add, AddAssign},
15};
16use zinc_utils::{add, from_ref::FromRef, mul_by_scalar::MulByScalar};
17
18/// Pseudo Reed-Solomon encoder over the integers. Internally uses a
19/// radix-8 NTT-style recursion with a base Vandermonde matrix sized
20/// `base_len x base_dim` (defaults to 64x32).
21#[derive(Clone)]
22pub struct IprsCode<Zt: ZipTypes, Config: PnttConfig, const REP: usize, const CHECK: bool> {
23    pntt_params: Radix8PnttParams<Config>,
24    _phantom: PhantomData<Zt>,
25}
26
27impl<Zt, Config, const REP: usize, const CHECK: bool> IprsCode<Zt, Config, REP, CHECK>
28where
29    Zt: ZipTypes,
30    Config: PnttConfig,
31{
32    pub fn new(row_len: usize, depth: usize) -> Result<Self, ZipError> {
33        // TODO(alex): Calculate max expected Zt::Cw::COEFF_BIT_WIDTH to ensure in
34        //             advance that the encoding will not overflow
35        Ok(Self {
36            pntt_params: Radix8PnttParams::new(row_len, depth, REP)?,
37            _phantom: Default::default(),
38        })
39    }
40
41    /// Create a new IPRS code with the optimal depth heuristics trying to keep
42    /// number of columns in the base matrix small.
43    /// Currently, keeps number of columns <= 2^8 but this might be tweaked in
44    /// the future.
45    pub fn new_with_optimal_depth(row_len: usize) -> Result<Self, ZipError> {
46        const MAX_BASE_COLS_LOG2: usize = 6;
47
48        let target_base_len = 1 << MAX_BASE_COLS_LOG2;
49        // We want depth to be at least 1.
50        let depth = 1.max(((1.max(row_len / target_base_len)).ilog2() as usize).div_ceil(3));
51
52        Self::new(row_len, depth)
53    }
54
55    /// Encode without modular reduction, purely over the integers.
56    fn encode_inner<In, Out>(&self, row: &[In]) -> Vec<Out>
57    where
58        In: MulByScalar<PnttInt, Out> + Clone + Send + Sync,
59        Out: CheckedAdd
60            + for<'a> AddAssign<&'a Out>
61            + for<'a> Add<&'a Out, Output = Out>
62            + CheckedMul
63            + MulByScalar<PnttInt>
64            + Sum
65            + FromRef<In>
66            + Clone
67            + Debug
68            + Send
69            + Sync,
70    {
71        assert_eq!(
72            row.len(),
73            self.pntt_params.row_len,
74            "Input length {} does not match expected row length {}",
75            row.len(),
76            self.pntt_params.row_len,
77        );
78
79        macro_rules! mul_fn {
80            () => {
81                |v: &_, tw: &PnttInt| {
82                    v.clone()
83                        .mul_by_scalar::<CHECK>(tw)
84                        .expect("Multiplication by twiddle should not overflow")
85                }
86            };
87        }
88
89        #[allow(clippy::arithmetic_side_effects)] // intended
90        let add_fn = |mut a: Out, b: &Out| {
91            if CHECK {
92                add!(a, b)
93            } else {
94                a += b;
95                a
96            }
97        };
98
99        pntt::radix8::pntt(
100            row,
101            &self.pntt_params,
102            Out::from_ref,
103            mul_fn!(),
104            mul_fn!(),
105            add_fn,
106        )
107    }
108
109    // Do the encoding but make use of the fact
110    // that we are dealing with a field.
111    fn encode_inner_f<C>(&self, cfg: &C, row: &[C::Element]) -> Vec<C::Element>
112    where
113        C: BaseFieldConfig + ProjectElementWithConfig<PnttInt>,
114    {
115        assert_eq!(
116            row.len(),
117            self.pntt_params.row_len,
118            "Input length {} does not match expected row length {}",
119            row.len(),
120            self.pntt_params.row_len,
121        );
122
123        let mul_fn = |f: &C::Element, tw: &PnttInt| cfg.mul(f, &cfg.project(tw));
124        let add_fn = |mut a: C::Element, b: &C::Element| {
125            cfg.add_assign(&mut a, b);
126            a
127        };
128
129        pntt::radix8::pntt(
130            row,
131            &self.pntt_params,
132            |v| v.clone(),
133            mul_fn,
134            mul_fn,
135            add_fn,
136        )
137    }
138}
139
140impl<Zt: ZipTypes, Config, const REP: usize, const CHECK: bool> LinearCode<Zt>
141    for IprsCode<Zt, Config, REP, CHECK>
142where
143    Zt: ZipTypes,
144    Config: PnttConfig,
145    Zt::Eval: MulByScalar<PnttInt, Zt::Cw>,
146    Zt::Cw: MulByScalar<PnttInt> + CheckedAdd,
147    Zt::CombR: MulByScalar<PnttInt>,
148{
149    const REPETITION_FACTOR: usize = REP;
150
151    fn encode(&self, row: &[Zt::Eval]) -> Vec<Zt::Cw> {
152        assert_eq!(
153            row.len(),
154            self.pntt_params.row_len,
155            "Input length {} does not match expected row length {}",
156            row.len(),
157            self.pntt_params.row_len,
158        );
159
160        self.encode_inner(row)
161    }
162
163    fn row_len(&self) -> usize {
164        self.pntt_params.row_len
165    }
166
167    fn codeword_len(&self) -> usize {
168        self.pntt_params.codeword_len
169    }
170
171    fn params_string(&self) -> String {
172        format!(
173            "row_len={}, rate=1/{REP}, depth={}",
174            self.row_len(),
175            self.pntt_params.depth
176        )
177    }
178
179    fn encode_wide(&self, row: &[Zt::CombR]) -> Vec<Zt::CombR> {
180        self.encode_inner(row)
181    }
182
183    fn encode_f<C>(&self, cfg: &C, row: &[C::Element]) -> Vec<C::Element>
184    where
185        C: BaseFieldConfig + ProjectPrimitiveIntegersWithConfig,
186    {
187        self.encode_inner_f(cfg, row)
188    }
189}
190
191impl<Zt, Config, const REP: usize, const CHECK: bool> Debug for IprsCode<Zt, Config, REP, CHECK>
192where
193    Zt: ZipTypes,
194    Config: PnttConfig,
195{
196    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
197        f.debug_struct("IprsCode")
198            .field("pntt_params", &self.pntt_params)
199            .finish()
200    }
201}
202
203impl<Zt, Config, const REP: usize, const CHECK: bool> PartialEq for IprsCode<Zt, Config, REP, CHECK>
204where
205    Config: PnttConfig,
206    Zt: ZipTypes,
207{
208    fn eq(&self, other: &Self) -> bool {
209        self.pntt_params == other.pntt_params
210    }
211}
212
213impl<Zt, Config, const REP: usize, const CHECK: bool> Eq for IprsCode<Zt, Config, REP, CHECK>
214where
215    Zt: ZipTypes,
216    Config: PnttConfig,
217{
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223    use crate::pcs::{structs::ZipPlus, test_utils::*};
224    use crypto_bigint::U64;
225    use crypto_primitives::{
226        Semiring, boolean::Boolean, crypto_bigint_int::Int, crypto_bigint_uint::Uint,
227    };
228    use rand::{
229        distr::{Distribution, StandardUniform},
230        prelude::ThreadRng,
231    };
232    use zinc_poly::{
233        mle::{DenseMultilinearExtension, MultilinearExtensionRand},
234        univariate::{
235            binary::{BinaryPoly, BinaryPolyInnerProduct},
236            dense::{DensePolyInnerProduct, DensePolynomial},
237        },
238    };
239    use zinc_primality::MillerRabin;
240    use zinc_transcript::traits::ConstTranscribable;
241    use zinc_utils::{
242        CHECKED,
243        inner_product::{MBSInnerProduct, ScalarProduct},
244        named::Named,
245    };
246
247    const INT_LIMBS: usize = U64::LIMBS;
248    const N: usize = INT_LIMBS;
249    const K: usize = INT_LIMBS * 4;
250    const M: usize = INT_LIMBS * 8;
251    type Zt = TestZipTypes<N, K, M>;
252
253    type Code = IprsCode<Zt, PnttConfigF65537, REP_FACTOR, CHECKED>;
254
255    #[test]
256    fn new_with_different_params() {
257        assert!(Code::new(1, 0).is_ok());
258        assert!(Code::new(8, 0).is_ok());
259        assert!(Code::new(1, 1).is_err());
260        assert!(Code::new(8, 1).is_ok());
261
262        assert!(Code::new_with_optimal_depth(1).is_err());
263        assert!(Code::new_with_optimal_depth(8).is_ok());
264        assert!(Code::new_with_optimal_depth(12).is_err());
265        assert!(Code::new_with_optimal_depth(16).is_ok());
266    }
267
268    fn do_encode<Zt, const REP: usize>(num_vars: usize)
269    where
270        Zt: ZipTypes,
271        Zt::Eval: MulByScalar<PnttInt, Zt::Cw>,
272        Zt::Cw: MulByScalar<PnttInt> + CheckedAdd,
273        Zt::CombR: MulByScalar<PnttInt>,
274        StandardUniform: Distribution<Zt::Eval>,
275    {
276        let mut rng = ThreadRng::default();
277        let poly_size: usize = 1 << num_vars;
278        let mle = DenseMultilinearExtension::rand(num_vars, &mut rng);
279
280        let code = IprsCode::<Zt, PnttConfigF65537, 4, CHECKED>::new_with_optimal_depth(poly_size)
281            .unwrap();
282        let pp = ZipPlus::setup(poly_size, code);
283        ZipPlus::<Zt, _>::encode_rows(&pp, &mle.evaluations);
284    }
285
286    /// Test the widest integer encoding used in benchmarks
287    #[test]
288    #[cfg_attr(miri, ignore)] // long running
289    fn encode_bench_int() {
290        #[derive(Clone, Debug)]
291        struct BenchZipTypes {}
292        impl ZipTypes for BenchZipTypes {
293            const NUM_COLUMN_OPENINGS: usize = 100;
294            type Eval = i32;
295            type Cw = i128;
296            type Fmod = Uint<{ INT_LIMBS * 4 }>;
297            type PrimeTest = MillerRabin;
298            type Chal = i128;
299            type Pt = i128;
300            type CombR = Int<{ INT_LIMBS * 3 }>;
301            type Comb = Self::CombR;
302            type EvalDotChal = ScalarProduct;
303            type CombDotChal = ScalarProduct;
304            type ArrCombRDotChal = MBSInnerProduct;
305        }
306
307        do_encode::<BenchZipTypes, 4>(14);
308    }
309
310    /// Test the widest binary polynomial encoding used in benchmarks
311    #[test]
312    #[cfg_attr(miri, ignore)] // long running
313    fn encode_bench_poly() {
314        const D_PLUS_ONE: usize = 32;
315
316        #[derive(Clone, Debug)]
317        struct BenchZipPlusTypes<CwCoeff>(PhantomData<CwCoeff>);
318        impl<CwCoeff> ZipTypes for BenchZipPlusTypes<CwCoeff>
319        where
320            CwCoeff: ConstTranscribable
321                + Copy
322                + Default
323                + FromRef<Boolean>
324                + Named
325                + Semiring
326                + Send
327                + Sync,
328            Int<5>: FromRef<CwCoeff>,
329        {
330            const NUM_COLUMN_OPENINGS: usize = 100;
331            type Eval = BinaryPoly<D_PLUS_ONE>;
332            type Cw = DensePolynomial<CwCoeff, D_PLUS_ONE>;
333            type Fmod = Uint<{ INT_LIMBS * 4 }>;
334            type PrimeTest = MillerRabin;
335            type Chal = i128;
336            type Pt = i128;
337            type CombR = Int<{ INT_LIMBS * 5 }>;
338            type Comb = DensePolynomial<Self::CombR, D_PLUS_ONE>;
339            type EvalDotChal = BinaryPolyInnerProduct<Self::Chal, D_PLUS_ONE>;
340            type CombDotChal = DensePolyInnerProduct<
341                (),
342                Self::CombR,
343                Self::Chal,
344                Self::CombR,
345                MBSInnerProduct,
346                D_PLUS_ONE,
347            >;
348            type ArrCombRDotChal = MBSInnerProduct;
349        }
350
351        do_encode::<BenchZipPlusTypes<i64>, 4>(12);
352    }
353}
zip_plus/code/iprs.rs

zip_plus/code/
iprs.rs
