poulpy/base2k/src/vec_znx.rs
base2k/src/vec_znx.rsRaw
1use crate::cast_mut_u8_to_mut_i64_slice;
2use crate::ffi::znx::{
3 znx_automorphism_i64, znx_automorphism_inplace_i64, znx_normalize, znx_zero_i64_ref,
4};
5use crate::module::Module;
6use itertools::izip;
7use rand_distr::{Distribution, Normal};
8use sampling::source::Source;
9use std::cmp::min;
10
11impl Module {
12 pub fn new_vec_znx(&self, log_base2k: usize, limbs: usize) -> VecZnx {
13 VecZnx::new(self.n(), log_base2k, limbs)
14 }
15}
16
17#[derive(Clone)]
18pub struct VecZnx {
19 pub n: usize,
20 pub log_base2k: usize,
21 pub data: Vec<i64>,
22}
23
24impl VecZnx {
25 pub fn new(n: usize, log_base2k: usize, limbs: usize) -> Self {
26 Self {
27 n: n,
28 log_base2k: log_base2k,
29 data: vec![i64::default(); Self::buffer_size(n, limbs)],
30 }
31 }
32
33 pub fn buffer_size(n: usize, limbs: usize) -> usize {
34 n * limbs
35 }
36
37 pub fn from_buffer(&mut self, n: usize, log_base2k: usize, limbs: usize, buf: &[i64]) {
38 let size = Self::buffer_size(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 self.n = n;
48 self.log_base2k = log_base2k;
49 self.data = Vec::from(&buf[..size])
50 }
51
52 pub fn log_n(&self) -> u64 {
53 (u64::BITS - (self.n - 1).leading_zeros()) as _
54 }
55
56 pub fn n(&self) -> usize {
57 self.n
58 }
59
60 pub fn limbs(&self) -> usize {
61 self.data.len() / self.n
62 }
63
64 pub fn copy_from(&mut self, a: &VecZnx) {
65 let size = min(self.data.len(), a.data.len());
66 self.data[..size].copy_from_slice(&a.data[..size])
67 }
68
69 pub fn as_ptr(&self) -> *const i64 {
70 self.data.as_ptr()
71 }
72
73 pub fn as_mut_ptr(&mut self) -> *mut i64 {
74 self.data.as_mut_ptr()
75 }
76
77 pub fn at(&self, i: usize) -> &[i64] {
78 &self.data[i * self.n..(i + 1) * self.n]
79 }
80
81 pub fn at_ptr(&self, i: usize) -> *const i64 {
82 &self.data[i * self.n] as *const i64
83 }
84
85 pub fn at_mut_ptr(&mut self, i: usize) -> *mut i64 {
86 &mut self.data[i * self.n] as *mut i64
87 }
88
89 pub fn at_mut(&mut self, i: usize) -> &mut [i64] {
90 &mut self.data[i * self.n..(i + 1) * self.n]
91 }
92
93 pub fn zero(&mut self) {
94 unsafe { znx_zero_i64_ref(self.data.len() as u64, self.data.as_mut_ptr()) }
95 }
96
97 pub fn from_i64(&mut self, data: &[i64], log_max: usize, log_k: usize) {
98 let limbs: usize = (log_k + self.log_base2k - 1) / self.log_base2k;
99
100 assert!(limbs <= self.limbs(), "invalid argument log_k: (log_k + self.log_base2k - 1)/self.log_base2k={} > self.limbs()={}", limbs, self.limbs());
101
102 let size: usize = min(data.len(), self.n());
103 let log_k_rem: usize = self.log_base2k - (log_k % self.log_base2k);
104
105 // If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
106 // values on the last limb.
107 // Else we decompose values base2k.
108 if log_max + log_k_rem < 63 || log_k_rem == self.log_base2k {
109 (0..limbs - 1).for_each(|i| unsafe {
110 znx_zero_i64_ref(size as u64, self.at_mut(i).as_mut_ptr());
111 });
112 self.at_mut(self.limbs() - 1)[..size].copy_from_slice(&data[..size]);
113 } else {
114 let mask: i64 = (1 << self.log_base2k) - 1;
115 let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
116
117 (0..steps).for_each(|i| unsafe {
118 znx_zero_i64_ref(size as u64, self.at_mut(i).as_mut_ptr());
119 });
120
121 (limbs - steps..limbs)
122 .rev()
123 .enumerate()
124 .for_each(|(i, i_rev)| {
125 let shift: usize = i * self.log_base2k;
126 izip!(self.at_mut(i_rev)[..size].iter_mut(), data[..size].iter())
127 .for_each(|(y, x)| *y = (x >> shift) & mask);
128 })
129 }
130
131 // Case where self.prec % self.k != 0.
132 if log_k_rem != self.log_base2k {
133 let limbs = self.limbs();
134 let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
135 (limbs - steps..limbs).rev().for_each(|i| {
136 self.at_mut(i)[..size]
137 .iter_mut()
138 .for_each(|x| *x <<= log_k_rem);
139 })
140 }
141 }
142
143 pub fn from_i64_single(&mut self, i: usize, value: i64, log_max: usize, log_k: usize) {
144 assert!(i < self.n());
145 let limbs: usize = (log_k + self.log_base2k - 1) / self.log_base2k;
146 assert!(limbs <= self.limbs(), "invalid argument log_k: (log_k + self.log_base2k - 1)/self.log_base2k={} > self.limbs()={}", limbs, self.limbs());
147 let log_k_rem: usize = self.log_base2k - (log_k % self.log_base2k);
148 let limbs = self.limbs();
149
150 // If 2^{log_base2k} * 2^{log_k_rem} < 2^{63}-1, then we can simply copy
151 // values on the last limb.
152 // Else we decompose values base2k.
153 if log_max + log_k_rem < 63 || log_k_rem == self.log_base2k {
154 (0..limbs - 1).for_each(|j| self.at_mut(j)[i] = 0);
155
156 self.at_mut(self.limbs() - 1)[i] = value;
157 } else {
158 let mask: i64 = (1 << self.log_base2k) - 1;
159 let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
160
161 (0..limbs - steps).for_each(|j| self.at_mut(j)[i] = 0);
162
163 (limbs - steps..limbs)
164 .rev()
165 .enumerate()
166 .for_each(|(j, j_rev)| {
167 self.at_mut(j_rev)[i] = (value >> (j * self.log_base2k)) & mask;
168 })
169 }
170
171 // Case where self.prec % self.k != 0.
172 if log_k_rem != self.log_base2k {
173 let limbs = self.limbs();
174 let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
175 (limbs - steps..limbs).rev().for_each(|j| {
176 self.at_mut(j)[i] <<= log_k_rem;
177 })
178 }
179 }
180
181 pub fn normalize(&mut self, carry: &mut [u8]) {
182 assert!(
183 carry.len() >= self.n * 8,
184 "invalid carry: carry.len()={} < self.n()={}",
185 carry.len(),
186 self.n()
187 );
188
189 let carry_i64: &mut [i64] = cast_mut_u8_to_mut_i64_slice(carry);
190
191 unsafe {
192 znx_zero_i64_ref(self.n() as u64, carry_i64.as_mut_ptr());
193 (0..self.limbs()).rev().for_each(|i| {
194 znx_normalize(
195 self.n as u64,
196 self.log_base2k as u64,
197 self.at_mut_ptr(i),
198 carry_i64.as_mut_ptr(),
199 self.at_mut_ptr(i),
200 carry_i64.as_mut_ptr(),
201 )
202 });
203 }
204 }
205
206 pub fn to_i64(&self, data: &mut [i64], log_k: usize) {
207 let limbs: usize = (log_k + self.log_base2k - 1) / self.log_base2k;
208 assert!(
209 data.len() >= self.n,
210 "invalid data: data.len()={} < self.n()={}",
211 data.len(),
212 self.n
213 );
214 data.copy_from_slice(self.at(0));
215 let rem: usize = self.log_base2k - (log_k % self.log_base2k);
216 (1..limbs).for_each(|i| {
217 if i == limbs - 1 && rem != self.log_base2k {
218 let k_rem: usize = self.log_base2k - rem;
219 izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
220 *y = (*y << k_rem) + (x >> rem);
221 });
222 } else {
223 izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
224 *y = (*y << self.log_base2k) + x;
225 });
226 }
227 })
228 }
229
230 pub fn to_i64_single(&self, i: usize, log_k: usize) -> i64 {
231 let limbs: usize = (log_k + self.log_base2k - 1) / self.log_base2k;
232 assert!(i < self.n());
233 let mut res: i64 = self.data[i];
234 let rem: usize = self.log_base2k - (log_k % self.log_base2k);
235 (1..limbs).for_each(|i| {
236 let x = self.data[i * self.n];
237 if i == limbs - 1 && rem != self.log_base2k {
238 let k_rem: usize = self.log_base2k - rem;
239 res = (res << k_rem) + (x >> rem);
240 } else {
241 res = (res << self.log_base2k) + x;
242 }
243 });
244 res
245 }
246
247 pub fn automorphism_inplace(&mut self, gal_el: i64) {
248 unsafe {
249 (0..self.limbs()).for_each(|i| {
250 znx_automorphism_inplace_i64(self.n as u64, gal_el, self.at_mut_ptr(i))
251 })
252 }
253 }
254 pub fn automorphism(&mut self, gal_el: i64, a: &mut VecZnx) {
255 unsafe {
256 (0..self.limbs()).for_each(|i| {
257 znx_automorphism_i64(self.n as u64, gal_el, a.at_mut_ptr(i), self.at_ptr(i))
258 })
259 }
260 }
261
262 pub fn fill_uniform(&mut self, source: &mut Source, log_k: usize) {
263 let mut base2k: u64 = 1 << self.log_base2k;
264 let mut mask: u64 = base2k - 1;
265 let mut base2k_half: i64 = (base2k >> 1) as i64;
266
267 let size: usize = self.n() * (self.limbs() - 1);
268
269 self.data[..size]
270 .iter_mut()
271 .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
272
273 let log_base2k_rem: usize = log_k % self.log_base2k;
274
275 if log_base2k_rem != 0 {
276 base2k = 1 << log_base2k_rem;
277 mask = (base2k - 1) << (self.log_base2k - log_base2k_rem);
278 base2k_half = ((mask >> 1) + 1) as i64;
279 }
280
281 self.data[size..]
282 .iter_mut()
283 .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
284 }
285
286 pub fn add_dist_f64<T: Distribution<f64>>(
287 &mut self,
288 source: &mut Source,
289 dist: T,
290 bound: f64,
291 log_k: usize,
292 ) {
293 let log_base2k_rem: usize = log_k % self.log_base2k;
294
295 if log_base2k_rem != 0 {
296 self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
297 let mut dist_f64: f64 = dist.sample(source);
298 while dist_f64.abs() > bound {
299 dist_f64 = dist.sample(source)
300 }
301 *a += (dist_f64.round() as i64) << log_base2k_rem
302 });
303 } else {
304 self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
305 let mut dist_f64: f64 = dist.sample(source);
306 while dist_f64.abs() > bound {
307 dist_f64 = dist.sample(source)
308 }
309 *a += dist_f64.round() as i64
310 });
311 }
312 }
313
314 pub fn add_normal(&mut self, source: &mut Source, sigma: f64, bound: f64, log_k: usize) {
315 self.add_dist_f64(source, Normal::new(0.0, sigma).unwrap(), bound, log_k);
316 }
317
318 pub fn trunc_pow2(&mut self, k: usize) {
319 if k == 0 {
320 return;
321 }
322
323 self.data
324 .truncate((self.limbs() - k / self.log_base2k) * self.n());
325
326 let k_rem: usize = k % self.log_base2k;
327
328 if k_rem != 0 {
329 let mask: i64 = ((1 << (self.log_base2k - k_rem - 1)) - 1) << k_rem;
330 self.at_mut(self.limbs() - 1)
331 .iter_mut()
332 .for_each(|x: &mut i64| *x &= mask)
333 }
334 }
335
336 pub fn rsh(&mut self, k: usize, carry: &mut [u8]) {
337 assert!(
338 carry.len() >> 3 >= self.n(),
339 "invalid carry: carry.len()/8={} < self.n()={}",
340 carry.len() >> 3,
341 self.n()
342 );
343
344 let limbs: usize = self.limbs();
345 let limbs_steps: usize = k / self.log_base2k;
346
347 self.data.rotate_right(self.n * limbs_steps);
348 unsafe {
349 znx_zero_i64_ref((self.n * limbs_steps) as u64, self.data.as_mut_ptr());
350 }
351
352 let k_rem = k % self.log_base2k;
353
354 if k_rem != 0 {
355 let carry_i64: &mut [i64] = cast_mut_u8_to_mut_i64_slice(carry);
356
357 unsafe {
358 znx_zero_i64_ref(self.n() as u64, carry_i64.as_mut_ptr());
359 }
360
361 let mask: i64 = (1 << k_rem) - 1;
362 let log_base2k: usize = self.log_base2k;
363
364 (limbs_steps..limbs).for_each(|i| {
365 izip!(carry_i64.iter_mut(), self.at_mut(i).iter_mut()).for_each(|(ci, xi)| {
366 *xi += *ci << log_base2k;
367 *ci = *xi & mask;
368 *xi /= 1 << k_rem;
369 });
370 })
371 }
372 }
373
374 pub fn switch_degree(&self, a: &mut VecZnx) {
375 let (n_in, n_out) = (self.n(), a.n());
376 let (gap_in, gap_out): (usize, usize);
377
378 if n_in > n_out {
379 (gap_in, gap_out) = (n_in / n_out, 1)
380 } else {
381 (gap_in, gap_out) = (1, n_out / n_in);
382 a.zero();
383 }
384
385 let limbs = min(self.limbs(), a.limbs());
386
387 (0..limbs).for_each(|i| {
388 izip!(
389 self.at(i).iter().step_by(gap_in),
390 a.at_mut(i).iter_mut().step_by(gap_out)
391 )
392 .for_each(|(x_in, x_out)| *x_out = *x_in);
393 });
394 }
395
396 pub fn print_limbs(&self, limbs: usize, n: usize) {
397 (0..limbs).for_each(|i| println!("{}: {:?}", i, &self.at(i)[..n]))
398 }
399}
400
401#[cfg(test)]
402mod tests {
403 use crate::VecZnx;
404 use itertools::izip;
405 use sampling::source::Source;
406
407 #[test]
408 fn test_set_get_i64_lo_norm() {
409 let n: usize = 8;
410 let log_base2k: usize = 17;
411 let limbs: usize = 5;
412 let log_k: usize = limbs * log_base2k - 5;
413 let mut a: VecZnx = VecZnx::new(n, log_base2k, limbs);
414 let mut have: Vec<i64> = vec![i64::default(); n];
415 have.iter_mut()
416 .enumerate()
417 .for_each(|(i, x)| *x = (i as i64) - (n as i64) / 2);
418 a.from_i64(&have, 10, log_k);
419 let mut want = vec![i64::default(); n];
420 a.to_i64(&mut want, log_k);
421 izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
422 }
423
424 #[test]
425 fn test_set_get_i64_hi_norm() {
426 let n: usize = 8;
427 let log_base2k: usize = 17;
428 let limbs: usize = 5;
429 let log_k: usize = limbs * log_base2k - 5;
430 let mut a: VecZnx = VecZnx::new(n, log_base2k, limbs);
431 let mut have: Vec<i64> = vec![i64::default(); n];
432 let mut source = Source::new([1; 32]);
433 have.iter_mut().for_each(|x| {
434 *x = source
435 .next_u64n(u64::MAX, u64::MAX)
436 .wrapping_sub(u64::MAX / 2 + 1) as i64;
437 });
438 a.from_i64(&have, 63, log_k);
439 //(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
440 let mut want = vec![i64::default(); n];
441 //(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
442 a.to_i64(&mut want, log_k);
443 izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
444 }
445 #[test]
446 fn test_normalize() {
447 let n: usize = 8;
448 let log_base2k: usize = 17;
449 let limbs: usize = 5;
450 let log_k: usize = limbs * log_base2k - 5;
451 let mut a: VecZnx = VecZnx::new(n, log_base2k, limbs);
452 let mut have: Vec<i64> = vec![i64::default(); n];
453 let mut source = Source::new([1; 32]);
454 have.iter_mut().for_each(|x| {
455 *x = source
456 .next_u64n(u64::MAX, u64::MAX)
457 .wrapping_sub(u64::MAX / 2 + 1) as i64;
458 });
459 a.from_i64(&have, 63, log_k);
460 let mut carry: Vec<u8> = vec![u8::default(); n * 8];
461 a.normalize(&mut carry);
462
463 let base_half = 1 << (log_base2k - 1);
464 a.data
465 .iter()
466 .for_each(|x| assert!(x.abs() <= base_half, "|x|={} > 2^(k-1)={}", x, base_half));
467 let mut want = vec![i64::default(); n];
468 a.to_i64(&mut want, log_k);
469 izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
470 }
471}