poulpy/base2k/src/vec_znx.rs
base2k/src/vec_znx.rsRaw
1use crate::cast_mut_u8_to_mut_i64_slice;
2use crate::ffi::vec_znx;
3use crate::ffi::znx;
4use crate::{alias_mut_slice_to_vec, alloc_aligned};
5use crate::{Infos, Module};
6use itertools::izip;
7use std::cmp::min;
8
9/// [VecZnx] represents a vector of small norm polynomials of Zn\[X\] with [i64] coefficients.
10/// A [VecZnx] is composed of multiple Zn\[X\] polynomials stored in a single contiguous array
11/// in the memory.
12#[derive(Clone)]
13pub struct VecZnx {
14 /// Polynomial degree.
15 pub n: usize,
16 /// Polynomial coefficients, as a contiguous array. Each limb is equally spaced by n.
17 pub data: Vec<i64>,
18}
19
20impl VecZnx {
21 /// Allocates a new [VecZnx] composed of #limbs polynomials of Z\[X\].
22 pub fn new(n: usize, limbs: usize) -> Self {
23 Self {
24 n: n,
25 data: alloc_aligned::<i64>(n * limbs, 64),
26 }
27 }
28
29 /// Returns the minimum size of the [u8] array required to assign a
30 /// new backend array to a [VecZnx] through [VecZnx::from_bytes].
31 pub fn bytes(n: usize, limbs: usize) -> usize {
32 n * limbs * 8
33 }
34
35 /// Returns a new [VecZnx] with the provided data as backing array.
36 /// User must ensure that data is properly alligned and that
37 /// the size of data is at least equal to [Module::bytes_of_vec_znx].
38 pub fn from_bytes(n: usize, limbs: usize, buf: &mut [u8]) -> VecZnx {
39 let size = Self::bytes(n, limbs);
40 assert!(
41 buf.len() >= size,
42 "invalid buffer: buf.len()={} < self.buffer_size(n={}, limbs={})={}",
43 buf.len(),
44 n,
45 limbs,
46 size
47 );
48
49 VecZnx {
50 n: n,
51 data: alias_mut_slice_to_vec(cast_mut_u8_to_mut_i64_slice(&mut buf[..size])),
52 }
53 }
54
55 /// Copies the coefficients of `a` on the receiver.
56 /// Copy is done with the minimum size matching both backing arrays.
57 pub fn copy_from(&mut self, a: &VecZnx) {
58 let size = min(self.data.len(), a.data.len());
59 self.data[..size].copy_from_slice(&a.data[..size])
60 }
61
62 /// Returns a non-mutable pointer to the backing array of the [VecZnx].
63 pub fn as_ptr(&self) -> *const i64 {
64 self.data.as_ptr()
65 }
66
67 /// Returns a mutable pointer to the backing array of the [VecZnx].
68 pub fn as_mut_ptr(&mut self) -> *mut i64 {
69 self.data.as_mut_ptr()
70 }
71
72 /// Returns a non-mutable reference to the i-th limb of the [VecZnx].
73 pub fn at(&self, i: usize) -> &[i64] {
74 &self.data[i * self.n..(i + 1) * self.n]
75 }
76
77 /// Returns a mutable reference to the i-th limb of the [VecZnx].
78 pub fn at_mut(&mut self, i: usize) -> &mut [i64] {
79 &mut self.data[i * self.n..(i + 1) * self.n]
80 }
81
82 /// Returns a non-mutable pointer to the i-th limb of the [VecZnx].
83 pub fn at_ptr(&self, i: usize) -> *const i64 {
84 &self.data[i * self.n] as *const i64
85 }
86
87 /// Returns a mutable pointer to the i-th limb of the [VecZnx].
88 pub fn at_mut_ptr(&mut self, i: usize) -> *mut i64 {
89 &mut self.data[i * self.n] as *mut i64
90 }
91
92 /// Zeroes the backing array of the [VecZnx].
93 pub fn zero(&mut self) {
94 unsafe { znx::znx_zero_i64_ref(self.data.len() as u64, self.data.as_mut_ptr()) }
95 }
96
97 /// Normalizes the [VecZnx], ensuring all coefficients are in the interval \[-2^log_base2k, 2^log_base2k].
98 ///
99 /// # Arguments
100 ///
101 /// * `log_base2k`: the base two logarithm of the base to reduce to.
102 /// * `carry`: scratch space of size at least self.n()<<3.
103 ///
104 /// # Panics
105 ///
106 /// The method will panic if carry.len() < self.data.len()*8.
107 ///
108 /// # Example
109 /// ```
110 /// use base2k::{VecZnx, Encoding, Infos, alloc_aligned};
111 /// use itertools::izip;
112 /// use sampling::source::Source;
113 ///
114 /// let n: usize = 8; // polynomial degree
115 /// let log_base2k: usize = 17; // base two logarithm of the coefficients decomposition
116 /// let limbs: usize = 5; // number of limbs (i.e. can store coeffs in the range +/- 2^{limbs * log_base2k - 1})
117 /// let log_k: usize = limbs * log_base2k - 5;
118 /// let mut a: VecZnx = VecZnx::new(n, limbs);
119 /// let mut carry: Vec<u8> = alloc_aligned::<u8>(a.n()<<3, 64);
120 /// let mut have: Vec<i64> = alloc_aligned::<i64>(a.n(), 64);
121 /// let mut source = Source::new([1; 32]);
122 ///
123 /// // Populates the first limb of the of polynomials with random i64 values.
124 /// have.iter_mut().for_each(|x| {
125 /// *x = source
126 /// .next_u64n(u64::MAX, u64::MAX)
127 /// .wrapping_sub(u64::MAX / 2 + 1) as i64;
128 /// });
129 /// a.encode_vec_i64(log_base2k, log_k, &have, 63);
130 /// a.normalize(log_base2k, &mut carry);
131 ///
132 /// // Ensures normalized values are in the range +/- 2^{log_base2k-1}
133 /// let base_half = 1 << (log_base2k - 1);
134 /// a.data
135 /// .iter()
136 /// .for_each(|x| assert!(x.abs() <= base_half, "|x|={} > 2^(k-1)={}", x, base_half));
137 ///
138 /// // Ensures reconstructed normalized values are equal to non-normalized values.
139 /// let mut want = alloc_aligned::<i64>(a.n(), 64);
140 /// a.decode_vec_i64(log_base2k, log_k, &mut want);
141 /// izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
142 /// ```
143 pub fn normalize(&mut self, log_base2k: usize, carry: &mut [u8]) {
144 assert!(
145 carry.len() >= self.n * 8,
146 "invalid carry: carry.len()={} < self.n()={}",
147 carry.len(),
148 self.n()
149 );
150
151 let carry_i64: &mut [i64] = cast_mut_u8_to_mut_i64_slice(carry);
152
153 unsafe {
154 znx::znx_zero_i64_ref(self.n() as u64, carry_i64.as_mut_ptr());
155 (0..self.limbs()).rev().for_each(|i| {
156 znx::znx_normalize(
157 self.n as u64,
158 log_base2k as u64,
159 self.at_mut_ptr(i),
160 carry_i64.as_mut_ptr(),
161 self.at_mut_ptr(i),
162 carry_i64.as_mut_ptr(),
163 )
164 });
165 }
166 }
167
168 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each limb.
169 ///
170 /// # Arguments
171 ///
172 /// * `k`: the power to which to map each coefficients.
173 /// * `limbs`: the number of limbs on which to apply the mapping.
174 ///
175 /// # Panics
176 ///
177 /// The method will panic if the argument `limbs` is greater than `self.limbs()`.
178 ///
179 /// # Example
180 /// ```
181 /// use base2k::{VecZnx, Encoding, Infos};
182 /// use itertools::izip;
183 ///
184 /// let n: usize = 8; // polynomial degree
185 /// let mut a: VecZnx = VecZnx::new(n, 2);
186 /// let mut b: VecZnx = VecZnx::new(n, 2);
187 ///
188 /// (0..a.limbs()).for_each(|i|{
189 /// a.at_mut(i).iter_mut().enumerate().for_each(|(i, x)|{
190 /// *x = i as i64
191 /// })
192 /// });
193 ///
194 /// b.copy_from(&a);
195 ///
196 /// a.automorphism_inplace(-1, 1); // X^i -> X^(-i)
197 /// let limb = b.at_mut(0);
198 /// (1..limb.len()).for_each(|i|{
199 /// limb[n-i] = -(i as i64)
200 /// });
201 /// izip!(a.data.iter(), b.data.iter()).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
202 /// ```
203 pub fn automorphism_inplace(&mut self, k: i64, limbs: usize) {
204 assert!(
205 limbs <= self.limbs(),
206 "invalid limbs argument: limbs={} > self.limbs()={}",
207 limbs,
208 self.limbs()
209 );
210 unsafe {
211 (0..limbs).for_each(|i| {
212 znx::znx_automorphism_inplace_i64(self.n as u64, k, self.at_mut_ptr(i))
213 })
214 }
215 }
216
217 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each limb.
218 ///
219 /// # Arguments
220 ///
221 /// * `a`: the receiver.
222 /// * `k`: the power to which to map each coefficients.
223 /// * `limbs`: the number of limbs on which to apply the mapping.
224 ///
225 /// # Panics
226 ///
227 /// The method will panic if the argument `limbs` is greater than `self.limbs()` or `a.limbs()`.
228 ///
229 /// # Example
230 /// ```
231 /// use base2k::{VecZnx, Encoding, Infos};
232 /// use itertools::izip;
233 ///
234 /// let n: usize = 8; // polynomial degree
235 /// let mut a: VecZnx = VecZnx::new(n, 2);
236 /// let mut b: VecZnx = VecZnx::new(n, 2);
237 /// let mut c: VecZnx = VecZnx::new(n, 2);
238 ///
239 /// (0..a.limbs()).for_each(|i|{
240 /// a.at_mut(i).iter_mut().enumerate().for_each(|(i, x)|{
241 /// *x = i as i64
242 /// })
243 /// });
244 ///
245 /// a.automorphism(&mut b, -1, 1); // X^i -> X^(-i)
246 /// let limb = c.at_mut(0);
247 /// (1..limb.len()).for_each(|i|{
248 /// limb[n-i] = -(i as i64)
249 /// });
250 /// izip!(b.data.iter(), c.data.iter()).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
251 /// ```
252 pub fn automorphism(&mut self, a: &mut VecZnx, k: i64, limbs: usize) {
253 assert!(
254 limbs <= self.limbs(),
255 "invalid limbs argument: limbs={} > self.limbs()={}",
256 limbs,
257 self.limbs()
258 );
259 assert!(
260 limbs <= a.limbs(),
261 "invalid limbs argument: limbs={} > a.limbs()={}",
262 limbs,
263 a.limbs()
264 );
265 unsafe {
266 (0..limbs).for_each(|i| {
267 znx::znx_automorphism_i64(self.n as u64, k, a.at_mut_ptr(i), self.at_ptr(i))
268 })
269 }
270 }
271
272 /// Truncates the precision of the [VecZnx] by k bits.
273 ///
274 /// # Arguments
275 ///
276 /// * `log_base2k`: the base two logarithm of the coefficients decomposition.
277 /// * `k`: the number of bits of precision to drop.
278 pub fn trunc_pow2(&mut self, log_base2k: usize, k: usize) {
279 if k == 0 {
280 return;
281 }
282
283 self.data
284 .truncate((self.limbs() - k / log_base2k) * self.n());
285
286 let k_rem: usize = k % log_base2k;
287
288 if k_rem != 0 {
289 let mask: i64 = ((1 << (log_base2k - k_rem - 1)) - 1) << k_rem;
290 self.at_mut(self.limbs() - 1)
291 .iter_mut()
292 .for_each(|x: &mut i64| *x &= mask)
293 }
294 }
295
296 /// Right shifts the coefficients by k bits.
297 ///
298 /// # Arguments
299 ///
300 /// * `log_base2k`: the base two logarithm of the coefficients decomposition.
301 /// * `k`: the shift amount.
302 /// * `carry`: scratch space of size at least equal to self.n() * self.limbs() << 3.
303 ///
304 /// # Panics
305 ///
306 /// The method will panic if carry.len() < self.n() * self.limbs() << 3.
307 pub fn rsh(&mut self, log_base2k: usize, k: usize, carry: &mut [u8]) {
308 assert!(
309 carry.len() >> 3 >= self.n(),
310 "invalid carry: carry.len()/8={} < self.n()={}",
311 carry.len() >> 3,
312 self.n()
313 );
314
315 let limbs: usize = self.limbs();
316 let limbs_steps: usize = k / log_base2k;
317
318 self.data.rotate_right(self.n * limbs_steps);
319 unsafe {
320 znx::znx_zero_i64_ref((self.n * limbs_steps) as u64, self.data.as_mut_ptr());
321 }
322
323 let k_rem = k % log_base2k;
324
325 if k_rem != 0 {
326 let carry_i64: &mut [i64] = cast_mut_u8_to_mut_i64_slice(carry);
327
328 unsafe {
329 znx::znx_zero_i64_ref(self.n() as u64, carry_i64.as_mut_ptr());
330 }
331
332 let mask: i64 = (1 << k_rem) - 1;
333 let log_base2k: usize = log_base2k;
334
335 (limbs_steps..limbs).for_each(|i| {
336 izip!(carry_i64.iter_mut(), self.at_mut(i).iter_mut()).for_each(|(ci, xi)| {
337 *xi += *ci << log_base2k;
338 *ci = *xi & mask;
339 *xi /= 1 << k_rem;
340 });
341 })
342 }
343 }
344
345 /// If self.n() > a.n(): Extracts X^{i*self.n()/a.n()} -> X^{i}.
346 /// If self.n() < a.n(): Extracts X^{i} -> X^{i*a.n()/self.n()}.
347 ///
348 /// # Arguments
349 ///
350 /// * `a`: the receiver polynomial in which the extracted coefficients are stored.
351 pub fn switch_degree(&self, a: &mut VecZnx) {
352 let (n_in, n_out) = (self.n(), a.n());
353 let (gap_in, gap_out): (usize, usize);
354
355 if n_in > n_out {
356 (gap_in, gap_out) = (n_in / n_out, 1)
357 } else {
358 (gap_in, gap_out) = (1, n_out / n_in);
359 a.zero();
360 }
361
362 let limbs = min(self.limbs(), a.limbs());
363
364 (0..limbs).for_each(|i| {
365 izip!(
366 self.at(i).iter().step_by(gap_in),
367 a.at_mut(i).iter_mut().step_by(gap_out)
368 )
369 .for_each(|(x_in, x_out)| *x_out = *x_in);
370 });
371 }
372
373 pub fn print_limbs(&self, limbs: usize, n: usize) {
374 (0..limbs).for_each(|i| println!("{}: {:?}", i, &self.at(i)[..n]))
375 }
376}
377
378pub trait VecZnxOps {
379 /// Allocates a new [VecZnx].
380 ///
381 /// # Arguments
382 ///
383 /// * `limbs`: the number of limbs.
384 fn new_vec_znx(&self, limbs: usize) -> VecZnx;
385
386 /// Returns the minimum number of bytes necessary to allocate
387 /// a new [VecZnx] through [VecZnx::from_bytes].
388 fn bytes_of_vec_znx(&self, limbs: usize) -> usize;
389
390 /// c <- a + b.
391 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
392
393 /// b <- b + a.
394 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx);
395
396 /// c <- a - b.
397 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
398
399 /// b <- b - a.
400 fn vec_znx_sub_inplace(&self, b: &mut VecZnx, a: &VecZnx);
401
402 /// b <- -a.
403 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx);
404
405 /// b <- -b.
406 fn vec_znx_negate_inplace(&self, a: &mut VecZnx);
407
408 /// b <- a * X^k (mod X^{n} + 1)
409 fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
410
411 /// a <- a * X^k (mod X^{n} + 1)
412 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx);
413
414 /// b <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
415 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
416
417 /// a <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
418 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx);
419
420 /// Splits b into subrings and copies them them into a.
421 ///
422 /// # Panics
423 ///
424 /// This method requires that all [VecZnx] of b have the same ring degree
425 /// and that b.n() * b.len() <= a.n()
426 fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx);
427
428 /// Merges the subrings a into b.
429 ///
430 /// # Panics
431 ///
432 /// This method requires that all [VecZnx] of a have the same ring degree
433 /// and that a.n() * a.len() <= b.n()
434 fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>);
435}
436
437impl VecZnxOps for Module {
438 fn new_vec_znx(&self, limbs: usize) -> VecZnx {
439 VecZnx::new(self.n(), limbs)
440 }
441
442 fn bytes_of_vec_znx(&self, limbs: usize) -> usize {
443 self.n() * limbs * 8
444 }
445
446 // c <- a + b
447 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
448 unsafe {
449 vec_znx::vec_znx_add(
450 self.0,
451 c.as_mut_ptr(),
452 c.limbs() as u64,
453 c.n() as u64,
454 a.as_ptr(),
455 a.limbs() as u64,
456 a.n() as u64,
457 b.as_ptr(),
458 b.limbs() as u64,
459 b.n() as u64,
460 )
461 }
462 }
463
464 // b <- a + b
465 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
466 unsafe {
467 vec_znx::vec_znx_add(
468 self.0,
469 b.as_mut_ptr(),
470 b.limbs() as u64,
471 b.n() as u64,
472 a.as_ptr(),
473 a.limbs() as u64,
474 a.n() as u64,
475 b.as_ptr(),
476 b.limbs() as u64,
477 b.n() as u64,
478 )
479 }
480 }
481
482 // c <- a + b
483 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
484 unsafe {
485 vec_znx::vec_znx_sub(
486 self.0,
487 c.as_mut_ptr(),
488 c.limbs() as u64,
489 c.n() as u64,
490 a.as_ptr(),
491 a.limbs() as u64,
492 a.n() as u64,
493 b.as_ptr(),
494 b.limbs() as u64,
495 b.n() as u64,
496 )
497 }
498 }
499
500 // b <- a + b
501 fn vec_znx_sub_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
502 unsafe {
503 vec_znx::vec_znx_sub(
504 self.0,
505 b.as_mut_ptr(),
506 b.limbs() as u64,
507 b.n() as u64,
508 a.as_ptr(),
509 a.limbs() as u64,
510 a.n() as u64,
511 b.as_ptr(),
512 b.limbs() as u64,
513 b.n() as u64,
514 )
515 }
516 }
517
518 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx) {
519 unsafe {
520 vec_znx::vec_znx_negate(
521 self.0,
522 b.as_mut_ptr(),
523 b.limbs() as u64,
524 b.n() as u64,
525 a.as_ptr(),
526 a.limbs() as u64,
527 a.n() as u64,
528 )
529 }
530 }
531
532 fn vec_znx_negate_inplace(&self, a: &mut VecZnx) {
533 unsafe {
534 vec_znx::vec_znx_negate(
535 self.0,
536 a.as_mut_ptr(),
537 a.limbs() as u64,
538 a.n() as u64,
539 a.as_ptr(),
540 a.limbs() as u64,
541 a.n() as u64,
542 )
543 }
544 }
545
546 fn vec_znx_rotate(&self, k: i64, a: &mut VecZnx, b: &VecZnx) {
547 unsafe {
548 vec_znx::vec_znx_rotate(
549 self.0,
550 k,
551 a.as_mut_ptr(),
552 a.limbs() as u64,
553 a.n() as u64,
554 b.as_ptr(),
555 b.limbs() as u64,
556 b.n() as u64,
557 )
558 }
559 }
560
561 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx) {
562 unsafe {
563 vec_znx::vec_znx_rotate(
564 self.0,
565 k,
566 a.as_mut_ptr(),
567 a.limbs() as u64,
568 a.n() as u64,
569 a.as_ptr(),
570 a.limbs() as u64,
571 a.n() as u64,
572 )
573 }
574 }
575
576 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
577 unsafe {
578 vec_znx::vec_znx_automorphism(
579 self.0,
580 k,
581 b.as_mut_ptr(),
582 b.limbs() as u64,
583 b.n() as u64,
584 a.as_ptr(),
585 a.limbs() as u64,
586 a.n() as u64,
587 );
588 }
589 }
590
591 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx) {
592 unsafe {
593 vec_znx::vec_znx_automorphism(
594 self.0,
595 k,
596 a.as_mut_ptr(),
597 a.limbs() as u64,
598 a.n() as u64,
599 a.as_ptr(),
600 a.limbs() as u64,
601 a.n() as u64,
602 );
603 }
604 }
605
606 fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx) {
607 let (n_in, n_out) = (a.n(), b[0].n());
608
609 assert!(
610 n_out < n_in,
611 "invalid a: output ring degree should be smaller"
612 );
613 b[1..].iter().for_each(|bi| {
614 assert_eq!(
615 bi.n(),
616 n_out,
617 "invalid input a: all VecZnx must have the same degree"
618 )
619 });
620
621 b.iter_mut().enumerate().for_each(|(i, bi)| {
622 if i == 0 {
623 a.switch_degree(bi);
624 self.vec_znx_rotate(-1, buf, a);
625 } else {
626 buf.switch_degree(bi);
627 self.vec_znx_rotate_inplace(-1, buf);
628 }
629 })
630 }
631
632 fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>) {
633 let (n_in, n_out) = (b.n(), a[0].n());
634
635 assert!(
636 n_out < n_in,
637 "invalid a: output ring degree should be smaller"
638 );
639 a[1..].iter().for_each(|ai| {
640 assert_eq!(
641 ai.n(),
642 n_out,
643 "invalid input a: all VecZnx must have the same degree"
644 )
645 });
646
647 a.iter().enumerate().for_each(|(_, ai)| {
648 ai.switch_degree(b);
649 self.vec_znx_rotate_inplace(-1, b);
650 });
651
652 self.vec_znx_rotate_inplace(a.len() as i64, b);
653 }
654}