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
411impl VecZnxOps for Module {
412 fn new_vec_znx(&self, limbs: usize) -> VecZnx {
413 VecZnx::new(self.n(), limbs)
414 }
415
416 // c <- a + b
417 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
418 unsafe {
419 vec_znx::vec_znx_add(
420 self.0,
421 c.as_mut_ptr(),
422 c.limbs() as u64,
423 c.n() as u64,
424 a.as_ptr(),
425 a.limbs() as u64,
426 a.n() as u64,
427 b.as_ptr(),
428 b.limbs() as u64,
429 b.n() as u64,
430 )
431 }
432 }
433
434 // b <- a + b
435 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
436 unsafe {
437 vec_znx::vec_znx_add(
438 self.0,
439 b.as_mut_ptr(),
440 b.limbs() as u64,
441 b.n() as u64,
442 a.as_ptr(),
443 a.limbs() as u64,
444 a.n() as u64,
445 b.as_ptr(),
446 b.limbs() as u64,
447 b.n() as u64,
448 )
449 }
450 }
451
452 // c <- a + b
453 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
454 unsafe {
455 vec_znx::vec_znx_sub(
456 self.0,
457 c.as_mut_ptr(),
458 c.limbs() as u64,
459 c.n() as u64,
460 a.as_ptr(),
461 a.limbs() as u64,
462 a.n() as u64,
463 b.as_ptr(),
464 b.limbs() as u64,
465 b.n() as u64,
466 )
467 }
468 }
469
470 // b <- a + b
471 fn vec_znx_sub_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
472 unsafe {
473 vec_znx::vec_znx_sub(
474 self.0,
475 b.as_mut_ptr(),
476 b.limbs() as u64,
477 b.n() as u64,
478 a.as_ptr(),
479 a.limbs() as u64,
480 a.n() as u64,
481 b.as_ptr(),
482 b.limbs() as u64,
483 b.n() as u64,
484 )
485 }
486 }
487
488 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx) {
489 unsafe {
490 vec_znx::vec_znx_negate(
491 self.0,
492 b.as_mut_ptr(),
493 b.limbs() as u64,
494 b.n() as u64,
495 a.as_ptr(),
496 a.limbs() as u64,
497 a.n() as u64,
498 )
499 }
500 }
501
502 fn vec_znx_negate_inplace(&self, a: &mut VecZnx) {
503 unsafe {
504 vec_znx::vec_znx_negate(
505 self.0,
506 a.as_mut_ptr(),
507 a.limbs() as u64,
508 a.n() as u64,
509 a.as_ptr(),
510 a.limbs() as u64,
511 a.n() as u64,
512 )
513 }
514 }
515
516 fn vec_znx_rotate(&self, k: i64, a: &mut VecZnx, b: &VecZnx) {
517 unsafe {
518 vec_znx::vec_znx_rotate(
519 self.0,
520 k,
521 a.as_mut_ptr(),
522 a.limbs() as u64,
523 a.n() as u64,
524 b.as_ptr(),
525 b.limbs() as u64,
526 b.n() as u64,
527 )
528 }
529 }
530
531 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx) {
532 unsafe {
533 vec_znx::vec_znx_rotate(
534 self.0,
535 k,
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_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
547 unsafe {
548 vec_znx::vec_znx_automorphism(
549 self.0,
550 k,
551 b.as_mut_ptr(),
552 b.limbs() as u64,
553 b.n() as u64,
554 a.as_ptr(),
555 a.limbs() as u64,
556 a.n() as u64,
557 );
558 }
559 }
560
561 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx) {
562 unsafe {
563 vec_znx::vec_znx_automorphism(
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}