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