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 from_i64(&mut self, data: &[i64], log_max: usize) {
102        let size: usize = min(data.len(), self.n());
103        let k_rem: usize = self.log_base2k - (self.log_q % self.log_base2k);
104
105        let limbs: usize = self.limbs();
106
107        // If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
108        // values on the last limb.
109        // Else we decompose values base2k.
110        if log_max + k_rem < 63 || k_rem == self.log_base2k {
111            (0..limbs - 1).for_each(|i| unsafe {
112                znx_zero_i64_ref(size as u64, self.at_mut(i).as_mut_ptr());
113            });
114
115            self.at_mut(self.limbs() - 1)[..size].copy_from_slice(&data[..size]);
116        } else {
117            let mask: i64 = (1 << self.log_base2k) - 1;
118            let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
119
120            (0..steps).for_each(|i| unsafe {
121                znx_zero_i64_ref(size as u64, self.at_mut(i).as_mut_ptr());
122            });
123
124            (limbs - steps..limbs)
125                .rev()
126                .enumerate()
127                .for_each(|(i, i_rev)| {
128                    let shift: usize = i * self.log_base2k;
129                    izip!(self.at_mut(i_rev)[..size].iter_mut(), data[..size].iter())
130                        .for_each(|(y, x)| *y = (x >> shift) & mask);
131                })
132        }
133
134        // Case where self.prec % self.k != 0.
135        if k_rem != self.log_base2k {
136            let limbs = self.limbs();
137            let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
138            (limbs - steps..limbs).rev().for_each(|i| {
139                self.at_mut(i)[..size].iter_mut().for_each(|x| *x <<= k_rem);
140            })
141        }
142    }
143
144    pub fn from_i64_single(&mut self, i: usize, value: i64, log_max: usize) {
145        assert!(i < self.n());
146        let k_rem: usize = self.log_base2k - (self.log_q % self.log_base2k);
147        let limbs = self.limbs();
148
149        // If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
150        // values on the last limb.
151        // Else we decompose values base2k.
152        if log_max + k_rem < 63 || k_rem == self.log_base2k {
153            (0..limbs - 1).for_each(|j| self.at_mut(j)[i] = 0);
154
155            self.at_mut(self.limbs() - 1)[i] = value;
156        } else {
157            let mask: i64 = (1 << self.log_base2k) - 1;
158            let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
159
160            (0..limbs - steps).for_each(|j| self.at_mut(j)[i] = 0);
161
162            (limbs - steps..limbs)
163                .rev()
164                .enumerate()
165                .for_each(|(j, j_rev)| {
166                    self.at_mut(j_rev)[i] = (value >> (j * self.log_base2k)) & mask;
167                })
168        }
169
170        // Case where self.prec % self.k != 0.
171        if k_rem != self.log_base2k {
172            let limbs = self.limbs();
173            let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
174            (limbs - steps..limbs).rev().for_each(|j| {
175                self.at_mut(j)[i] <<= k_rem;
176            })
177        }
178    }
179
180    pub fn normalize(&mut self, carry: &mut [u8]) {
181        assert!(
182            carry.len() >= self.n * 8,
183            "invalid carry: carry.len()={} < self.n()={}",
184            carry.len(),
185            self.n()
186        );
187
188        let carry_i64: &mut [i64] = cast_mut_u8_to_mut_i64_slice(carry);
189
190        unsafe {
191            znx_zero_i64_ref(self.n() as u64, carry_i64.as_mut_ptr());
192            (0..self.limbs()).rev().for_each(|i| {
193                znx_normalize(
194                    self.n as u64,
195                    self.log_base2k as u64,
196                    self.at_mut_ptr(i),
197                    carry_i64.as_mut_ptr(),
198                    self.at_mut_ptr(i),
199                    carry_i64.as_mut_ptr(),
200                )
201            });
202        }
203    }
204
205    pub fn to_i64(&self, data: &mut [i64]) {
206        assert!(
207            data.len() >= self.n,
208            "invalid data: data.len()={} < self.n()={}",
209            data.len(),
210            self.n
211        );
212        data.copy_from_slice(self.at(0));
213        let rem: usize = self.log_base2k - (self.log_q % self.log_base2k);
214        (1..self.limbs()).for_each(|i| {
215            if i == self.limbs() - 1 && rem != self.log_base2k {
216                let k_rem: usize = self.log_base2k - rem;
217                izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
218                    *y = (*y << k_rem) + (x >> rem);
219                });
220            } else {
221                izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
222                    *y = (*y << self.log_base2k) + x;
223                });
224            }
225        })
226    }
227
228    pub fn to_i64_single(&self, i: usize) -> i64 {
229        assert!(i < self.n());
230        let mut res: i64 = self.data[i];
231        let rem: usize = self.log_base2k - (self.log_q % self.log_base2k);
232        (1..self.limbs()).for_each(|i| {
233            let x = self.data[i * self.n];
234            if i == self.limbs() - 1 && rem != self.log_base2k {
235                let k_rem: usize = self.log_base2k - rem;
236                res = (res << k_rem) + (x >> rem);
237            } else {
238                res = (res << self.log_base2k) + x;
239            }
240        });
241        res
242    }
243
244    pub fn automorphism_inplace(&mut self, gal_el: i64) {
245        unsafe {
246            (0..self.limbs()).for_each(|i| {
247                znx_automorphism_inplace_i64(self.n as u64, gal_el, self.at_mut_ptr(i))
248            })
249        }
250    }
251    pub fn automorphism(&mut self, gal_el: i64, a: &mut VecZnx) {
252        unsafe {
253            (0..self.limbs()).for_each(|i| {
254                znx_automorphism_i64(self.n as u64, gal_el, a.at_mut_ptr(i), self.at_ptr(i))
255            })
256        }
257    }
258
259    pub fn fill_uniform(&mut self, source: &mut Source) {
260        let mut base2k: u64 = 1 << self.log_base2k;
261        let mut mask: u64 = base2k - 1;
262        let mut base2k_half: i64 = (base2k >> 1) as i64;
263
264        let size: usize = self.n() * (self.limbs() - 1);
265
266        self.data[..size]
267            .iter_mut()
268            .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
269
270        let log_base2k_rem: usize = self.log_q % self.log_base2k;
271
272        if log_base2k_rem != 0 {
273            base2k = 1 << log_base2k_rem;
274            mask = (base2k - 1) << (self.log_base2k - log_base2k_rem);
275            base2k_half = ((mask >> 1) + 1) as i64;
276        }
277
278        self.data[size..]
279            .iter_mut()
280            .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
281    }
282
283    pub fn add_dist_f64<T: Distribution<f64>>(&mut self, source: &mut Source, dist: T, bound: f64) {
284        let log_base2k_rem: usize = self.log_q % self.log_base2k;
285
286        if log_base2k_rem != 0 {
287            self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
288                let mut dist_f64: f64 = dist.sample(source);
289                while dist_f64.abs() > bound {
290                    dist_f64 = dist.sample(source)
291                }
292                *a += (dist_f64.round() as i64) << log_base2k_rem
293            });
294        } else {
295            self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
296                let mut dist_f64: f64 = dist.sample(source);
297                while dist_f64.abs() > bound {
298                    dist_f64 = dist.sample(source)
299                }
300                *a += dist_f64.round() as i64
301            });
302        }
303    }
304
305    pub fn add_normal(&mut self, source: &mut Source, sigma: f64, bound: f64) {
306        self.add_dist_f64(source, Normal::new(0.0, sigma).unwrap(), bound);
307    }
308
309    pub fn trunc_pow2(&mut self, k: usize) {
310        if k == 0 {
311            return;
312        }
313
314        assert!(
315            k <= self.log_q,
316            "invalid argument k: k={} > self.prec()={}",
317            k,
318            self.log_q()
319        );
320
321        self.log_q -= k;
322        self.data
323            .truncate((self.limbs() - k / self.log_base2k) * self.n());
324
325        let k_rem: usize = k % self.log_base2k;
326
327        if k_rem != 0 {
328            let mask: i64 = ((1 << (self.log_base2k - k_rem - 1)) - 1) << k_rem;
329            self.at_mut(self.limbs() - 1)
330                .iter_mut()
331                .for_each(|x: &mut i64| *x &= mask)
332        }
333    }
334
335    pub fn rsh(&mut self, k: usize, carry: &mut [u8]) {
336        let limbs: usize = self.limbs();
337        let limbs_steps: usize = k / self.log_base2k;
338
339        self.data.rotate_right(self.n * limbs_steps);
340        unsafe {
341            znx_zero_i64_ref((self.n * limbs_steps) as u64, self.data.as_mut_ptr());
342        }
343
344        let k_rem = k % self.log_base2k;
345
346        if k_rem != 0 {
347            let carry_i64: &mut [i64] = cast_mut_u8_to_mut_i64_slice(carry);
348
349            unsafe {
350                znx_zero_i64_ref(self.n() as u64, carry_i64.as_mut_ptr());
351            }
352
353            let mask: i64 = (1 << k_rem) - 1;
354            let log_base2k: usize = self.log_base2k;
355
356            (limbs_steps..limbs).for_each(|i| {
357                izip!(carry_i64.iter_mut(), self.at_mut(i).iter_mut()).for_each(|(ci, xi)| {
358                    *xi += *ci << log_base2k;
359                    *ci = *xi & mask;
360                    *xi /= 1 << k_rem;
361                });
362            })
363        }
364    }
365}
366
367#[cfg(test)]
368mod tests {
369    use crate::VecZnx;
370    use itertools::izip;
371    use sampling::source::Source;
372
373    #[test]
374    fn test_set_get_i64_lo_norm() {
375        let n: usize = 32;
376        let k: usize = 19;
377        let prec: usize = 128;
378        let mut a: VecZnx = VecZnx::new(n, k, prec);
379        let mut have: Vec<i64> = vec![i64::default(); n];
380        have.iter_mut()
381            .enumerate()
382            .for_each(|(i, x)| *x = (i as i64) - (n as i64) / 2);
383        a.from_i64(&have, 10);
384        let mut want = vec![i64::default(); n];
385        a.to_i64(&mut want);
386        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
387    }
388
389    #[test]
390    fn test_set_get_i64_hi_norm() {
391        let n: usize = 8;
392        let k: usize = 17;
393        let prec: usize = 84;
394        let mut a: VecZnx = VecZnx::new(n, k, prec);
395        let mut have: Vec<i64> = vec![i64::default(); n];
396        let mut source = Source::new([1; 32]);
397        have.iter_mut().for_each(|x| {
398            *x = source
399                .next_u64n(u64::MAX, u64::MAX)
400                .wrapping_sub(u64::MAX / 2 + 1) as i64;
401        });
402        a.from_i64(&have, 63);
403        //(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
404        let mut want = vec![i64::default(); n];
405        //(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
406        a.to_i64(&mut want);
407        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
408    }
409    #[test]
410    fn test_normalize() {
411        let n: usize = 8;
412        let k: usize = 17;
413        let prec: usize = 84;
414        let mut a: VecZnx = VecZnx::new(n, k, prec);
415        let mut have: Vec<i64> = vec![i64::default(); n];
416        let mut source = Source::new([1; 32]);
417        have.iter_mut().for_each(|x| {
418            *x = source
419                .next_u64n(u64::MAX, u64::MAX)
420                .wrapping_sub(u64::MAX / 2 + 1) as i64;
421        });
422        a.from_i64(&have, 63);
423        let mut carry: Vec<u8> = vec![u8::default(); n * 8];
424        a.normalize(&mut carry);
425
426        let base_half = 1 << (k - 1);
427        a.data
428            .iter()
429            .for_each(|x| assert!(x.abs() <= base_half, "|x|={} > 2^(k-1)={}", x, base_half));
430        let mut want = vec![i64::default(); n];
431        a.to_i64(&mut want);
432        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
433    }
434}