1use crate::{znx_normalize, znx_zero_i64_ref};
  2use itertools::izip;
  3use std::cmp::min;
  4
  5pub struct Poly {
  6    pub n: usize,
  7    pub k: usize,
  8    pub prec: usize,
  9    pub data: Vec<i64>,
 10}
 11
 12impl Poly {
 13    pub fn new(n: usize, k: usize, prec: usize) -> Self {
 14        Self {
 15            n: n,
 16            k: k,
 17            prec: prec,
 18            data: vec![i64::default(); Self::buffer_size(n, k, prec)],
 19        }
 20    }
 21
 22    pub fn buffer_size(n: usize, k: usize, prec: usize) -> usize {
 23        n * ((prec + k - 1) / k)
 24    }
 25
 26    pub fn from_buffer(&mut self, n: usize, k: usize, prec: usize, buf: &[i64]) {
 27        let size = Self::buffer_size(n, k, prec);
 28        assert!(
 29            buf.len() >= size,
 30            "invalid buffer: buf.len()={} < self.buffer_size(n={}, k={}, prec={})={}",
 31            buf.len(),
 32            n,
 33            k,
 34            prec,
 35            size
 36        );
 37        self.n = n;
 38        self.k = k;
 39        self.prec = prec;
 40        self.data = Vec::from(&buf[..size])
 41    }
 42
 43    pub fn log_n(&self) -> usize {
 44        (u64::BITS - (self.n - 1).leading_zeros()) as _
 45    }
 46
 47    pub fn n(&self) -> usize {
 48        self.n
 49    }
 50
 51    pub fn limbs(&self) -> usize {
 52        self.data.len() / self.n
 53    }
 54
 55    pub fn at(&self, i: usize) -> &[i64] {
 56        &self.data[i * self.n..(i + 1) * self.n]
 57    }
 58
 59    pub fn at_ptr(&self, i: usize) -> *const i64 {
 60        &self.data[i * self.n] as *const i64
 61    }
 62
 63    pub fn at_mut_ptr(&mut self, i: usize) -> *mut i64 {
 64        &mut self.data[i * self.n] as *mut i64
 65    }
 66
 67    pub fn at_mut(&mut self, i: usize) -> &mut [i64] {
 68        &mut self.data[i * self.n..(i + 1) * self.n]
 69    }
 70
 71    pub fn set_i64(&mut self, data: &[i64], log_max: usize) {
 72        let size: usize = min(data.len(), self.n());
 73        let k_rem: usize = self.k - (self.prec % self.k);
 74
 75        // If 2^{base} * 2^{k_rem} < 2^{63}-1, then we can simply copy
 76        // values on the last limb.
 77        // Else we decompose values base k.
 78        if log_max + k_rem < 63 || k_rem == self.k {
 79            self.at_mut(self.limbs() - 1).copy_from_slice(&data[..size]);
 80        } else {
 81            let mask: i64 = (1 << self.k) - 1;
 82            let limbs = self.limbs();
 83            let steps: usize = min(limbs, (log_max + k_rem + self.k - 1) / self.k);
 84            (limbs - steps..limbs)
 85                .rev()
 86                .enumerate()
 87                .for_each(|(i, i_rev)| {
 88                    let shift: usize = i * self.k;
 89                    izip!(self.at_mut(i_rev)[..size].iter_mut(), data[..size].iter())
 90                        .for_each(|(y, x)| *y = (x >> shift) & mask);
 91                })
 92        }
 93
 94        // Case where self.prec % self.k != 0.
 95        if k_rem != self.k {
 96            let limbs = self.limbs();
 97            let steps: usize = min(limbs, (log_max + k_rem + self.k - 1) / self.k);
 98            (limbs - steps..limbs).rev().for_each(|i| {
 99                self.at_mut(i)[..size].iter_mut().for_each(|x| *x <<= k_rem);
100            })
101        }
102    }
103
104    pub fn normalize(&mut self, carry: &mut [i64]) {
105        assert!(
106            carry.len() >= self.n,
107            "invalid carry: carry.len()={} < self.n()={}",
108            carry.len(),
109            self.n()
110        );
111        unsafe {
112            znx_zero_i64_ref(self.n() as u64, carry.as_mut_ptr());
113            (0..self.limbs()).rev().for_each(|i| {
114                znx_normalize(
115                    self.n as u64,
116                    self.k as u64,
117                    self.at_mut_ptr(i),
118                    carry.as_mut_ptr(),
119                    self.at_mut_ptr(i),
120                    carry.as_mut_ptr(),
121                )
122            });
123        }
124    }
125
126    pub fn get_i64(&self, data: &mut [i64]) {
127        assert!(
128            data.len() >= self.n,
129            "invalid data: data.len()={} < self.n()={}",
130            data.len(),
131            self.n
132        );
133        data.copy_from_slice(self.at(0));
134        let rem: usize = self.k - (self.prec % self.k);
135        (1..self.limbs()).for_each(|i| {
136            if i == self.limbs() - 1 && rem != self.k {
137                let k_rem: usize = self.k - rem;
138                izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
139                    *y = (*y << k_rem) + (x >> rem);
140                });
141            } else {
142                izip!(self.at(i).iter(), data.iter_mut()).for_each(|(x, y)| {
143                    *y = (*y << self.k) + x;
144                });
145            }
146        })
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use crate::poly::Poly;
153    use itertools::izip;
154    use sampling::source::Source;
155
156    #[test]
157    fn test_set_get_i64_lo_norm() {
158        let n: usize = 32;
159        let k: usize = 19;
160        let prec: usize = 128;
161        let mut a: Poly = Poly::new(n, k, prec);
162        let mut have: Vec<i64> = vec![i64::default(); n];
163        have.iter_mut()
164            .enumerate()
165            .for_each(|(i, x)| *x = (i as i64) - (n as i64) / 2);
166        a.set_i64(&have, 10);
167        let mut want = vec![i64::default(); n];
168        a.get_i64(&mut want);
169        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
170    }
171
172    #[test]
173    fn test_set_get_i64_hi_norm() {
174        let n: usize = 1;
175        let k: usize = 19;
176        let prec: usize = 128;
177        let mut a: Poly = Poly::new(n, k, prec);
178        let mut have: Vec<i64> = vec![i64::default(); n];
179        let mut source = Source::new([1; 32]);
180        have.iter_mut().for_each(|x| {
181            *x = source
182                .next_u64n(u64::MAX, u64::MAX)
183                .wrapping_sub(u64::MAX / 2 + 1) as i64;
184        });
185        a.set_i64(&have, 63);
186        let mut want = vec![i64::default(); n];
187        a.get_i64(&mut want);
188        izip!(want, have).for_each(|(a, b)| assert_eq!(a, b));
189    }
190}