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