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