poulpy/base2k/src/vmp.rs
base2k/src/vmp.rsRaw
1use crate::ffi::vmp;
2use crate::{Infos, Module, VecZnx, VecZnxDft};
3use std::cmp::min;
4
5/// Vector Matrix Product Prepared Matrix: a vector of [VecZnx],
6/// stored as a 3D matrix in the DFT domain in a single contiguous array.
7/// Each row of the [VmpPMat] can be seen as a [VecZnxDft].
8///
9/// The backend array of [VmpPMat] is allocate in C,
10/// and thus must be manually freed.
11///
12/// [VmpPMat] is used to permform a vector matrix product between a [VecZnx] and a [VmpPMat].
13/// See the trait [VmpPMatOps] for additional information.
14pub struct VmpPMat {
15 /// The pointer to the C memory.
16 pub data: *mut vmp::vmp_pmat_t,
17 /// The number of [VecZnxDft].
18 pub rows: usize,
19 /// The number of limbs in each [VecZnxDft].
20 pub cols: usize,
21 /// The ring degree of each [VecZnxDft].
22 pub n: usize,
23}
24
25impl VmpPMat {
26 /// Returns the pointer to the [vmp_pmat_t].
27 pub fn data(&self) -> *mut vmp::vmp_pmat_t {
28 self.data
29 }
30
31 /// Returns a copy of the backend array at index (i, j) of the [VmpPMat].
32 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
33 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
34 ///
35 /// # Arguments
36 ///
37 /// * `row`: row index (i).
38 /// * `col`: col index (j).
39 pub fn at<T: Default + Copy>(&self, row: usize, col: usize) -> Vec<T> {
40 let mut res: Vec<T> = vec![T::default(); self.n];
41
42 if self.n < 8 {
43 res.copy_from_slice(
44 &self.get_backend_array::<T>()[(row + col * self.rows()) * self.n()
45 ..(row + col * self.rows()) * (self.n() + 1)],
46 );
47 } else {
48 (0..self.n >> 3).for_each(|blk| {
49 res[blk * 8..(blk + 1) * 8].copy_from_slice(&self.get_array(row, col, blk)[..8]);
50 });
51 }
52
53 res
54 }
55
56 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
57 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
58 fn get_array<T>(&self, row: usize, col: usize, blk: usize) -> &[T] {
59 let nrows: usize = self.rows();
60 let ncols: usize = self.cols();
61 if col == (ncols - 1) && (ncols & 1 == 1) {
62 &self.get_backend_array::<T>()[blk * nrows * ncols * 8 + col * nrows * 8 + row * 8..]
63 } else {
64 &self.get_backend_array::<T>()[blk * nrows * ncols * 8
65 + (col / 2) * (2 * nrows) * 8
66 + row * 2 * 8
67 + (col % 2) * 8..]
68 }
69 }
70
71 /// Returns a non-mutable reference of `T` of the entire contiguous array of the [VmpPMat].
72 /// When using [`crate::FFT64`] as backend, `T` should be [f64].
73 /// When using [`crate::NTT120`] as backend, `T` should be [i64].
74 /// The length of the returned array is rows * cols * n.
75 pub fn get_backend_array<T>(&self) -> &[T] {
76 let ptr: *const T = self.data as *const T;
77 let len: usize = (self.rows() * self.cols() * self.n() * 8) / std::mem::size_of::<T>();
78 unsafe { &std::slice::from_raw_parts(ptr, len) }
79 }
80}
81
82/// This trait implements methods for vector matrix product,
83/// that is, multiplying a [VecZnx] with a [VmpPMat].
84pub trait VmpPMatOps {
85 /// Allocates a new [VmpPMat] with the given number of rows and columns.
86 ///
87 /// # Arguments
88 ///
89 /// * `rows`: number of rows (number of [VecZnxDft]).
90 /// * `cols`: number of cols (number of limbs of each [VecZnxDft]).
91 fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat;
92
93 /// Returns the number of bytes needed as scratch space for [VmpPMatOps::vmp_prepare_contiguous].
94 ///
95 /// # Arguments
96 ///
97 /// * `rows`: number of rows of the [VmpPMat] used in [VmpPMatOps::vmp_prepare_contiguous].
98 /// * `cols`: number of cols of the [VmpPMat] used in [VmpPMatOps::vmp_prepare_contiguous].
99 fn vmp_prepare_tmp_bytes(&self, rows: usize, cols: usize) -> usize;
100
101 /// Prepares a [VmpPMat] from a contiguous array of [i64].
102 /// The helper struct [Matrix3D] can be used to contruct and populate
103 /// the appropriate contiguous array.
104 ///
105 /// # Arguments
106 ///
107 /// * `b`: [VmpPMat] on which the values are encoded.
108 /// * `a`: the contiguous array of [i64] of the 3D matrix to encode on the [VmpPMat].
109 /// * `buf`: scratch space, the size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
110 ///
111 /// # Example
112 /// ```
113 /// use base2k::{Module, Matrix3D, VmpPMat, VmpPMatOps, FFT64, Free};
114 /// use std::cmp::min;
115 ///
116 /// let n: usize = 1024;
117 /// let module = Module::new::<FFT64>(n);
118 /// let rows = 5;
119 /// let cols = 6;
120 ///
121 /// let mut b_mat: Matrix3D<i64> = Matrix3D::new(rows, cols, n);
122 ///
123 /// // Populates the i-th row of b_math with X^1 * 2^(i * log_w) (here log_w is undefined)
124 /// (0..min(rows, cols)).for_each(|i| {
125 /// b_mat.at_mut(i, i)[1] = 1 as i64;
126 /// });
127 ///
128 /// let mut buf: Vec<u8> = vec![u8::default(); module.vmp_prepare_tmp_bytes(rows, cols)];
129 ///
130 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
131 /// module.vmp_prepare_contiguous(&mut vmp_pmat, &b_mat.data, &mut buf);
132 ///
133 /// vmp_pmat.free() // don't forget to free the memory once vmp_pmat is not needed anymore.
134 /// ```
135 fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], buf: &mut [u8]);
136
137 /// Prepares a [VmpPMat] from a vector of [VecZnx].
138 ///
139 /// # Arguments
140 ///
141 /// * `b`: [VmpPMat] on which the values are encoded.
142 /// * `a`: the vector of [VecZnx] to encode on the [VmpPMat].
143 /// * `buf`: scratch space, the size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
144 ///
145 /// The size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
146 ///
147 /// # Example
148 /// ```
149 /// use base2k::{Module, FFT64, Matrix3D, VmpPMat, VmpPMatOps, VecZnx, VecZnxOps, Free};
150 /// use std::cmp::min;
151 ///
152 /// let n: usize = 1024;
153 /// let module: Module = Module::new::<FFT64>(n);
154 /// let rows: usize = 5;
155 /// let cols: usize = 6;
156 ///
157 /// let mut vecznx: Vec<VecZnx>= Vec::new();
158 /// (0..rows).for_each(|_|{
159 /// vecznx.push(module.new_vec_znx(cols));
160 /// });
161 ///
162 /// let dble: Vec<&[i64]> = vecznx.iter().map(|v| v.data.as_slice()).collect();
163 ///
164 /// let mut buf: Vec<u8> = vec![u8::default(); module.vmp_prepare_tmp_bytes(rows, cols)];
165 ///
166 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
167 /// module.vmp_prepare_dblptr(&mut vmp_pmat, &dble, &mut buf);
168 ///
169 /// vmp_pmat.free();
170 /// module.free();
171 /// ```
172 fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &[&[i64]], buf: &mut [u8]);
173
174 /// Prepares the ith-row of [VmpPMat] from a vector of [VecZnx].
175 ///
176 /// # Arguments
177 ///
178 /// * `b`: [VmpPMat] on which the values are encoded.
179 /// * `a`: the vector of [VecZnx] to encode on the [VmpPMat].
180 /// * `row_i`: the index of the row to prepare.
181 /// * `buf`: scratch space, the size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
182 ///
183 /// The size of buf can be obtained with [VmpPMatOps::vmp_prepare_tmp_bytes].
184 /// /// # Example
185 /// ```
186 /// use base2k::{Module, FFT64, Matrix3D, VmpPMat, VmpPMatOps, VecZnx, VecZnxOps, Free};
187 /// use std::cmp::min;
188 ///
189 /// let n: usize = 1024;
190 /// let module: Module = Module::new::<FFT64>(n);
191 /// let rows: usize = 5;
192 /// let cols: usize = 6;
193 ///
194 /// let vecznx = vec![0i64; cols*n];
195 ///
196 /// let mut buf: Vec<u8> = vec![u8::default(); module.vmp_prepare_tmp_bytes(rows, cols)];
197 ///
198 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
199 /// module.vmp_prepare_row(&mut vmp_pmat, &vecznx, 0, &mut buf);
200 ///
201 /// vmp_pmat.free();
202 /// module.free();
203 /// ```
204 fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &[i64], row_i: usize, tmp_bytes: &mut [u8]);
205
206 /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft].
207 ///
208 /// # Arguments
209 ///
210 /// * `c_limbs`: number of limbs of the output [VecZnxDft].
211 /// * `a_limbs`: number of limbs of the input [VecZnx].
212 /// * `rows`: number of rows of the input [VmpPMat].
213 /// * `cols`: number of cols of the input [VmpPMat].
214 fn vmp_apply_dft_tmp_bytes(
215 &self,
216 c_limbs: usize,
217 a_limbs: usize,
218 rows: usize,
219 cols: usize,
220 ) -> usize;
221
222 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
223 ///
224 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
225 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
226 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
227 ///
228 /// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
229 /// `j` cols, the output is a [VecZnx] of `j` limbs.
230 ///
231 /// If there is a mismatch between the dimensions the largest valid ones are used.
232 ///
233 /// ```text
234 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
235 /// |h i j|
236 /// |k l m|
237 /// ```
238 /// where each element is a [VecZnxDft].
239 ///
240 /// # Arguments
241 ///
242 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
243 /// * `a`: the left operand [VecZnx] of the vector matrix product.
244 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
245 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_tmp_bytes].
246 ///
247 /// # Example
248 /// ```
249 /// use base2k::{Module, VecZnx, VecZnxOps, VecZnxDft, VmpPMat, VmpPMatOps, FFT64, Free};
250 ///
251 /// let n = 1024;
252 ///
253 /// let module: Module = Module::new::<FFT64>(n);
254 /// let limbs: usize = 5;
255 ///
256 /// let rows: usize = limbs;
257 /// let cols: usize = limbs + 1;
258 /// let c_limbs: usize = cols;
259 /// let a_limbs: usize = limbs;
260 /// let tmp_bytes: usize = module.vmp_apply_dft_tmp_bytes(c_limbs, a_limbs, rows, cols);
261 ///
262 /// let mut buf: Vec<u8> = vec![0; tmp_bytes];
263 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
264 ///
265 /// let a: VecZnx = module.new_vec_znx(limbs);
266 /// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols);
267 /// module.vmp_apply_dft(&mut c_dft, &a, &vmp_pmat, &mut buf);
268 ///
269 /// c_dft.free();
270 /// vmp_pmat.free();
271 /// module.free();
272 /// ```
273 fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]);
274
275 /// Returns the size of the stratch space necessary for [VmpPMatOps::vmp_apply_dft_to_dft].
276 ///
277 /// # Arguments
278 ///
279 /// * `c_limbs`: number of limbs of the output [VecZnxDft].
280 /// * `a_limbs`: number of limbs of the input [VecZnxDft].
281 /// * `rows`: number of rows of the input [VmpPMat].
282 /// * `cols`: number of cols of the input [VmpPMat].
283 fn vmp_apply_dft_to_dft_tmp_bytes(
284 &self,
285 c_limbs: usize,
286 a_limbs: usize,
287 rows: usize,
288 cols: usize,
289 ) -> usize;
290
291 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat].
292 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
293 ///
294 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
295 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
296 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
297 ///
298 /// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
299 /// `j` cols, the output is a [VecZnx] of `j` limbs.
300 ///
301 /// If there is a mismatch between the dimensions the largest valid ones are used.
302 ///
303 /// ```text
304 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
305 /// |h i j|
306 /// |k l m|
307 /// ```
308 /// where each element is a [VecZnxDft].
309 ///
310 /// # Arguments
311 ///
312 /// * `c`: the output of the vector matrix product, as a [VecZnxDft].
313 /// * `a`: the left operand [VecZnxDft] of the vector matrix product.
314 /// * `b`: the right operand [VmpPMat] of the vector matrix product.
315 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
316 ///
317 /// # Example
318 /// ```
319 /// use base2k::{Module, VecZnx, VecZnxDft, VmpPMat, VmpPMatOps, FFT64, Free};
320 ///
321 /// let n = 1024;
322 ///
323 /// let module: Module = Module::new::<FFT64>(n);
324 /// let limbs: usize = 5;
325 ///
326 /// let rows: usize = limbs;
327 /// let cols: usize = limbs + 1;
328 /// let c_limbs: usize = cols;
329 /// let a_limbs: usize = limbs;
330 /// let tmp_bytes: usize = module.vmp_apply_dft_to_dft_tmp_bytes(c_limbs, a_limbs, rows, cols);
331 ///
332 /// let mut buf: Vec<u8> = vec![0; tmp_bytes];
333 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
334 ///
335 /// let a_dft: VecZnxDft = module.new_vec_znx_dft(limbs);
336 /// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(cols);
337 /// module.vmp_apply_dft_to_dft(&mut c_dft, &a_dft, &vmp_pmat, &mut buf);
338 ///
339 /// a_dft.free();
340 /// c_dft.free();
341 /// vmp_pmat.free();
342 /// module.free();
343 /// ```
344 fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, buf: &mut [u8]);
345
346 /// Applies the vector matrix product [VecZnxDft] x [VmpPMat] in place.
347 /// The size of `buf` is given by [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
348 ///
349 /// A vector matrix product is equivalent to a sum of [crate::SvpPPolOps::svp_apply_dft]
350 /// where each [crate::Scalar] is a limb of the input [VecZnxDft] (equivalent to an [crate::SvpPPol])
351 /// and each vector a [VecZnxDft] (row) of the [VmpPMat].
352 ///
353 /// As such, given an input [VecZnx] of `i` limbs and a [VmpPMat] of `i` rows and
354 /// `j` cols, the output is a [VecZnx] of `j` limbs.
355 ///
356 /// If there is a mismatch between the dimensions the largest valid ones are used.
357 ///
358 /// ```text
359 /// |a b c d| x |e f g| = (a * |e f g| + b * |h i j| + c * |k l m|) = |n o p|
360 /// |h i j|
361 /// |k l m|
362 /// ```
363 /// where each element is a [VecZnxDft].
364 ///
365 /// # Arguments
366 ///
367 /// * `b`: the input and output of the vector matrix product, as a [VecZnxDft].
368 /// * `a`: the right operand [VmpPMat] of the vector matrix product.
369 /// * `buf`: scratch space, the size can be obtained with [VmpPMatOps::vmp_apply_dft_to_dft_tmp_bytes].
370 ///
371 /// # Example
372 /// ```
373 /// use base2k::{Module, VecZnx, VecZnxOps, VecZnxDft, VmpPMat, VmpPMatOps, FFT64, Free};
374 ///
375 /// let n = 1024;
376 ///
377 /// let module: Module = Module::new::<FFT64>(n);
378 /// let limbs: usize = 5;
379 ///
380 /// let rows: usize = limbs;
381 /// let cols: usize = limbs + 1;
382 /// let tmp_bytes: usize = module.vmp_apply_dft_to_dft_tmp_bytes(limbs, limbs, rows, cols);
383 ///
384 /// let mut buf: Vec<u8> = vec![0; tmp_bytes];
385 /// let a: VecZnx = module.new_vec_znx(limbs);
386 /// let mut vmp_pmat: VmpPMat = module.new_vmp_pmat(rows, cols);
387 ///
388 /// let mut c_dft: VecZnxDft = module.new_vec_znx_dft(limbs);
389 /// module.vmp_apply_dft_to_dft_inplace(&mut c_dft, &vmp_pmat, &mut buf);
390 ///
391 /// c_dft.free();
392 /// vmp_pmat.free();
393 /// module.free();
394 /// ```
395 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]);
396}
397
398impl VmpPMatOps for Module {
399 fn new_vmp_pmat(&self, rows: usize, cols: usize) -> VmpPMat {
400 unsafe {
401 VmpPMat {
402 data: vmp::new_vmp_pmat(self.0, rows as u64, cols as u64),
403 rows,
404 cols,
405 n: self.n(),
406 }
407 }
408 }
409
410 fn vmp_prepare_tmp_bytes(&self, rows: usize, cols: usize) -> usize {
411 unsafe { vmp::vmp_prepare_tmp_bytes(self.0, rows as u64, cols as u64) as usize }
412 }
413
414 fn vmp_prepare_contiguous(&self, b: &mut VmpPMat, a: &[i64], buf: &mut [u8]) {
415 unsafe {
416 vmp::vmp_prepare_contiguous(
417 self.0,
418 b.data(),
419 a.as_ptr(),
420 b.rows() as u64,
421 b.cols() as u64,
422 buf.as_mut_ptr(),
423 );
424 }
425 }
426
427 fn vmp_prepare_dblptr(&self, b: &mut VmpPMat, a: &[&[i64]], buf: &mut [u8]) {
428 let ptrs: Vec<*const i64> = a.iter().map(|v| v.as_ptr()).collect();
429 unsafe {
430 vmp::vmp_prepare_dblptr(
431 self.0,
432 b.data(),
433 ptrs.as_ptr(),
434 b.rows() as u64,
435 b.cols() as u64,
436 buf.as_mut_ptr(),
437 );
438 }
439 }
440
441 fn vmp_prepare_row(&self, b: &mut VmpPMat, a: &[i64], row_i: usize, buf: &mut [u8]) {
442 unsafe {
443 vmp::vmp_prepare_row(
444 self.0,
445 b.data(),
446 a.as_ptr(),
447 row_i as u64,
448 b.rows() as u64,
449 b.cols() as u64,
450 buf.as_mut_ptr(),
451 );
452 }
453 }
454
455 fn vmp_apply_dft_tmp_bytes(
456 &self,
457 c_limbs: usize,
458 a_limbs: usize,
459 rows: usize,
460 cols: usize,
461 ) -> usize {
462 unsafe {
463 vmp::vmp_apply_dft_tmp_bytes(
464 self.0,
465 c_limbs as u64,
466 a_limbs as u64,
467 rows as u64,
468 cols as u64,
469 ) as usize
470 }
471 }
472
473 fn vmp_apply_dft(&self, c: &mut VecZnxDft, a: &VecZnx, b: &VmpPMat, buf: &mut [u8]) {
474 unsafe {
475 vmp::vmp_apply_dft(
476 self.0,
477 c.0,
478 c.limbs() as u64,
479 a.as_ptr(),
480 a.limbs() as u64,
481 a.n() as u64,
482 b.data(),
483 b.rows() as u64,
484 b.cols() as u64,
485 buf.as_mut_ptr(),
486 )
487 }
488 }
489
490 fn vmp_apply_dft_to_dft_tmp_bytes(
491 &self,
492 c_limbs: usize,
493 a_limbs: usize,
494 rows: usize,
495 cols: usize,
496 ) -> usize {
497 unsafe {
498 vmp::vmp_apply_dft_to_dft_tmp_bytes(
499 self.0,
500 c_limbs as u64,
501 a_limbs as u64,
502 rows as u64,
503 cols as u64,
504 ) as usize
505 }
506 }
507
508 fn vmp_apply_dft_to_dft(&self, c: &mut VecZnxDft, a: &VecZnxDft, b: &VmpPMat, buf: &mut [u8]) {
509 unsafe {
510 vmp::vmp_apply_dft_to_dft(
511 self.0,
512 c.0,
513 c.limbs() as u64,
514 a.0,
515 a.limbs() as u64,
516 b.data(),
517 b.rows() as u64,
518 b.cols() as u64,
519 buf.as_mut_ptr(),
520 )
521 }
522 }
523
524 fn vmp_apply_dft_to_dft_inplace(&self, b: &mut VecZnxDft, a: &VmpPMat, buf: &mut [u8]) {
525 unsafe {
526 vmp::vmp_apply_dft_to_dft(
527 self.0,
528 b.0,
529 b.limbs() as u64,
530 b.0,
531 b.limbs() as u64,
532 a.data(),
533 a.rows() as u64,
534 a.cols() as u64,
535 buf.as_mut_ptr(),
536 )
537 }
538 }
539}
540
541/// A helper struture that stores a 3D matrix as a contiguous array.
542/// To be passed to [VmpPMatOps::vmp_prepare_contiguous].
543///
544/// rows: index of the i-th base2K power.
545/// cols: index of the j-th limb of the i-th row.
546/// n : polynomial degree.
547///
548/// A [Matrix3D] can be seen as a vector of [VecZnx].
549pub struct Matrix3D<T> {
550 pub data: Vec<T>,
551 pub rows: usize,
552 pub cols: usize,
553 pub n: usize,
554}
555
556impl<T: Default + Clone + std::marker::Copy> Matrix3D<T> {
557 /// Allocates a new [Matrix3D] with the respective dimensions.
558 ///
559 /// # Arguments
560 ///
561 /// * `rows`: the number of rows of the matrix.
562 /// * `cols`: the number of cols of the matrix.
563 /// # `n`: the size of each entry of the matrix.
564 ///
565 /// # Example
566 /// ```
567 /// use base2k::Matrix3D;
568 ///
569 /// let rows = 5; // #decomp
570 /// let cols = 5; // #limbs
571 /// let n = 1024; // #coeffs
572 ///
573 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
574 /// ```
575 pub fn new(rows: usize, cols: usize, n: usize) -> Self {
576 let size = rows * cols * n;
577 Self {
578 data: vec![T::default(); size],
579 rows,
580 cols,
581 n,
582 }
583 }
584
585 /// Returns a non-mutable reference to the entry (row, col) of the [Matrix3D].
586 /// The returned array is of size n.
587 ///
588 /// # Arguments
589 ///
590 /// * `row`: the index of the row.
591 /// * `col`: the index of the col.
592 ///
593 /// # Example
594 /// ```
595 /// use base2k::Matrix3D;
596 ///
597 /// let rows = 5; // #decomp
598 /// let cols = 5; // #limbs
599 /// let n = 1024; // #coeffs
600 ///
601 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
602 ///
603 /// let elem: &[i64] = mat.at(4, 4); // size n
604 /// ```
605 pub fn at(&self, row: usize, col: usize) -> &[T] {
606 assert!(row < self.rows && col < self.cols);
607 let idx: usize = row * (self.n * self.cols) + col * self.n;
608 &self.data[idx..idx + self.n]
609 }
610
611 /// Returns a mutable reference of the array at the (row, col) entry of the [Matrix3D].
612 /// The returned array is of size n.
613 ///
614 /// # Arguments
615 ///
616 /// * `row`: the index of the row.
617 /// * `col`: the index of the col.
618 ///
619 /// # Example
620 /// ```
621 /// use base2k::Matrix3D;
622 ///
623 /// let rows = 5; // #decomp
624 /// let cols = 5; // #limbs
625 /// let n = 1024; // #coeffs
626 ///
627 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
628 ///
629 /// let elem: &mut [i64] = mat.at_mut(4, 4); // size n
630 /// ```
631 pub fn at_mut(&mut self, row: usize, col: usize) -> &mut [T] {
632 assert!(row < self.rows && col < self.cols);
633 let idx: usize = row * (self.n * self.cols) + col * self.n;
634 &mut self.data[idx..idx + self.n]
635 }
636
637 /// Sets the entry \[row\] of the [Matrix3D].
638 /// Typicall this is used to assign a [VecZnx] to the i-th row
639 /// of the [Matrix3D].
640 ///
641 /// # Arguments
642 ///
643 /// * `row`: the index of the row.
644 /// * `a`: the data to encode onthe row.
645 ///
646 /// # Example
647 /// ```
648 /// use base2k::{Matrix3D, VecZnx};
649 ///
650 /// let rows = 5; // #decomp
651 /// let cols = 5; // #limbs
652 /// let n = 1024; // #coeffs
653 ///
654 /// let mut mat = Matrix3D::<i64>::new(rows, cols, n);
655 ///
656 /// let a: VecZnx = VecZnx::new(n, cols);
657 ///
658 /// mat.set_row(1, &a.data);
659 /// ```
660 pub fn set_row(&mut self, row: usize, a: &[T]) {
661 assert!(
662 row < self.rows,
663 "invalid argument row: row={} > self.rows={}",
664 row,
665 self.rows
666 );
667 let idx: usize = row * (self.n * self.cols);
668 let size: usize = min(a.len(), self.cols * self.n);
669 self.data[idx..idx + size].copy_from_slice(&a[..size]);
670 }
671}