1use backend::{FFT64, Module, VecZnx, VecZnxAlloc, VecZnxOps, ZnxInfos, ZnxViewMut, alloc_aligned};
  2
  3pub struct LookUpTable {
  4    pub(crate) data: Vec<VecZnx<Vec<u8>>>,
  5    pub(crate) basek: usize,
  6    pub(crate) k: usize,
  7}
  8
  9impl LookUpTable {
 10    pub fn alloc(module: &Module<FFT64>, basek: usize, k: usize, extension_factor: usize) -> Self {
 11        #[cfg(debug_assertions)]
 12        {
 13            assert!(
 14                extension_factor & (extension_factor - 1) == 0,
 15                "extension_factor must be a power of two but is: {}",
 16                extension_factor
 17            );
 18        }
 19        let size: usize = k.div_ceil(basek);
 20        let mut data: Vec<VecZnx<Vec<u8>>> = Vec::with_capacity(extension_factor);
 21        (0..extension_factor).for_each(|_| {
 22            data.push(module.new_vec_znx(1, size));
 23        });
 24        Self { data, basek, k }
 25    }
 26
 27    pub fn log_extension_factor(&self) -> usize {
 28        (usize::BITS - (self.extension_factor() - 1).leading_zeros()) as _
 29    }
 30
 31    pub fn extension_factor(&self) -> usize {
 32        self.data.len()
 33    }
 34
 35    pub fn domain_size(&self) -> usize {
 36        self.data.len() * self.data[0].n()
 37    }
 38
 39    pub fn set(&mut self, module: &Module<FFT64>, f: &Vec<i64>, k: usize) {
 40        assert!(f.len() <= module.n());
 41
 42        let basek: usize = self.basek;
 43
 44        // Get the number minimum limb to store the message modulus
 45        let limbs: usize = k.div_ceil(1 << basek);
 46
 47        #[cfg(debug_assertions)]
 48        {
 49            assert!(limbs <= self.data[0].size());
 50        }
 51
 52        // Scaling factor
 53        let scale: i64 = 1 << (k % basek) as i64;
 54
 55        // #elements in lookup table
 56        let f_len: usize = f.len();
 57
 58        // If LUT size > module.n()
 59        let domain_size: usize = self.domain_size();
 60
 61        let size: usize = self.k.div_ceil(self.basek);
 62
 63        // Equivalent to AUTO([f(0), -f(n-1), -f(n-2), ..., -f(1)], -1)
 64        let mut lut_full: VecZnx<Vec<u8>> = VecZnx::new::<i64>(domain_size, 1, size);
 65
 66        let lut_at: &mut [i64] = lut_full.at_mut(0, limbs - 1);
 67
 68        f.iter().enumerate().for_each(|(i, fi)| {
 69            let start: usize = (i * domain_size).div_round(f_len);
 70            let end: usize = ((i + 1) * domain_size).div_round(f_len);
 71            lut_at[start..end].fill(fi * scale);
 72        });
 73
 74        // Rotates half the step to the left
 75        let half_step: usize = domain_size.div_round(f_len << 1);
 76
 77        lut_full.rotate(-(half_step as i64));
 78
 79        let mut tmp_bytes: Vec<u8> = alloc_aligned(lut_full.n() * size_of::<i64>());
 80        lut_full.normalize(self.basek, 0, &mut tmp_bytes);
 81
 82        if self.extension_factor() > 1 {
 83            (0..self.extension_factor()).for_each(|i| {
 84                module.switch_degree(&mut self.data[i], 0, &lut_full, 0);
 85                if i < self.extension_factor() {
 86                    lut_full.rotate(-1);
 87                }
 88            });
 89        } else {
 90            module.vec_znx_copy(&mut self.data[0], 0, &lut_full, 0);
 91        }
 92    }
 93
 94    #[allow(dead_code)]
 95    pub(crate) fn rotate(&mut self, k: i64) {
 96        let extension_factor: usize = self.extension_factor();
 97        let two_n: usize = 2 * self.data[0].n();
 98        let two_n_ext: usize = two_n * extension_factor;
 99
100        let k_pos: usize = ((k + two_n_ext as i64) % two_n_ext as i64) as usize;
101
102        let k_hi: usize = k_pos / extension_factor;
103        let k_lo: usize = k_pos % extension_factor;
104
105        (0..extension_factor - k_lo).for_each(|i| {
106            self.data[i].rotate(k_hi as i64);
107        });
108
109        (extension_factor - k_lo..extension_factor).for_each(|i| {
110            self.data[i].rotate(k_hi as i64 + 1);
111        });
112
113        self.data.rotate_right(k_lo as usize);
114    }
115}
116
117pub(crate) trait DivRound {
118    fn div_round(self, rhs: Self) -> Self;
119}
120
121impl DivRound for usize {
122    #[inline]
123    fn div_round(self, rhs: Self) -> Self {
124        (self + rhs / 2) / rhs
125    }
126}