crystals_rs/
polymat.rs

1use core::ops::{Index, IndexMut};
2
3use crate::{
4    poly::{kyber::KyberPoly, SizedPolynomial, UNIFORM_SEED_BYTES},
5    polyvec::*,
6};
7
8#[derive(Debug)]
9pub struct PolyMat<P, const N: usize, const K: usize, const L: usize>([PolyVec<P, N, L>; K])
10where
11    P: SizedPolynomial<N>;
12
13impl<P, const N: usize, const K: usize, const L: usize> Default for PolyMat<P, N, K, L>
14where
15    P: SizedPolynomial<N>,
16{
17    #[inline]
18    fn default() -> Self {
19        Self([PolyVec::default(); K])
20    }
21}
22
23impl<P, const N: usize, const K: usize, const L: usize> Index<usize> for PolyMat<P, N, K, L>
24where
25    P: SizedPolynomial<N>,
26{
27    type Output = PolyVec<P, N, L>;
28
29    #[inline(always)]
30    fn index(&self, i: usize) -> &Self::Output {
31        &self.0[i]
32    }
33}
34
35impl<P, const N: usize, const K: usize, const L: usize> IndexMut<usize> for PolyMat<P, N, K, L>
36where
37    P: SizedPolynomial<N>,
38{
39    #[inline(always)]
40    fn index_mut(&mut self, i: usize) -> &mut Self::Output {
41        &mut self.0[i]
42    }
43}
44
45impl<P, const N: usize, const K: usize, const L: usize> AsRef<[PolyVec<P, N, L>; K]>
46    for PolyMat<P, N, K, L>
47where
48    P: SizedPolynomial<N>,
49{
50    #[inline(always)]
51    fn as_ref(&self) -> &[PolyVec<P, N, L>; K] {
52        &self.0
53    }
54}
55
56impl<P, const N: usize, const K: usize, const L: usize> AsMut<[PolyVec<P, N, L>; K]>
57    for PolyMat<P, N, K, L>
58where
59    P: SizedPolynomial<N>,
60{
61    #[inline(always)]
62    fn as_mut(&mut self) -> &mut [PolyVec<P, N, L>; K] {
63        &mut self.0
64    }
65}
66
67impl<P, const N: usize, const K: usize, const L: usize> PolyMat<P, N, K, L>
68where
69    P: SizedPolynomial<N>,
70{
71    //     #[inline]
72    //     pub fn mult_vec(&self, polyvec: &PolyVec<T, N, K>) -> PolyVec<T, N, K> {
73    //         let mut r = PolyVec::default();
74
75    //         for i in 0..K {
76    //             self[i].basemul_acc(&polyvec, &mut r[i]);
77    //         }
78    //         r
79    //     }
80
81    /// Expand seed to A matrix (or A^T if TRANSPOSED is true)
82    /// For Kyber:
83    ///
84    /// For dilithium:
85    ///                polyvec_matrix_expand TRANSPOSED = false
86    #[inline(always)]
87    pub fn gen_matrix<const TRANSPOSED: bool>(seed: &[u8; UNIFORM_SEED_BYTES]) -> Self {
88        let mut a = Self::default();
89        a.gen_matrix_into::<TRANSPOSED>(seed);
90        a
91    }
92
93    #[inline(always)]
94    pub fn gen_a(seed: &[u8; UNIFORM_SEED_BYTES]) -> Self {
95        Self::gen_matrix::<false>(seed)
96    }
97
98    #[inline(always)]
99    pub fn gen_at(seed: &[u8; UNIFORM_SEED_BYTES]) -> Self {
100        Self::gen_matrix::<true>(seed)
101    }
102
103    #[inline]
104    pub fn gen_matrix_into<const TRANSPOSED: bool>(&mut self, seed: &[u8; UNIFORM_SEED_BYTES]) {
105        for (i, vec) in self.as_mut().iter_mut().enumerate() {
106            vec.uniform_xof::<TRANSPOSED>(seed, i as u8);
107        }
108    }
109}
110
111pub type KyberMatrix<const K: usize> = PolyMat<KyberPoly, { KyberPoly::N }, K, K>;
112
113#[cfg(test)]
114mod tests {
115    use rand::Rng;
116
117    use crate::poly::kyber::KYBER_N;
118
119    use super::*;
120
121    #[test]
122    fn gen_matrix() {
123        gen_matrix_x::<2, true>();
124        gen_matrix_x::<3, true>();
125        gen_matrix_x::<4, true>();
126        gen_matrix_x::<2, false>();
127        gen_matrix_x::<3, false>();
128        gen_matrix_x::<4, false>();
129    }
130
131    fn gen_matrix_x<const K: usize, const TRANSPOSED: bool>() {
132        let mut seed = [0u8; 32];
133        let mut rng = rand::thread_rng();
134
135        const NUM_TESTS: usize = if cfg!(miri) { 3 } else { 111 };
136
137        let mut a_ref = [[[0i16; KYBER_N]; K]; K];
138
139        for _ in 0..NUM_TESTS {
140            rng.fill(&mut seed);
141
142            let a = KyberMatrix::<K>::gen_matrix::<TRANSPOSED>(&seed);
143
144            crystals_cref::kyber::gen_matrix(&mut a_ref, &seed, TRANSPOSED);
145
146            for i in 0..K {
147                for j in 0..K {
148                    for k in 0..KYBER_N {
149                        assert_eq!(
150                            a[i][j][k / 2].0[k % 2],
151                            a_ref[i][j][k],
152                            "i={}, j={}, k={}",
153                            i,
154                            j,
155                            k
156                        );
157                    }
158                }
159            }
160        }
161    }
162}
crystals_rs/polymat.rs

crystals_rs/
polymat.rs
