various API uniformization

base2k/src/vec_znx.rs

@@ -0,0 +1,389 @@
+use crate::bindings::{
+    znx_automorphism_i64, znx_automorphism_inplace_i64, znx_normalize, znx_zero_i64_ref,
+};
+use crate::cast_mut_u8_to_mut_i64_slice;
+use crate::module::Module;
+use itertools::izip;
+use rand_distr::{Distribution, Normal};
+use sampling::source::Source;
+use std::cmp::min;
+
+impl Module {
+    pub fn new_vec_znx(&self, log_base2k: usize, log_q: usize) -> VecZnx {
+        VecZnx::new(self.n(), log_base2k, log_q)
+    }
+}
+
+#[derive(Clone)]
+pub struct VecZnx {
+    pub n: usize,
+    pub log_base2k: usize,
+    pub log_q: usize,
+    pub data: Vec<i64>,
+}
+
+impl VecZnx {
+    pub fn new(n: usize, log_base2k: usize, log_q: usize) -> Self {
+        Self {
+            n: n,
+            log_base2k: log_base2k,
+            log_q: log_q,
+            data: vec![i64::default(); Self::buffer_size(n, log_base2k, log_q)],
+        }
+    }
+
+    pub fn buffer_size(n: usize, log_base2k: usize, log_q: usize) -> usize {
+        n * ((log_q + log_base2k - 1) / log_base2k)
+    }
+
+    pub fn from_buffer(&mut self, n: usize, log_base2k: usize, log_q: usize, buf: &[i64]) {
+        let size = Self::buffer_size(n, log_base2k, log_q);
+        assert!(
+            buf.len() >= size,
+            "invalid buffer: buf.len()={} < self.buffer_size(n={}, k={}, log_q={})={}",
+            buf.len(),
+            n,
+            log_base2k,
+            log_q,
+            size
+        );
+        self.n = n;
+        self.log_base2k = log_base2k;
+        self.log_q = log_q;
+        self.data = Vec::from(&buf[..size])
+    }
+
+    pub fn log_n(&self) -> u64 {
+        (u64::BITS - (self.n - 1).leading_zeros()) as _
+    }
+
+    pub fn n(&self) -> usize {
+        self.n
+    }
+
+    pub fn log_q(&self) -> usize {
+        self.log_q
+    }
+
+    pub fn limbs(&self) -> usize {
+        self.data.len() / self.n
+    }
+
+    pub fn copy_from(&mut self, a: &VecZnx) {
+        let size = min(self.data.len(), a.data.len());
+        self.data[..size].copy_from_slice(&a.data[..size])
+    }
+
+    pub fn as_ptr(&self) -> *const i64 {
+        self.data.as_ptr()
+    }
+
+    pub fn as_mut_ptr(&mut self) -> *mut i64 {
+        self.data.as_mut_ptr()
+    }
+
+    pub fn at(&self, i: usize) -> &[i64] {
+        &self.data[i * self.n..(i + 1) * self.n]
+    }
+
+    pub fn at_ptr(&self, i: usize) -> *const i64 {
+        &self.data[i * self.n] as *const i64
+    }
+
+    pub fn at_mut_ptr(&mut self, i: usize) -> *mut i64 {
+        &mut self.data[i * self.n] as *mut i64
+    }
+
+    pub fn at_mut(&mut self, i: usize) -> &mut [i64] {
+        &mut self.data[i * self.n..(i + 1) * self.n]
+    }
+
+    pub fn set_i64(&mut self, data: &[i64], log_max: usize) {
+        let size: usize = min(data.len(), self.n());
+        let k_rem: usize = self.log_base2k - (self.log_q % self.log_base2k);
+
+        // If 2^{log_base2k} * 2^{k_rem} < 2^{63}-1, then we can simply copy
+        // values on the last limb.
+        // Else we decompose values base2k.
+        if log_max + k_rem < 63 || k_rem == self.log_base2k {
+            self.at_mut(self.limbs() - 1).copy_from_slice(&data[..size]);
+        } else {
+            let mask: i64 = (1 << self.log_base2k) - 1;
+            let limbs = self.limbs();
+            let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
+            (limbs - steps..limbs)
+                .rev()
+                .enumerate()
+                .for_each(|(i, i_rev)| {
+                    let shift: usize = i * self.log_base2k;
+                    izip!(self.at_mut(i_rev)[..size].iter_mut(), data[..size].iter())
+                        .for_each(|(y, x)| *y = (x >> shift) & mask);
+                })
+        }
+
+        // Case where self.prec % self.k != 0.
+        if k_rem != self.log_base2k {
+            let limbs = self.limbs();
+            let steps: usize = min(limbs, (log_max + self.log_base2k - 1) / self.log_base2k);
+            (limbs - steps..limbs).rev().for_each(|i| {
+                self.at_mut(i)[..size].iter_mut().for_each(|x| *x <<= k_rem);
+            })
+        }
+    }
+
+    pub fn normalize(&mut self, carry: &mut [u8]) {
+        assert!(
+            carry.len() >= self.n * 8,
+            "invalid carry: carry.len()={} < self.n()={}",
+            carry.len(),
+            self.n()
+        );
+
+        let carry_i64: &mut [i64] = cast_mut_u8_to_mut_i64_slice(carry);
+
+        unsafe {
+            znx_zero_i64_ref(self.n() as u64, carry_i64.as_mut_ptr());
+            (0..self.limbs()).rev().for_each(|i| {
+                znx_normalize(
+                    self.n as u64,
+                    self.log_base2k as u64,
+                    self.at_mut_ptr(i),
+                    carry_i64.as_mut_ptr(),
+                    self.at_mut_ptr(i),
+                    carry_i64.as_mut_ptr(),
+                )
+            });
+        }
+    }
+
+    pub fn get_i64(&self, data: &mut [i64]) {
+        assert!(
+            data.len() >= self.n,
+            "invalid data: data.len()={} < self.n()={}",
+            data.len(),
+            self.n
+        );
+        data.copy_from_slice(self.at(0));
+        let rem: usize = self.log_base2k - (self.log_q % self.log_base2k);
+        (1..self.limbs()).for_each(|i| {
+            if i == self.limbs() - 1 && rem != self.log_base2k {
+                let k_rem: usize = self.log_base2k - rem;
+                izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
+                    *y = (*y << k_rem) + (x >> rem);
+                });
+            } else {
+                izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
+                    *y = (*y << self.log_base2k) + x;
+                });
+            }
+        })
+    }
+
+    pub fn get_single_i64(&self, i: usize) -> i64 {
+        assert!(i < self.n());
+        let mut res: i64 = self.data[i];
+        let rem: usize = self.log_base2k - (self.log_q % self.log_base2k);
+        (1..self.limbs()).for_each(|i| {
+            let x = self.data[i * self.n];
+            if i == self.limbs() - 1 && rem != self.log_base2k {
+                let k_rem: usize = self.log_base2k - rem;
+                res = (res << k_rem) + (x >> rem);
+            } else {
+                res = (res << self.log_base2k) + x;
+            }
+        });
+        res
+    }
+
+    pub fn automorphism_inplace(&mut self, gal_el: i64) {
+        unsafe {
+            (0..self.limbs()).for_each(|i| {
+                znx_automorphism_inplace_i64(self.n as u64, gal_el, self.at_mut_ptr(i))
+            })
+        }
+    }
+    pub fn automorphism(&mut self, gal_el: i64, a: &mut VecZnx) {
+        unsafe {
+            (0..self.limbs()).for_each(|i| {
+                znx_automorphism_i64(self.n as u64, gal_el, a.at_mut_ptr(i), self.at_ptr(i))
+            })
+        }
+    }
+
+    pub fn fill_uniform(&mut self, source: &mut Source) {
+        let mut base2k: u64 = 1 << self.log_base2k;
+        let mut mask: u64 = base2k - 1;
+        let mut base2k_half: i64 = (base2k >> 1) as i64;
+
+        let size: usize = self.n() * (self.limbs() - 1);
+
+        self.data[..size]
+            .iter_mut()
+            .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
+
+        let log_base2k_rem: usize = self.log_q % self.log_base2k;
+
+        if log_base2k_rem != 0 {
+            base2k = 1 << log_base2k_rem;
+            mask = (base2k - 1) << (self.log_base2k - log_base2k_rem);
+            base2k_half = ((mask >> 1) + 1) as i64;
+        }
+
+        self.data[size..]
+            .iter_mut()
+            .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
+    }
+
+    pub fn add_dist_f64<T: Distribution<f64>>(&mut self, source: &mut Source, dist: T, bound: f64) {
+        let log_base2k_rem: usize = self.log_q % self.log_base2k;
+
+        if log_base2k_rem != 0 {
+            self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
+                let mut dist_f64: f64 = dist.sample(source);
+                while dist_f64.abs() > bound {
+                    dist_f64 = dist.sample(source)
+                }
+                *a += (dist_f64.round() as i64) << log_base2k_rem
+            });
+        } else {
+            self.at_mut(self.limbs() - 1).iter_mut().for_each(|a| {
+                let mut dist_f64: f64 = dist.sample(source);
+                while dist_f64.abs() > bound {
+                    dist_f64 = dist.sample(source)
+                }
+                *a += dist_f64.round() as i64
+            });
+        }
+    }
+
+    pub fn add_normal(&mut self, source: &mut Source, sigma: f64, bound: f64) {
+        self.add_dist_f64(source, Normal::new(0.0, sigma).unwrap(), bound);
+    }
+
+    pub fn trunc_pow2(&mut self, k: usize) {
+        if k == 0 {
+            return;
+        }
+
+        assert!(
+            k <= self.log_q,
+            "invalid argument k: k={} > self.prec()={}",
+            k,
+            self.log_q()
+        );
+
+        self.log_q -= k;
+        self.data
+            .truncate((self.limbs() - k / self.log_base2k) * self.n());
+
+        let k_rem: usize = k % self.log_base2k;
+
+        if k_rem != 0 {
+            let mask: i64 = ((1 << (self.log_base2k - k_rem - 1)) - 1) << k_rem;
+            self.at_mut(self.limbs() - 1)
+                .iter_mut()
+                .for_each(|x: &mut i64| *x &= mask)
+        }
+    }
+
+    pub fn rsh(&mut self, k: usize, carry: &mut [u8]) {
+        let limbs: usize = self.limbs();
+        let limbs_steps: usize = k / self.log_base2k;
+
+        self.data.rotate_right(self.n * limbs_steps);
+        unsafe {
+            znx_zero_i64_ref((self.n * limbs_steps) as u64, self.data.as_mut_ptr());
+        }
+
+        let k_rem = k % self.log_base2k;
+
+        if k_rem != 0 {
+            let carry_i64: &mut [i64] = cast_mut_u8_to_mut_i64_slice(carry);
+
+            unsafe {
+                znx_zero_i64_ref(self.n() as u64, carry_i64.as_mut_ptr());
+            }
+
+            let mask: i64 = (1 << k_rem) - 1;
+            let log_base2k: usize = self.log_base2k;
+            println!("mask: {} log_base2k: {}", mask, log_base2k);
+
+            (limbs_steps..limbs).for_each(|i| {
+                izip!(carry_i64.iter_mut(), self.at_mut(i).iter_mut()).for_each(|(ci, xi)| {
+                    *xi += *ci << log_base2k;
+                    *ci = *xi & mask;
+                    *xi /= 1 << k_rem;
+                });
+            })
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::VecZnx;
+    use itertools::izip;
+    use sampling::source::Source;
+
+    #[test]
+    fn test_set_get_i64_lo_norm() {
+        let n: usize = 32;
+        let k: usize = 19;
+        let prec: usize = 128;
+        let mut a: VecZnx = VecZnx::new(n, k, prec);
+        let mut have: Vec<i64> = vec![i64::default(); n];
+        have.iter_mut()
+            .enumerate()
+            .for_each(|(i, x)| *x = (i as i64) - (n as i64) / 2);
+        a.set_i64(&have, 10);
+        let mut want = vec![i64::default(); n];
+        a.get_i64(&mut want);
+        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
+    }
+
+    #[test]
+    fn test_set_get_i64_hi_norm() {
+        let n: usize = 8;
+        let k: usize = 17;
+        let prec: usize = 84;
+        let mut a: VecZnx = VecZnx::new(n, k, prec);
+        let mut have: Vec<i64> = vec![i64::default(); n];
+        let mut source = Source::new([1; 32]);
+        have.iter_mut().for_each(|x| {
+            *x = source
+                .next_u64n(u64::MAX, u64::MAX)
+                .wrapping_sub(u64::MAX / 2 + 1) as i64;
+        });
+        a.set_i64(&have, 63);
+        //(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
+        let mut want = vec![i64::default(); n];
+        //(0..a.limbs()).for_each(|i| println!("i:{} -> {:?}", i, a.at(i)));
+        a.get_i64(&mut want);
+        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
+    }
+    #[test]
+    fn test_normalize() {
+        let n: usize = 8;
+        let k: usize = 17;
+        let prec: usize = 84;
+        let mut a: VecZnx = VecZnx::new(n, k, prec);
+        let mut have: Vec<i64> = vec![i64::default(); n];
+        let mut source = Source::new([1; 32]);
+        have.iter_mut().for_each(|x| {
+            *x = source
+                .next_u64n(u64::MAX, u64::MAX)
+                .wrapping_sub(u64::MAX / 2 + 1) as i64;
+        });
+        a.set_i64(&have, 63);
+        let mut carry: Vec<u8> = vec![u8::default(); n * 8];
+        a.normalize(&mut carry);
+
+        let base_half = 1 << (k - 1);
+        a.data
+            .iter()
+            .for_each(|x| assert!(x.abs() <= base_half, "|x|={} > 2^(k-1)={}", x, base_half));
+        let mut want = vec![i64::default(); n];
+        a.get_i64(&mut want);
+        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
+    }
+}