zip_plus/code/
raa.rs

1use crate::{code::LinearCode, pcs::structs::ZipTypes, utils::shuffle_seeded};
2use crypto_primitives::{
3    BaseFieldConfig, FixedConfig, ProjectPrimitiveIntegersWithConfig, SemiringConfig,
4};
5use std::{fmt::Debug, marker::PhantomData};
6use zinc_poly::ConstCoeffBitWidth;
7use zinc_utils::{add, from_ref::FromRef, mul};
8
9pub trait RaaConfig: Copy + Send + Sync {
10    /// Whether to permute the codeword in place, instead of copying it using a
11    /// precomputed permutation.
12    const PERMUTE_IN_PLACE: bool;
13    /// Whether to check for overflows during encoding
14    // TODO: Unify with `CHECK_FOR_OVERFLOW` in `zinc_poly`
15    const CHECK_FOR_OVERFLOWS: bool;
16}
17
18/// Implementation of a repeat-accumulate-accumulate (RAA) codes over the binary
19/// field, as defined by the Blaze paper (https://eprint.iacr.org/2024/1609)
20#[derive(Clone)]
21pub struct RaaCode<Zt: ZipTypes, Config: RaaConfig, const REP: usize> {
22    pub(crate) row_len: usize,
23    /// Randomness seed for the first permutation
24    pub(crate) perm_1_seed: u64,
25
26    /// Randomness seed for the second permutation
27    pub(crate) perm_2_seed: u64,
28
29    /// First permutation
30    pub(crate) perm_1: Vec<usize>,
31
32    /// Second permutation
33    pub(crate) perm_2: Vec<usize>,
34
35    phantom: PhantomData<(Zt, Config)>,
36}
37
38impl<Zt: ZipTypes, Config: RaaConfig, const REP: usize> RaaCode<Zt, Config, REP> {
39    pub fn new(row_len: usize) -> Self {
40        assert!(
41            REP.is_power_of_two(),
42            "Repetition factor must be a power of two"
43        );
44        assert!(
45            row_len.is_power_of_two(),
46            "Row length must be a power of two"
47        );
48
49        // Width of each entry in codeword vector, in bits.
50        // For RAA it's initial_bits + 2*log2(codeword_len),
51        // where codeword_len = row_len * REP and the factor of 2
52        // comes from the two accumulation steps.
53        let codeword_width_bits = {
54            let initial_bits =
55                u32::try_from(Zt::Eval::COEFF_BIT_WIDTH).expect("Size of EvalR type is too large");
56
57            let row_len_log = row_len.ilog2();
58            let rep_factor_log = REP.ilog2();
59            add!(
60                initial_bits,
61                add!(mul!(row_len_log, 2), mul!(rep_factor_log, 2))
62            )
63        };
64        let codeword_type_bits =
65            u32::try_from(Zt::Cw::COEFF_BIT_WIDTH).expect("Size of CwR type is too large");
66        assert!(
67            codeword_type_bits >= codeword_width_bits,
68            "Cannot fit {codeword_width_bits}-bit wide codeword entries in {} bits entries",
69            codeword_type_bits
70        );
71
72        // We don't need a secure/unpredictable randomness here, so use fixed seeds
73        const PERM_1_SEED: u64 = 1;
74        const PERM_2_SEED: u64 = 2;
75
76        let codeword_len = mul!(row_len, REP);
77
78        let mut perm_1: Vec<usize> = (0..codeword_len).collect();
79        shuffle_seeded(&mut perm_1, PERM_1_SEED);
80        let mut perm_2: Vec<usize> = (0..codeword_len).collect();
81        shuffle_seeded(&mut perm_2, PERM_2_SEED);
82
83        Self {
84            row_len,
85            perm_1_seed: PERM_1_SEED,
86            perm_2_seed: PERM_2_SEED,
87            perm_1,
88            perm_2,
89            phantom: PhantomData,
90        }
91    }
92
93    /// Do the actual encoding, as per RAA spec
94    fn encode_inner<In, C, Map>(&self, cfg: &C, row: &[In], map: Map) -> Vec<C::Element>
95    where
96        C: SemiringConfig,
97        Map: Fn(&In) -> C::Element + Clone,
98    {
99        debug_assert_eq!(
100            row.len(),
101            self.row_len,
102            "Row length must match the code's row length"
103        );
104
105        let mut result: Vec<C::Element> = repeat(row, REP, map);
106        if Config::PERMUTE_IN_PLACE {
107            shuffle_seeded(&mut result, self.perm_1_seed);
108        } else {
109            result = clone_shuffled(&result, &self.perm_1);
110        }
111        if Config::CHECK_FOR_OVERFLOWS {
112            accumulate(cfg, &mut result);
113        } else {
114            accumulate_unchecked(cfg, &mut result);
115        }
116        if Config::PERMUTE_IN_PLACE {
117            shuffle_seeded(&mut result, self.perm_2_seed);
118        } else {
119            result = clone_shuffled(&result, &self.perm_2);
120        }
121        if Config::CHECK_FOR_OVERFLOWS {
122            accumulate(cfg, &mut result);
123        } else {
124            accumulate_unchecked(cfg, &mut result);
125        }
126        debug_assert_eq!(result.len(), self.codeword_len());
127        result
128    }
129}
130
131impl<Zt: ZipTypes, Config: RaaConfig, const REP: usize> LinearCode<Zt>
132    for RaaCode<Zt, Config, REP>
133{
134    const REPETITION_FACTOR: usize = REP;
135
136    fn row_len(&self) -> usize {
137        self.row_len
138    }
139
140    #[allow(clippy::arithmetic_side_effects)]
141    fn codeword_len(&self) -> usize {
142        self.row_len * REP
143    }
144
145    fn params_string(&self) -> String {
146        format!("row_len={}, rate=1/{REP}", self.row_len())
147    }
148
149    fn encode(&self, row: &[Zt::Eval]) -> Vec<Zt::Cw> {
150        self.encode_inner(&FixedConfig::default(), row, Zt::Cw::from_ref)
151    }
152
153    fn encode_wide(&self, row: &[Zt::CombR]) -> Vec<Zt::CombR> {
154        self.encode_inner(&FixedConfig::default(), row, |v| v.clone())
155    }
156
157    fn encode_f<C>(&self, cfg: &C, row: &[C::Element]) -> Vec<C::Element>
158    where
159        C: BaseFieldConfig + ProjectPrimitiveIntegersWithConfig,
160    {
161        self.encode_inner(cfg, row, |v| v.clone())
162    }
163}
164
165impl<Zt: ZipTypes, Config: RaaConfig, const REP: usize> Debug for RaaCode<Zt, Config, REP> {
166    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
167        f.debug_struct("RaaCode")
168            .field("row_len", &self.row_len)
169            .field("perm_1_seed", &self.perm_1_seed)
170            .field("perm_2_seed", &self.perm_2_seed)
171            .finish()
172    }
173}
174
175impl<Zt: ZipTypes, Config: RaaConfig, const REP: usize> PartialEq for RaaCode<Zt, Config, REP> {
176    fn eq(&self, other: &Self) -> bool {
177        self.row_len == other.row_len
178            && self.perm_1_seed == other.perm_1_seed
179            && self.perm_2_seed == other.perm_2_seed
180    }
181}
182
183impl<Zt: ZipTypes, Config: RaaConfig, const REP: usize> Eq for RaaCode<Zt, Config, REP> {}
184
185/// Repeat the given slice N times, e.g `[1,2,3] => [1,2,3,1,2,3]`
186#[allow(clippy::arithmetic_side_effects)]
187#[inline]
188pub(crate) fn repeat<In, Out: Clone>(
189    input: &[In],
190    repetition_factor: usize,
191    map: impl Fn(&In) -> Out + Clone,
192) -> Vec<Out> {
193    input
194        .iter()
195        .map(map)
196        .cycle()
197        .take(input.len() * repetition_factor)
198        .collect()
199}
200
201/// Perform an operation equivalent to multiplying the slice in-place by the
202/// accumulation matrix from the RAA code - a lower triangular matrix of the
203/// appropriate size, i.e. a matrix looking like this:
204///
205/// ```text
206/// 1 0 0 0
207/// 1 1 0 0
208/// 1 1 1 0
209/// 1 1 1 1
210/// ```
211pub(crate) fn accumulate<C: SemiringConfig>(cfg: &C, input: &mut [C::Element]) {
212    if let Some(first) = input.first().cloned() {
213        let mut acc = first;
214        for curr in input.iter_mut().skip(1) {
215            acc = cfg.checked_add(curr, &acc).expect("Accumulation overflow");
216            *curr = acc.clone();
217        }
218    }
219}
220
221pub(crate) fn accumulate_unchecked<C: SemiringConfig>(cfg: &C, input: &mut [C::Element]) {
222    if let Some(first) = input.first().cloned() {
223        let mut acc = first;
224        for i in 1..input.len() {
225            // Avoid bound checking
226            unsafe {
227                cfg.add_assign(&mut acc, input.get_unchecked(i));
228                *input.get_unchecked_mut(i) = acc.clone();
229            };
230        }
231    }
232}
233
234/// Clone the data using a precomputed permutation.
235pub(crate) fn clone_shuffled<T>(data: &[T], perm: &[usize]) -> Vec<T>
236where
237    T: Clone,
238{
239    perm.iter().map(|&i| data[i].clone()).collect()
240}
241
242#[cfg(test)]
243#[allow(clippy::arithmetic_side_effects, clippy::clone_on_copy)]
244mod tests {
245    use super::*;
246    use crate::{code::LinearCode, pcs::test_utils::TestZipTypes, utils::shuffle_seeded};
247    use crypto_primitives::{crypto_bigint_int::Int, crypto_bigint_uint::U64};
248    use num_traits::Zero;
249
250    const REPETITION_FACTOR: usize = 4;
251
252    // Define common types for testing
253    const INT_LIMBS: usize = U64::LIMBS;
254
255    const N: usize = INT_LIMBS;
256    const K: usize = INT_LIMBS * 4;
257    const M: usize = INT_LIMBS * 8;
258
259    #[derive(Clone, Copy)]
260    struct RaaConfigGeneric<const PERMUTE_IN_PLACE: bool, const CHECK_FOR_OVERFLOWS: bool>;
261
262    impl<const PERMUTE_IN_PLACE: bool, const CHECK_FOR_OVERFLOWS: bool> RaaConfig
263        for RaaConfigGeneric<PERMUTE_IN_PLACE, CHECK_FOR_OVERFLOWS>
264    {
265        const PERMUTE_IN_PLACE: bool = PERMUTE_IN_PLACE;
266        const CHECK_FOR_OVERFLOWS: bool = CHECK_FOR_OVERFLOWS;
267    }
268
269    macro_rules! cfg {
270        () => {
271            &FixedConfig::default()
272        };
273    }
274
275    macro_rules! test_raa {
276        ($zt:ty, $row_len:expr, $f:expr) => {
277            test_raa!($zt, $row_len, $f, RaaConfigGeneric<false, false>);
278            test_raa!($zt, $row_len, $f, RaaConfigGeneric<false, true>);
279            test_raa!($zt, $row_len, $f, RaaConfigGeneric<true, false>);
280            test_raa!($zt, $row_len, $f, RaaConfigGeneric<true, true>);
281        };
282
283        ($zt:ty, $row_len:expr, $f:expr, $config:ty) => {
284            {
285                let code = RaaCode::<$zt, $config, REPETITION_FACTOR>::new($row_len);
286                $f(&code)
287            }
288        };
289    }
290
291    #[test]
292    fn repeat_function_duplicates_row_correctly() {
293        type I = Int<N>;
294        let input = [10, 20].map(I::from);
295
296        let repetition_factor = 3;
297
298        let repeated_output = repeat::<_, I>(&input, repetition_factor, |v| v.clone());
299
300        let expected_output: Vec<_> = [10, 20, 10, 20, 10, 20].into_iter().map(I::from).collect();
301        assert_eq!(
302            repeated_output, expected_output,
303            "Failed on repetition factor > 1"
304        );
305
306        let empty_input: Vec<I> = vec![];
307        let repeated_empty = repeat::<_, I>(&empty_input, 5, |v| v.clone());
308        assert!(repeated_empty.is_empty(), "Failed on empty input vector");
309
310        let repeated_once = repeat::<_, I>(&input, 1, |v| v.clone());
311        assert_eq!(repeated_once, input, "Failed on repetition factor of 1");
312    }
313
314    #[test]
315    fn accumulate_function_computes_cumulative_sum() {
316        type I = Int<N>;
317        let mut input1: Vec<I> = [1, 2, 3, 4].into_iter().map(I::from).collect();
318        let expected1: Vec<I> = [1, 3, 6, 10].into_iter().map(I::from).collect();
319        accumulate(cfg!(), &mut input1);
320        assert_eq!(input1, expected1, "Failed on positive integers");
321
322        let mut input1: Vec<I> = [1, 2, 3, 4].into_iter().map(I::from).collect();
323        accumulate_unchecked(cfg!(), &mut input1);
324        assert_eq!(input1, expected1, "Failed on positive integers");
325
326        let mut input2: Vec<I> = [5, 0, 2, 0].into_iter().map(I::from).collect();
327        let expected2: Vec<I> = [5, 5, 7, 7].into_iter().map(I::from).collect();
328        accumulate(cfg!(), &mut input2);
329        assert_eq!(input2, expected2, "Failed on vector with zeros");
330
331        let mut input3: Vec<I> = [-1, 5, -10, 2].into_iter().map(I::from).collect();
332        let expected3: Vec<I> = [-1, 4, -6, -4].into_iter().map(I::from).collect();
333        accumulate(cfg!(), &mut input3);
334        assert_eq!(input3, expected3, "Failed on vector with negative numbers");
335
336        let mut empty_input: Vec<I> = vec![];
337        let expected_empty: Vec<I> = vec![];
338        accumulate(cfg!(), &mut empty_input);
339        assert_eq!(empty_input, expected_empty, "Failed on empty vector");
340    }
341
342    #[test]
343    fn shuffle_is_deterministic_for_a_given_seed() {
344        type I = Int<N>;
345        let original: Vec<I> = (1..=10).map(I::from).collect();
346        let mut vec1 = original.clone();
347        let mut vec2 = original.clone();
348        let mut vec3 = original.clone();
349
350        let seed1 = 12345;
351        let seed2 = 54321;
352
353        shuffle_seeded(&mut vec1, seed1);
354        shuffle_seeded(&mut vec2, seed1);
355        shuffle_seeded(&mut vec3, seed2);
356
357        assert_eq!(
358            vec1, vec2,
359            "Shuffling with the same seed should produce the same result"
360        );
361        assert_ne!(
362            vec1, vec3,
363            "Shuffling with different seeds should produce different results"
364        );
365        assert_ne!(
366            vec1, original,
367            "Shuffled vector should not be the same as the original"
368        );
369        assert_ne!(
370            vec3, original,
371            "Shuffled vector should not be the same as the original"
372        );
373    }
374
375    #[test]
376    fn encoding_preserves_linearity() {
377        test_raa!(TestZipTypes<N, K, M>, 4, |code: &RaaCode<_, _, _>| {
378            let a: Vec<Int<N>> = (1..=4).map(Int::<N>::from).collect();
379            let b: Vec<Int<N>> = (5..=8).map(Int::<N>::from).collect();
380            let sum_ab: Vec<Int<N>> = a.iter().zip(b.iter()).map(|(x, y)| *x + y).collect();
381
382            let encode_a: Vec<Int<K>> = code.encode(&a);
383            let encode_b: Vec<Int<K>> = code.encode(&b);
384            let encode_sum_ab: Vec<Int<K>> = code.encode(&sum_ab);
385
386            let sum_encode_ab: Vec<Int<K>> = encode_a
387                .iter()
388                .zip(encode_b.iter())
389                .map(|(x, y)| *x + y)
390                .collect();
391
392            assert_eq!(encode_sum_ab, sum_encode_ab);
393        });
394    }
395
396    /// Since our shuffle seeds are fixed, we can test the encoding
397    /// against a known output.
398    #[test]
399    fn encoding_produces_predictable_results() {
400        let a: Vec<Int<N>> = (1..=4).map(Int::<N>::from).collect();
401
402        test_raa!(TestZipTypes<N, K, M>, 4, |code: &RaaCode<_, _, _>| {
403            let encode_a: Vec<Int<K>> = code.encode(&a);
404            assert_eq!(
405                encode_a,
406                [
407                    0x1E, 0x36, 0x39, 0x5A, 0x70, 0x7E, 0xA5, 0xC1, 0xCB, 0xDC, 0xF9, 0x11E, 0x124,
408                    0x14C, 0x14D, 0x160
409                ]
410                .map(Int::<K>::from)
411            );
412        });
413    }
414
415    #[test]
416    fn encoding_zero_vector_results_in_zero_codeword() {
417        test_raa!(TestZipTypes<N, K, M>, 4, |code: &RaaCode<_, _, _>| {
418            let zero_vector: Vec<_> = vec![Int::<N>::zero(); code.row_len()];
419            let encoded_vector: Vec<Int<K>> = code.encode(&zero_vector);
420
421            let expected_codeword: Vec<Int<K>> = vec![Int::zero(); code.codeword_len()];
422
423            assert_eq!(
424                encoded_vector, expected_codeword,
425                "Encoding a zero vector should result in a zero codeword"
426            );
427        });
428    }
429
430    #[test]
431    fn in_place_permutation_should_not_affect_order() {
432        let data: Vec<Int<N>> = (1..=1024).map(Int::<N>::from).collect();
433        let row_len = data.len();
434        let codeword_1: Vec<Int<K>> = {
435            let code_in_place = RaaCode::<
436                TestZipTypes<N, K, M>,
437                RaaConfigGeneric<true, true>,
438                REPETITION_FACTOR,
439            >::new(row_len);
440            code_in_place.encode(&data)
441        };
442
443        let codeword_2: Vec<Int<K>> = {
444            let code_cloning = RaaCode::<
445                TestZipTypes<N, K, M>,
446                RaaConfigGeneric<false, true>,
447                REPETITION_FACTOR,
448            >::new(row_len);
449            code_cloning.encode(&data)
450        };
451        assert_eq!(
452            codeword_1, codeword_2,
453            "In-place permutation should not affect the final codeword"
454        );
455    }
456
457    #[test]
458    #[should_panic]
459    fn constructor_panics_on_insufficient_codeword_width() {
460        const N: usize = 1;
461        const K: usize = 1;
462
463        let _code = RaaCode::<
464            TestZipTypes<N, K, N>,
465            RaaConfigGeneric<false, true>,
466            REPETITION_FACTOR,
467        >::new(1 << 15);
468    }
469
470    #[test]
471    #[should_panic(expected = "Row length must match the code's row length")]
472    #[cfg(debug_assertions)]
473    fn encode_panics_on_mismatched_row_length() {
474        test_raa!(TestZipTypes<N, K, M>, 4, |code: &RaaCode<_, _, _>| {
475            let incorrect_row = vec![Int::<N>::from(1), Int::<N>::from(2), Int::<N>::from(3)];
476            let _: Vec<Int<K>> = code.encode(&incorrect_row);
477        });
478    }
479}
zip_plus/code/raa.rs

zip_plus/code/
raa.rs
