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 itertools::izip;
  6use rand_distr::{Distribution, Normal};
  7use sampling::source::Source;
  8use std::cmp::min;
  9
 10pub struct Vector {
 11    pub n: usize,
 12    pub log_base2k: usize,
 13    pub prec: usize,
 14    pub data: Vec<i64>,
 15}
 16
 17impl Vector {
 18    pub fn new(n: usize, log_base2k: usize, prec: usize) -> Self {
 19        Self {
 20            n: n,
 21            log_base2k: log_base2k,
 22            prec: prec,
 23            data: vec![i64::default(); Self::buffer_size(n, log_base2k, prec)],
 24        }
 25    }
 26
 27    pub fn buffer_size(n: usize, log_base2k: usize, prec: usize) -> usize {
 28        n * ((prec + log_base2k - 1) / log_base2k)
 29    }
 30
 31    pub fn from_buffer(&mut self, n: usize, log_base2k: usize, prec: usize, buf: &[i64]) {
 32        let size = Self::buffer_size(n, log_base2k, prec);
 33        assert!(
 34            buf.len() >= size,
 35            "invalid buffer: buf.len()={} < self.buffer_size(n={}, k={}, prec={})={}",
 36            buf.len(),
 37            n,
 38            log_base2k,
 39            prec,
 40            size
 41        );
 42        self.n = n;
 43        self.log_base2k = log_base2k;
 44        self.prec = prec;
 45        self.data = Vec::from(&buf[..size])
 46    }
 47
 48    pub fn log_n(&self) -> u64 {
 49        (u64::BITS - (self.n - 1).leading_zeros()) as _
 50    }
 51
 52    pub fn n(&self) -> usize {
 53        self.n
 54    }
 55
 56    pub fn prec(&self) -> usize {
 57        self.prec
 58    }
 59
 60    pub fn limbs(&self) -> usize {
 61        self.data.len() / self.n
 62    }
 63
 64    pub fn as_ptr(&self) -> *const i64 {
 65        self.data.as_ptr()
 66    }
 67
 68    pub fn as_mut_ptr(&mut self) -> *mut i64 {
 69        self.data.as_mut_ptr()
 70    }
 71
 72    pub fn at(&self, i: usize) -> &[i64] {
 73        &self.data[i * self.n..(i + 1) * self.n]
 74    }
 75
 76    pub fn at_ptr(&self, i: usize) -> *const i64 {
 77        &self.data[i * self.n] as *const i64
 78    }
 79
 80    pub fn at_mut_ptr(&mut self, i: usize) -> *mut i64 {
 81        &mut self.data[i * self.n] as *mut i64
 82    }
 83
 84    pub fn at_mut(&mut self, i: usize) -> &mut [i64] {
 85        &mut self.data[i * self.n..(i + 1) * self.n]
 86    }
 87
 88    pub fn set_i64(&mut self, data: &[i64], log_max: usize) {
 89        let size: usize = min(data.len(), self.n());
 90        let k_rem: usize = self.log_base2k - (self.prec % self.log_base2k);
 91
 92        // If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
 93        // values on the last limb.
 94        // Else we decompose values base2k.
 95        if log_max + k_rem < 63 || k_rem == self.log_base2k {
 96            self.at_mut(self.limbs() - 1).copy_from_slice(&data[..size]);
 97        } else {
 98            let mask: i64 = (1 << self.log_base2k) - 1;
 99            let limbs = self.limbs();
100            let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
101            (limbs - steps..limbs)
102                .rev()
103                .enumerate()
104                .for_each(|(i, i_rev)| {
105                    let shift: usize = i * self.log_base2k;
106                    izip!(self.at_mut(i_rev)[..size].iter_mut(), data[..size].iter())
107                        .for_each(|(y, x)| *y = (x >> shift) & mask);
108                })
109        }
110
111        // Case where self.prec % self.k != 0.
112        if k_rem != self.log_base2k {
113            let limbs = self.limbs();
114            let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
115            (limbs - steps..limbs).rev().for_each(|i| {
116                self.at_mut(i)[..size].iter_mut().for_each(|x| *x <<= k_rem);
117            })
118        }
119    }
120
121    pub fn normalize(&mut self, carry: &mut [u8]) {
122        assert!(
123            carry.len() >= self.n * 8,
124            "invalid carry: carry.len()={} < self.n()={}",
125            carry.len(),
126            self.n()
127        );
128
129        let carry_i64: &mut [i64] = cast_mut_u8_to_mut_i64_slice(carry);
130
131        unsafe {
132            znx_zero_i64_ref(self.n() as u64, carry_i64.as_mut_ptr());
133            (0..self.limbs()).rev().for_each(|i| {
134                znx_normalize(
135                    self.n as u64,
136                    self.log_base2k as u64,
137                    self.at_mut_ptr(i),
138                    carry_i64.as_mut_ptr(),
139                    self.at_mut_ptr(i),
140                    carry_i64.as_mut_ptr(),
141                )
142            });
143        }
144    }
145
146    pub fn get_i64(&self, data: &mut [i64]) {
147        assert!(
148            data.len() >= self.n,
149            "invalid data: data.len()={} < self.n()={}",
150            data.len(),
151            self.n
152        );
153        data.copy_from_slice(self.at(0));
154        let rem: usize = self.log_base2k - (self.prec % self.log_base2k);
155        (1..self.limbs()).for_each(|i| {
156            if i == self.limbs() - 1 && rem != self.log_base2k {
157                let k_rem: usize = self.log_base2k - rem;
158                izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
159                    *y = (*y << k_rem) + (x >> rem);
160                });
161            } else {
162                izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
163                    *y = (*y << self.log_base2k) + x;
164                });
165            }
166        })
167    }
168
169    pub fn automorphism_inplace(&mut self, gal_el: i64) {
170        unsafe {
171            (0..self.limbs()).for_each(|i| {
172                znx_automorphism_inplace_i64(self.n as u64, gal_el, self.at_mut_ptr(i))
173            })
174        }
175    }
176    pub fn automorphism(&mut self, gal_el: i64, a: &mut Vector) {
177        unsafe {
178            (0..self.limbs()).for_each(|i| {
179                znx_automorphism_i64(self.n as u64, gal_el, a.at_mut_ptr(i), self.at_ptr(i))
180            })
181        }
182    }
183
184    pub fn fill_uniform(&mut self, source: &mut Source) {
185        let mut base2k: u64 = 1 << self.log_base2k;
186        let mut mask: u64 = base2k - 1;
187        let mut base2k_half: i64 = (base2k >> 1) as i64;
188
189        let size: usize = self.n() * (self.limbs() - 1);
190
191        self.data[..size]
192            .iter_mut()
193            .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
194
195        let log_base2k_rem: usize = self.prec % self.log_base2k;
196
197        if log_base2k_rem != 0 {
198            base2k = 1 << log_base2k_rem;
199            mask = (base2k - 1) << (self.log_base2k - log_base2k_rem);
200            base2k_half = ((mask >> 1) + 1) as i64;
201        }
202
203        self.data[size..]
204            .iter_mut()
205            .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
206    }
207
208    pub fn add_dist_f64<T: Distribution<f64>>(&mut self, source: &mut Source, dist: T, bound: f64) {
209        let log_base2k_rem: usize = self.prec % self.log_base2k;
210
211        if log_base2k_rem != 0 {
212            self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
213                let mut dist_f64: f64 = dist.sample(source);
214                while dist_f64.abs() > bound {
215                    dist_f64 = dist.sample(source)
216                }
217                *a += (dist_f64.round() as i64) << log_base2k_rem
218            });
219        } else {
220            self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
221                let mut dist_f64: f64 = dist.sample(source);
222                while dist_f64.abs() > bound {
223                    dist_f64 = dist.sample(source)
224                }
225                *a += dist_f64.round() as i64
226            });
227        }
228    }
229
230    pub fn add_normal(&mut self, source: &mut Source, sigma: f64, bound: f64) {
231        self.add_dist_f64(source, Normal::new(0.0, sigma).unwrap(), bound);
232    }
233
234    pub fn trunc_pow2(&mut self, k: usize) {
235        if k == 0 {
236            return;
237        }
238
239        assert!(
240            k <= self.prec,
241            "invalid argument k: k={} > self.prec()={}",
242            k,
243            self.prec()
244        );
245
246        self.prec -= k;
247        self.data
248            .truncate((self.limbs() - k / self.log_base2k) * self.n());
249
250        let k_rem: usize = k % self.log_base2k;
251
252        if k_rem != 0 {
253            let mask: i64 = ((1 << (self.log_base2k - k_rem - 1)) - 1) << k_rem;
254            self.at_mut(self.limbs() - 1)
255                .iter_mut()
256                .for_each(|x: &mut i64| *x &= mask)
257        }
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use crate::vector::Vector;
264    use itertools::izip;
265    use sampling::source::Source;
266
267    #[test]
268    fn test_set_get_i64_lo_norm() {
269        let n: usize = 32;
270        let k: usize = 19;
271        let prec: usize = 128;
272        let mut a: Vector = Vector::new(n, k, prec);
273        let mut have: Vec<i64> = vec![i64::default(); n];
274        have.iter_mut()
275            .enumerate()
276            .for_each(|(i, x)| *x = (i as i64) - (n as i64) / 2);
277        a.set_i64(&have, 10);
278        let mut want = vec![i64::default(); n];
279        a.get_i64(&mut want);
280        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
281    }
282
283    #[test]
284    fn test_set_get_i64_hi_norm() {
285        let n: usize = 8;
286        let k: usize = 17;
287        let prec: usize = 84;
288        let mut a: Vector = Vector::new(n, k, prec);
289        let mut have: Vec<i64> = vec![i64::default(); n];
290        let mut source = Source::new([1; 32]);
291        have.iter_mut().for_each(|x| {
292            *x = source
293                .next_u64n(u64::MAX, u64::MAX)
294                .wrapping_sub(u64::MAX / 2 + 1) as i64;
295        });
296        a.set_i64(&have, 63);
297        //(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
298        let mut want = vec![i64::default(); n];
299        //(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
300        a.get_i64(&mut want);
301        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
302    }
303    #[test]
304    fn test_normalize() {
305        let n: usize = 8;
306        let k: usize = 17;
307        let prec: usize = 84;
308        let mut a: Vector = Vector::new(n, k, prec);
309        let mut have: Vec<i64> = vec![i64::default(); n];
310        let mut source = Source::new([1; 32]);
311        have.iter_mut().for_each(|x| {
312            *x = source
313                .next_u64n(u64::MAX, u64::MAX)
314                .wrapping_sub(u64::MAX / 2 + 1) as i64;
315        });
316        a.set_i64(&have, 63);
317        let mut carry: Vec<u8> = vec![u8::default(); n * 8];
318        a.normalize(&mut carry);
319
320        let base_half = 1 << (k - 1);
321        a.data
322            .iter()
323            .for_each(|x| assert!(x.abs() <= base_half, "|x|={} > 2^(k-1)={}", x, base_half));
324        let mut want = vec![i64::default(); n];
325        a.get_i64(&mut want);
326        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
327    }
328}