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