kopis/arithmetic/
matrix_arith.rs

1//! This file defines and implements matrices over our ring
2
3use crate::{arithmetic::RingElem, consts::RING_DEG};
4
5use crate::arithmetic::ring_arith::ring_mul_acc;
6
7use zeroize::Zeroize;
8
9/// An element of R^{x×y} where R is a [`RingElem`], stored in row-major order
10// We store the matrix in row-major order, so the outer array is the number of rows, i.e., the
11// height, i.e., X
12#[derive(Eq, PartialEq, Debug, Clone, Zeroize)]
13pub(crate) struct Matrix<const X: usize, const Y: usize>(pub(crate) [[RingElem; Y]; X]);
14
15impl<const X: usize, const Y: usize> Default for Matrix<X, Y> {
16    fn default() -> Self {
17        Matrix([[RingElem::default(); Y]; X])
18    }
19}
20
21impl<const X: usize, const Y: usize> Matrix<X, Y> {
22    /// Applies [`RingElem::shift_right`] to each element in the matrix
23    pub(crate) fn shift_right(&mut self, shift: usize) {
24        for row in self.0.iter_mut() {
25            for elem in row.iter_mut() {
26                elem.shift_right(shift)
27            }
28        }
29    }
30
31    /// Multiplies two matrices, using multiply-accumulate to avoid intermediate temporaries.
32    /// Each ring product is accumulated directly into the result element.
33    //
34    // This is the *reference* multiplication. Production code multiplies via `NttMatrix` (see
35    // ntt.rs), which is proved to agree with this function; the Lean correspondence between
36    // this schoolbook form and the spec (`Kopis/Properties/MatrixMul.lean`,
37    // `MulTranspose.lean`) is what that agreement is composed with. So it must stay outside
38    // `cfg(test)` to remain visible to the extractor, even though nothing outside tests calls
39    // it — it is `pub(crate)` and unused, so codegen drops it from the binary.
40    #[allow(dead_code)]
41    pub(crate) fn mul<const Z: usize>(&self, other: &Matrix<Y, Z>) -> Matrix<X, Z> {
42        let mut result = Matrix::default();
43        for i in 0..X {
44            for j in 0..Y {
45                for k in 0..Z {
46                    ring_mul_acc(&mut result.0[i][k], &self.0[i][j], &other.0[j][k]);
47                }
48            }
49        }
50
51        result
52    }
53
54    /// Multiplies the transpose of this matrix by the given vector, using multiply-accumulate
55    /// to avoid intermediate temporaries.
56    // The reference multiplication; see the note on `mul` above for why it is not `cfg(test)`.
57    #[allow(dead_code)]
58    pub(crate) fn mul_transpose<const Z: usize>(&self, other: &Matrix<X, Z>) -> Matrix<Y, Z> {
59        let mut result = Matrix::default();
60        for i in 0..X {
61            for j in 0..Y {
62                for k in 0..Z {
63                    ring_mul_acc(&mut result.0[j][k], &self.0[i][j], &other.0[i][k]);
64                }
65            }
66        }
67
68        result
69    }
70
71    /// Adds a given value to all coefficients of all elements of the matrix
72    pub(crate) fn wrapping_add_to_all(&mut self, val: u16) {
73        for row in self.0.iter_mut() {
74            for elem in row.iter_mut() {
75                elem.wrapping_add_to_all(val);
76            }
77        }
78    }
79
80    /// Serializes this matrix, ring element by ring element, treating each ring element
81    /// coefficient as having only `bits_per_elem` bits. In Saber terms, this runs POLYVECk2BS
82    /// where k = bits_per_elem
83    pub(crate) fn serialize(&self, out_buf: &mut [u8], bits_per_elem: usize) {
84        assert_eq!(out_buf.len(), X * Y * bits_per_elem * RING_DEG / 8);
85
86        /* More idiomatic version here. We have to use explicit indices because aeneas (the Lean
87         * extractor) has a problem with some iterator patterns
88        let mut chunk_iter = out_buf.chunks_mut(bits_per_elem * RING_DEG / 8);
89        for row in self.0.iter() {
90            for entry in row.iter() {
91                let out_chunk = chunk_iter
92                    .next()
93                    .expect("length check at beginning ensures X*Y many chunks");
94                entry.serialize(out_chunk, bits_per_elem);
95        */
96        let chunk_len = bits_per_elem * RING_DEG / 8;
97        for i in 0..X {
98            for j in 0..Y {
99                let idx = i * Y + j;
100                let out_chunk = &mut out_buf[idx * chunk_len..(idx + 1) * chunk_len];
101                self.0[i][j].serialize(out_chunk, bits_per_elem);
102            }
103        }
104    }
105
106    /// Deserializes a matrix of R10 values, element by element
107    pub(crate) fn deserialize_10(bytes: &[u8]) -> Self {
108        debug_assert_eq!(bytes.len(), X * Y * 10 * RING_DEG / 8);
109        let mut result = Matrix::default();
110
111        /* More idiomatic version here. We have to use explicit indices because aeneas (the Lean
112         * extractor) has a problem with some iterator patterns
113        let mut chunk_iter = bytes.chunks(bits_per_elem * RING_DEG / 8);
114        for row in result.0.iter_mut() {
115            for entry in row.iter_mut() {
116                let chunk = chunk_iter
117                    .next()
118                    .expect("length check at beginning ensures X*Y many chunks");
119                *entry = RingElem::deserialize(chunk, bits_per_elem);
120             }
121        }
122        */
123        let chunk_len = 10 * RING_DEG / 8;
124        for i in 0..X {
125            for j in 0..Y {
126                let idx = i * Y + j;
127                let chunk = &bytes[idx * chunk_len..(idx + 1) * chunk_len];
128                result.0[i][j] = RingElem::deserialize(chunk, 10);
129            }
130        }
131
132        result
133    }
134}
135
136impl<'a, const X: usize, const Y: usize> core::ops::Add<&'a Matrix<X, Y>> for &'a Matrix<X, Y> {
137    type Output = Matrix<X, Y>;
138
139    fn add(self, other: &'a Matrix<X, Y>) -> Self::Output {
140        let mut result = Matrix::default();
141        for i in 0..X {
142            for j in 0..Y {
143                result.0[i][j] = &self.0[i][j] + &other.0[i][j];
144            }
145        }
146
147        result
148    }
149}
150
151#[cfg(test)]
152mod test {
153    use super::*;
154
155    impl<const X: usize, const Y: usize> Matrix<X, Y> {
156        #[cfg(test)]
157        pub fn rand(rng: &mut impl rand_core::CryptoRng) -> Self {
158            let mut mat = Matrix::default();
159            for i in 0..X {
160                for j in 0..Y {
161                    mat.0[i][j] = RingElem::rand(rng);
162                }
163            }
164            mat
165        }
166
167        /// Returns the matrix transpose
168        pub(crate) fn transpose(&self) -> Matrix<Y, X> {
169            let mut ret = Matrix::default();
170            for i in 0..X {
171                for j in 0..Y {
172                    ret.0[j][i] = self.0[i][j];
173                }
174            }
175            ret
176        }
177    }
178
179    // Checks that mul and mul_transpose distribute over addition on the RHS
180    #[test]
181    fn distributivity() {
182        const X: usize = 4;
183        const Y: usize = 7;
184
185        let mut rng = rand::rng();
186
187        // Test mul_transpose
188        let mat = Matrix::<X, Y>::rand(&mut rng);
189        let vec1 = Matrix::<X, 1>::rand(&mut rng);
190        let vec2 = Matrix::<X, 1>::rand(&mut rng);
191        let prod1 = {
192            let vec_sum = &vec1 + &vec2;
193            mat.mul_transpose(&vec_sum)
194        };
195        let prod2 = &mat.mul_transpose(&vec1) + &mat.mul_transpose(&vec2);
196        assert_eq!(prod1, prod2);
197
198        // Now do the same with mul
199        let vec1 = Matrix::<Y, 1>::rand(&mut rng);
200        let vec2 = Matrix::<Y, 1>::rand(&mut rng);
201        let prod1 = {
202            let vec_sum = &vec1 + &vec2;
203            mat.mul(&vec_sum)
204        };
205        let prod2 = &mat.mul(&vec1) + &mat.mul(&vec2);
206        assert_eq!(prod1, prod2);
207    }
208
209    // Checks that mul, mul_transpose, and transpose are consistent with each other
210    #[test]
211    fn transpose() {
212        // Some arbitrary dimensions
213        const X: usize = 4;
214        const Y: usize = 7;
215        const Z: usize = 13;
216
217        let mut rng = rand::rng();
218
219        // Check that A^T B == (B^T A)^T
220        let mat1 = Matrix::<X, Y>::rand(&mut rng);
221        let mat2 = Matrix::<X, Z>::rand(&mut rng);
222        let prod1 = mat1.mul_transpose(&mat2);
223        let prod2 = mat2.mul_transpose(&mat1).transpose();
224        assert_eq!(prod1, prod2);
225
226        // Check that (AB)^T == A^T B^T
227        let mat1 = Matrix::<X, Y>::rand(&mut rng);
228        let mat2 = Matrix::<Y, Z>::rand(&mut rng);
229        let prod1 = &mat1.mul(&mat2).transpose();
230        let prod2 = &mat2.transpose().mul(&mat1.transpose());
231        assert_eq!(prod1, prod2);
232    }
233}
kopis/arithmetic/matrix_arith.rs

kopis/arithmetic/
matrix_arith.rs
