Skip to main content

kopis/arithmetic/
ring_arith.rs

1//! This file defines and implements Saber ring elements, specifically ℤ[X]/(X^256 + 1) mod n where
2//! n can be any power of two at most 2^16
3
4use crate::{
5    consts::RING_DEG,
6    ser::{deserialize_generic, serialize},
7};
8
9use core::ops::{Add, Mul, Sub};
10
11use zeroize::Zeroize;
12
13/// An element of the ring (Z/2^13 Z)[X] / (X^256 + 1)
14// The coefficients are in order of ascending powers, i.e., `self.0[0]` is the constant term
15#[derive(Eq, PartialEq, Debug, Clone, Copy, Zeroize)]
16pub struct RingElem(pub(crate) [u16; RING_DEG]);
17
18impl Default for RingElem {
19    fn default() -> Self {
20        RingElem([0u16; RING_DEG])
21    }
22}
23
24impl RingElem {
25    /// Creates a random ring element
26    #[cfg(test)]
27    pub(crate) fn rand(rng: &mut impl rand_core::CryptoRng) -> Self {
28        let modulus = 1 << crate::consts::MODULUS_Q_BITS as u32;
29
30        let mut result = [0; RING_DEG];
31        result.iter_mut().for_each(|coeff| {
32            let w = rng.next_u32() % modulus;
33            *coeff = w as u16;
34        });
35
36        RingElem(result)
37    }
38
39    /// Deserializes a ring element, treating each coefficient as having only `bits_per_elem` bits.
40    #[allow(clippy::unwrap_used)]
41    pub(crate) fn deserialize(bytes: &[u8], bits_per_elem: usize) -> Self {
42        assert_eq!(bytes.len(), bits_per_elem * RING_DEG / 8);
43
44        // Specialize based on bits_per_elem. unwraps are okay because of the check aboev
45        if bits_per_elem == crate::consts::MODULUS_Q_BITS {
46            let arr: &[u8; 13 * RING_DEG / 8] = bytes.try_into().unwrap();
47            RingElem(crate::ser::deserialize_13(arr))
48        } else if bits_per_elem == crate::consts::MODULUS_P_BITS {
49            let arr: &[u8; 10 * RING_DEG / 8] = bytes.try_into().unwrap();
50            RingElem(crate::ser::deserialize_10(arr))
51        } else {
52            RingElem(deserialize_generic(bytes, bits_per_elem))
53        }
54    }
55
56    /// Serializes this ring element, treating each coefficient as having only `bits_per_elem`
57    /// bits. In Saber terms, this runs POLYk2BS where k = bits_per_elem
58    #[allow(clippy::unwrap_used)]
59    pub(crate) fn serialize(&self, out_buf: &mut [u8], bits_per_elem: usize) {
60        assert_eq!(out_buf.len(), bits_per_elem * RING_DEG / 8);
61
62        // Specialize based on bits_per_elem. unwrap is okay because of the check above
63        if bits_per_elem == crate::consts::MODULUS_P_BITS {
64            let arr: &mut [u8; 10 * RING_DEG / 8] = out_buf.try_into().unwrap();
65            crate::ser::serialize_10(&self.0, arr)
66        } else {
67            serialize(&self.0, out_buf, bits_per_elem)
68        }
69    }
70
71    // Algorithm 8, ShiftRight
72    /// Right-shifts each coefficient by the specified amount, essentially dividing each coeff by a
73    /// power of two with rounding
74    pub(crate) fn shift_right(&mut self, shift: usize) {
75        for coeff in self.0.iter_mut() {
76            *coeff >>= shift;
77        }
78    }
79
80    // Algorithm 7, ShiftLeft
81    /// Left-shifts each coefficient by the specified amount, essentially multiplying each coeff by
82    /// a power of two, mod 2^16
83    pub(crate) fn shift_left(&mut self, shift: usize) {
84        for coeff in self.0.iter_mut() {
85            *coeff <<= shift;
86        }
87    }
88
89    /// Adds a given value to all coefficients
90    pub(crate) fn wrapping_add_to_all(&mut self, val: u16) {
91        for coeff in self.0.iter_mut() {
92            *coeff = coeff.wrapping_add(val);
93        }
94    }
95}
96
97impl<'a> Mul for &'a RingElem {
98    type Output = RingElem;
99
100    fn mul(self, other: &'a RingElem) -> Self::Output {
101        let mut ret = RingElem::default();
102        ring_mul_acc(&mut ret, self, other);
103        ret
104    }
105}
106
107// Half the ring degree. We split 256-coefficient polys into two 128-coefficient halves
108// for a single level of Karatsuba. Two levels (64×64 base) was benchmarked ~15% slower
109// due to increased overhead and less efficient vectorization of shorter inner loops.
110const HALF: usize = RING_DEG / 2; // 128
111
112/// Schoolbook multiplication of two 128-coefficient polynomials.
113/// The product of two degree-127 polys has degree at most 254, so all 256 output slots suffice.
114/// Writes result into `out[0..254]`; `out` must be zeroed on entry.
115///
116/// Takes fixed-size array references so the compiler knows the exact bounds and can
117/// eliminate all bounds checks and vectorize the inner loop.
118#[inline(never)] // Benchmarked: keeping this separate lets the compiler vectorize the inner loop better
119fn schoolbook_128(out: &mut [u16; RING_DEG], a: &[u16; HALF], b: &[u16; HALF]) {
120    // Standard O(n²) schoolbook. The inner loop over b is contiguous in memory, which is
121    // cache-friendly. The compiler can hoist a[i] as a loop-invariant broadcast.
122    for i in 0..HALF {
123        let ai = a[i];
124        for j in 0..HALF {
125            out[i + j] = out[i + j].wrapping_add(ai.wrapping_mul(b[j]));
126        }
127    }
128}
129
130/// Multiplies two ring elements using one level of Karatsuba, and **accumulates** the product
131/// into `acc`. This is the core hot function for Saber's matrix-vector multiplies.
132#[allow(clippy::unwrap_used)]
133// We use unwrap to split the slices. Once it's stable we should use split_array_ref()
134//   https://doc.rust-lang.org/std/primitive.array.html#method.split_array_ref
135pub(crate) fn ring_mul_acc(acc: &mut RingElem, a: &RingElem, b: &RingElem) {
136    // Convert slices to fixed-size array references for the schoolbook function.
137    // These are infallible since we split a RING_DEG array exactly in half.
138    let a_lo: &[u16; HALF] = a.0[..HALF].try_into().unwrap();
139    let a_hi: &[u16; HALF] = a.0[HALF..].try_into().unwrap();
140    let b_lo: &[u16; HALF] = b.0[..HALF].try_into().unwrap();
141    let b_hi: &[u16; HALF] = b.0[HALF..].try_into().unwrap();
142
143    // We split each input into low and high 128-coefficient halves:
144    //     a = a_lo + a_hi * X^128,   b = b_lo + b_hi * X^128
145    // Then use Karatsuba's identity:
146    //     a*b = z0 + z1*X^128 + z2*X^256
147    // where z0 = a_lo*b_lo, z2 = a_hi*b_hi, z3 = (a_lo+a_hi)*(b_lo+b_hi), z1 = z3 - z0 - z2.
148    //
149    // Since we work in Z[X]/(X^256 + 1), X^256 = -1, so:
150    //     a*b mod (X^256+1) = (z0 - z2) + z1*X^128  mod (X^256+1)
151    // And z1*X^128 wraps: coefficients 0..127 of z1 go to positions 128..255,
152    // while coefficients 128..255 of z1 wrap to positions 0..127 with a sign flip.
153
154    // Compute the three schoolbook products into flat arrays.
155    // Each is a product of two degree-127 polynomials, fitting in 256 coefficients.
156    let mut z0 = [0u16; RING_DEG];
157    let mut z2 = [0u16; RING_DEG];
158    schoolbook_128(&mut z0, a_lo, b_lo);
159    schoolbook_128(&mut z2, a_hi, b_hi);
160
161    // Compute (a_lo + a_hi) and (b_lo + b_hi) for the cross term
162    let mut a_sum = [0u16; HALF];
163    let mut b_sum = [0u16; HALF];
164    for i in 0..HALF {
165        a_sum[i] = a_lo[i].wrapping_add(a_hi[i]);
166        b_sum[i] = b_lo[i].wrapping_add(b_hi[i]);
167    }
168    let mut z3 = [0u16; RING_DEG];
169    schoolbook_128(&mut z3, &a_sum, &b_sum);
170
171    // Accumulate: acc += z0 - z2 + z1*X^128  mod (X^256+1)
172    //   where z1 = (z3 - z0 - z2)
173    for j in 0..HALF {
174        // For acc[j] where j in 0..HALF:
175        //   * z0[j] - z2[j] from the direct terms
176        //   * -(z3[j+HALF] - z0[j+HALF] - z2[j+HALF]) from z1[j+HALF] wrapping with negation
177        let z1_wrap = z3[j + HALF]
178            .wrapping_sub(z0[j + HALF])
179            .wrapping_sub(z2[j + HALF]);
180        acc.0[j] = acc.0[j]
181            .wrapping_add(z0[j])
182            .wrapping_sub(z2[j])
183            .wrapping_sub(z1_wrap);
184
185        // For acc[j] where j in HALF..RING_DEG:
186        //   * z0[j] - z2[j] from the direct terms
187        //   * +(z3[j-HALF] - z0[j-HALF] - z2[j-HALF]) from z1[j-HALF] (no wrap)
188        let z1_direct = z3[j].wrapping_sub(z0[j]).wrapping_sub(z2[j]);
189        acc.0[j + HALF] = acc.0[j + HALF]
190            .wrapping_add(z0[j + HALF])
191            .wrapping_sub(z2[j + HALF])
192            .wrapping_add(z1_direct);
193    }
194}
195
196impl<'a> Add for &'a RingElem {
197    type Output = RingElem;
198
199    fn add(self, other: &'a RingElem) -> Self::Output {
200        let mut ret = RingElem::default();
201        for i in 0..RING_DEG {
202            ret.0[i] = self.0[i].wrapping_add(other.0[i]);
203        }
204        ret
205    }
206}
207
208impl<'a> Sub for &'a RingElem {
209    type Output = RingElem;
210
211    fn sub(self, other: &'a RingElem) -> Self::Output {
212        let mut ret = RingElem::default();
213        for i in 0..RING_DEG {
214            ret.0[i] = self.0[i].wrapping_sub(other.0[i]);
215        }
216        ret
217    }
218}
219
220#[cfg(test)]
221mod test {
222    use super::*;
223    use crate::consts::RING_DEG;
224
225    use rand::{Rng, RngCore, rng};
226
227    // Checks that a * b == b * a and a + b == b + a for ring elements a, b
228    #[test]
229    fn commutativity() {
230        let mut rng = rng();
231
232        for _ in 0..100 {
233            let a = RingElem::rand(&mut rng);
234            let b = RingElem::rand(&mut rng);
235
236            let prod_1 = &a * &b;
237            let prod_2 = &b * &a;
238
239            let sum_1 = &a + &b;
240            let sum_2 = &b + &a;
241
242            assert_eq!(prod_1, prod_2);
243            assert_eq!(sum_1, sum_2);
244        }
245    }
246
247    /// Naive schoolbook multiplication directly in Z[X]/(X^256+1) for testing.
248    /// This is the simplest correct implementation: O(n^2) with explicit ring reduction.
249    fn reference_schoolbook_ring_mul(a: &RingElem, b: &RingElem) -> RingElem {
250        let mut result = RingElem::default();
251        for i in 0..RING_DEG {
252            for j in 0..RING_DEG {
253                let prod = a.0[i].wrapping_mul(b.0[j]);
254                let idx = i + j;
255                if idx < RING_DEG {
256                    result.0[idx] = result.0[idx].wrapping_add(prod);
257                } else {
258                    // X^256 = -1 in our ring, so wrap and negate
259                    result.0[idx - RING_DEG] = result.0[idx - RING_DEG].wrapping_sub(prod);
260                }
261            }
262        }
263        result
264    }
265
266    // Tests that our Karatsuba-based ring_mul_acc matches the naive schoolbook ring multiply
267    #[test]
268    fn karatsuba_vs_schoolbook() {
269        let mut rng = rng();
270
271        for _ in 0..100 {
272            let a = RingElem::rand(&mut rng);
273            let b = RingElem::rand(&mut rng);
274
275            let reference = reference_schoolbook_ring_mul(&a, &b);
276            let optimized = &a * &b;
277
278            assert_eq!(reference, optimized);
279        }
280    }
281
282    // Tests that ring_mul_acc correctly accumulates into a non-zero buffer
283    #[test]
284    fn mul_acc_accumulates() {
285        let mut rng = rng();
286
287        for _ in 0..50 {
288            let a = RingElem::rand(&mut rng);
289            let b = RingElem::rand(&mut rng);
290            let c = RingElem::rand(&mut rng);
291            let d = RingElem::rand(&mut rng);
292
293            // Compute a*b + c*d via ring_mul_acc
294            let mut acc = RingElem::default();
295            ring_mul_acc(&mut acc, &a, &b);
296            ring_mul_acc(&mut acc, &c, &d);
297
298            // Compute the same thing via separate multiplies + add
299            let expected = &(&a * &b) + &(&c * &d);
300
301            assert_eq!(acc, expected);
302        }
303    }
304
305    // Tests serialization and deserialization of ring elements
306    #[test]
307    fn deserialize() {
308        let mut rng = rng();
309
310        // The largest buffer we'll need for the following tests. We make 2 because we need to
311        // compare values in some places
312        let mut backing_buf1 = [0u8; 16 * RING_DEG / 8];
313        let mut backing_buf2 = [0u8; 16 * RING_DEG / 8];
314
315        for _ in 0..1000 {
316            // Check that deserialize matches the reference impl deserialize for N=2^13,2^10,2^1
317            let bits_per_elem = 13;
318            let bytes = &mut backing_buf1[..bits_per_elem * RING_DEG / 8];
319            rng.fill_bytes(bytes);
320            assert_eq!(
321                saber_ref_from_bytes_mod8192(&bytes),
322                RingElem::deserialize(&bytes, 13)
323            );
324
325            // Now check it matches the reference to_bytes impl
326            let elem = RingElem::rand(&mut rng);
327            let my_bytes = &mut backing_buf1[..bits_per_elem * RING_DEG / 8];
328            let ref_bytes = &mut backing_buf2[..bits_per_elem * RING_DEG / 8];
329            elem.serialize(my_bytes, bits_per_elem);
330            reference_impl_to_bytes_mod8192(&elem, ref_bytes);
331            assert_eq!(my_bytes, ref_bytes);
332
333            let bits_per_elem = 10;
334            let bytes = &mut backing_buf1[..bits_per_elem * RING_DEG / 8];
335            rng.fill_bytes(bytes);
336            assert_eq!(
337                saber_ref_from_bytes_mod1024(&bytes).0,
338                RingElem::deserialize(&bytes, 10).0,
339            );
340
341            let bits_per_elem = 1;
342            let bytes = &mut backing_buf1[..bits_per_elem * RING_DEG / 8];
343            rng.fill_bytes(bytes);
344            assert_eq!(
345                saber_ref_from_bytes_mod2(&bytes).0,
346                RingElem::deserialize(&bytes, 1).0,
347            );
348
349            // Now check it matches the reference to_bytes impl
350            let elem = RingElem::rand(&mut rng);
351            // The reference impl actually requires that the buffer is zeroed before use
352            backing_buf2.fill(0);
353            let my_bytes = &mut backing_buf1[..bits_per_elem * RING_DEG / 8];
354            let ref_bytes = &mut backing_buf2[..bits_per_elem * RING_DEG / 8];
355            elem.serialize(my_bytes, bits_per_elem);
356            saber_ref_to_bytes_mod2(&elem, ref_bytes);
357            assert_eq!(my_bytes, ref_bytes);
358
359            // Now check that to_bytes and from_bytes are inverses
360
361            // Pick a random bits_per_elem
362            for _ in 0..10 {
363                let bits_per_elem = rng.random_range(1..=13);
364                let bitmask = (1 << bits_per_elem) - 1;
365
366                // Generate a random element and make sure none of the values exceed 2^bits_per_elem
367                let mut p = RingElem::rand(&mut rng);
368                p.0.iter_mut().for_each(|e| *e &= bitmask);
369
370                // Check that a round trip preserves the polynomial
371                let p_bytes = &mut backing_buf1[..bits_per_elem * RING_DEG / 8];
372                p.serialize(p_bytes, bits_per_elem);
373                assert_eq!(p, RingElem::deserialize(&p_bytes, bits_per_elem));
374
375                // Now other way around
376                let p_bytes = &mut backing_buf1[..bits_per_elem * RING_DEG / 8];
377                rng.fill_bytes(p_bytes);
378                let p = RingElem::deserialize(&p_bytes, bits_per_elem);
379                let new_p_bytes = &mut backing_buf2[..bits_per_elem * RING_DEG / 8];
380                p.serialize(new_p_bytes, bits_per_elem);
381                assert_eq!(p_bytes, new_p_bytes);
382            }
383        }
384    }
385
386    /// A nearly verbatim copy of the C reference impl of BS2POL_N where N = 2^13
387    /// https://github.com/KULeuven-COSIC/SABER/blob/f7f39e4db2f3e22a21e1dd635e0601caae2b4510/Reference_Implementation_KEM/pack_unpack.c#L101
388    fn saber_ref_from_bytes_mod8192(b: &[u8]) -> RingElem {
389        let mut offset_byte;
390        let mut offset_data;
391        let mut poly = RingElem::default();
392        let data = &mut poly.0;
393
394        let b_arr: [u8; 13 * RING_DEG / 8] = b.try_into().unwrap();
395        let bytes = b_arr.map(|x| x as u16);
396
397        for j in 0..RING_DEG / 8 {
398            offset_byte = 13 * j;
399            offset_data = 8 * j;
400            data[offset_data] =
401                (bytes[offset_byte] & (0xff)) | ((bytes[offset_byte + 1] & 0x1f) << 8);
402            data[offset_data + 1] = (bytes[offset_byte + 1] >> 5 & (0x07))
403                | ((bytes[offset_byte + 2] & 0xff) << 3)
404                | ((bytes[offset_byte + 3] & 0x03) << 11);
405            data[offset_data + 2] =
406                (bytes[offset_byte + 3] >> 2 & (0x3f)) | ((bytes[offset_byte + 4] & 0x7f) << 6);
407            data[offset_data + 3] = (bytes[offset_byte + 4] >> 7 & (0x01))
408                | ((bytes[offset_byte + 5] & 0xff) << 1)
409                | ((bytes[offset_byte + 6] & 0x0f) << 9);
410            data[offset_data + 4] = (bytes[offset_byte + 6] >> 4 & (0x0f))
411                | ((bytes[offset_byte + 7] & 0xff) << 4)
412                | ((bytes[offset_byte + 8] & 0x01) << 12);
413            data[offset_data + 5] =
414                (bytes[offset_byte + 8] >> 1 & (0x7f)) | ((bytes[offset_byte + 9] & 0x3f) << 7);
415            data[offset_data + 6] = (bytes[offset_byte + 9] >> 6 & (0x03))
416                | ((bytes[offset_byte + 10] & 0xff) << 2)
417                | ((bytes[offset_byte + 11] & 0x07) << 10);
418            data[offset_data + 7] =
419                (bytes[offset_byte + 11] >> 3 & (0x1f)) | ((bytes[offset_byte + 12] & 0xff) << 5);
420        }
421
422        poly
423    }
424
425    /// A nearly verbatim copy of the C reference impl of BS2POL_N where N = 2^10
426    /// https://github.com/KULeuven-COSIC/SABER/blob/f7f39e4db2f3e22a21e1dd635e0601caae2b4510/Reference_Implementation_KEM/pack_unpack.c#L134
427    fn saber_ref_from_bytes_mod1024(b: &[u8]) -> RingElem {
428        let mut offset_byte;
429        let mut offset_data;
430        let mut poly = RingElem::default();
431        let data = &mut poly.0;
432
433        let b_arr: [u8; 10 * RING_DEG / 8] = b.try_into().unwrap();
434        let bytes = b_arr.map(|x| x as u16);
435
436        for j in 0..RING_DEG / 4 {
437            offset_byte = 5 * j;
438            offset_data = 4 * j;
439            data[offset_data] =
440                (bytes[offset_byte] & (0xff)) | ((bytes[offset_byte + 1] & 0x03) << 8);
441            data[offset_data + 1] =
442                ((bytes[offset_byte + 1] >> 2) & (0x3f)) | ((bytes[offset_byte + 2] & 0x0f) << 6);
443            data[offset_data + 2] =
444                ((bytes[offset_byte + 2] >> 4) & (0x0f)) | ((bytes[offset_byte + 3] & 0x3f) << 4);
445            data[offset_data + 3] =
446                ((bytes[offset_byte + 3] >> 6) & (0x03)) | ((bytes[offset_byte + 4] & 0xff) << 2);
447        }
448
449        poly
450    }
451
452    /// A nearly verbatim copy of the C reference impl of POL2BS_N where N = 2^13
453    /// https://github.com/KULeuven-COSIC/SABER/blob/f7f39e4db2f3e22a21e1dd635e0601caae2b4510/Reference_Implementation_KEM/pack_unpack.c#L78
454    fn reference_impl_to_bytes_mod8192(polyn: &RingElem, bytes: &mut [u8]) {
455        let mut offset_byte: usize;
456        let mut offset_data: usize;
457        let data = polyn.0;
458
459        for j in 0..RING_DEG / 8 {
460            offset_byte = 13 * j;
461            offset_data = 8 * j;
462            bytes[offset_byte] = (data[offset_data] & (0xff)) as u8;
463            bytes[offset_byte + 1] = ((data[offset_data] >> 8) & 0x1f) as u8
464                | ((data[offset_data + 1] & 0x07) << 5) as u8;
465            bytes[offset_byte + 2] = ((data[offset_data + 1] >> 3) & 0xff) as u8;
466            bytes[offset_byte + 3] = ((data[offset_data + 1] >> 11) & 0x03) as u8
467                | ((data[offset_data + 2] & 0x3f) << 2) as u8;
468            bytes[offset_byte + 4] = ((data[offset_data + 2] >> 6) & 0x7f) as u8
469                | ((data[offset_data + 3] & 0x01) << 7) as u8;
470            bytes[offset_byte + 5] = ((data[offset_data + 3] >> 1) & 0xff) as u8;
471            bytes[offset_byte + 6] = ((data[offset_data + 3] >> 9) & 0x0f) as u8
472                | ((data[offset_data + 4] & 0x0f) << 4) as u8;
473            bytes[offset_byte + 7] = ((data[offset_data + 4] >> 4) & 0xff) as u8;
474            bytes[offset_byte + 8] = ((data[offset_data + 4] >> 12) & 0x01) as u8
475                | ((data[offset_data + 5] & 0x7f) << 1) as u8;
476            bytes[offset_byte + 9] = ((data[offset_data + 5] >> 7) & 0x3f) as u8
477                | ((data[offset_data + 6] & 0x03) << 6) as u8;
478            bytes[offset_byte + 10] = ((data[offset_data + 6] >> 2) & 0xff) as u8;
479            bytes[offset_byte + 11] = ((data[offset_data + 6] >> 10) & 0x07) as u8
480                | ((data[offset_data + 7] & 0x1f) << 3) as u8;
481            bytes[offset_byte + 12] = ((data[offset_data + 7] >> 5) & 0xff) as u8;
482        }
483    }
484
485    /// A nearly verbatim copy of the C reference impl of BS2POL_N where N = 2
486    /// https://github.com/KULeuven-COSIC/SABER/blob/f7f39e4db2f3e22a21e1dd635e0601caae2b4510/Reference_Implementation_KEM/pack_unpack.c#L184
487    fn saber_ref_from_bytes_mod2(b: &[u8]) -> RingElem {
488        let mut poly = RingElem::default();
489        let data = &mut poly.0;
490
491        let b_arr: [u8; 1 * RING_DEG / 8] = b.try_into().unwrap();
492        let bytes = b_arr.map(|x| x as u16);
493
494        for j in 0..32 {
495            {
496                for i in 0..8 {
497                    data[j * 8 + i] = (bytes[j] >> i) & 0x01;
498                }
499            }
500        }
501
502        poly
503    }
504
505    /// A nearly verbatim copy of the C reference impl of POL2BS_N where N = 2
506    /// https://github.com/KULeuven-COSIC/SABER/blob/f7f39e4db2f3e22a21e1dd635e0601caae2b4510/Reference_Implementation_KEM/pack_unpack.c#L196
507    fn saber_ref_to_bytes_mod2(polyn: &RingElem, bytes: &mut [u8]) {
508        let data = polyn.0;
509
510        for j in 0..32 {
511            for i in 0..8 {
512                bytes[j] |= ((data[j * 8 + i] & 0x01) << i) as u8;
513            }
514        }
515    }
516}