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