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