1use backend::{FFT64, Module, ScalarZnx, VecZnx, VecZnxAlloc, VecZnxOps, ZnxInfos, ZnxView, 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        let size: usize = k.div_ceil(basek);
 12        let mut data: Vec<VecZnx<Vec<u8>>> = Vec::with_capacity(extension_factor);
 13        (0..extension_factor).for_each(|_| {
 14            data.push(module.new_vec_znx(1, size));
 15        });
 16        Self { data, basek, k }
 17    }
 18
 19    pub fn extension_factor(&self) -> usize {
 20        self.data.len()
 21    }
 22
 23    pub fn domain_size(&self) -> usize {
 24        self.data.len() * self.data[0].n()
 25    }
 26
 27    pub fn set(&mut self, module: &Module<FFT64>, f: fn(i64) -> i64, message_modulus: usize) {
 28        let basek: usize = self.basek;
 29
 30        // Get the number minimum limb to store the message modulus
 31        let limbs: usize = message_modulus.div_ceil(1 << basek);
 32
 33        // Scaling factor
 34        let scale: i64 = (1 << (basek * limbs - 1)).div_round(message_modulus) as i64;
 35
 36        // Updates function
 37        let f_scaled = |x: i64| (f(x) % message_modulus as i64) * scale;
 38
 39        // If LUT size > module.n()
 40        let domain_size: usize = self.domain_size();
 41
 42        let size: usize = self.k.div_ceil(self.basek);
 43
 44        // Equivalent to AUTO([f(0), -f(n-1), -f(n-2), ..., -f(1)], -1)
 45        let mut lut_full: VecZnx<Vec<u8>> = VecZnx::new::<i64>(domain_size, 1, size);
 46        {
 47            let lut_at: &mut [i64] = lut_full.at_mut(0, limbs - 1);
 48
 49            let start: usize = 0;
 50            let end: usize = (domain_size).div_round(message_modulus);
 51
 52            let y: i64 = f_scaled(0);
 53            (start..end).for_each(|i| {
 54                lut_at[i] = y;
 55            });
 56
 57            (1..message_modulus).for_each(|x| {
 58                let start: usize = (x * domain_size).div_round(message_modulus);
 59                let end: usize = ((x + 1) * domain_size).div_round(message_modulus);
 60                let y: i64 = f_scaled(x as i64);
 61                (start..end).for_each(|i| {
 62                    lut_at[i] = y;
 63                })
 64            });
 65        }
 66
 67        // Rotates half the step to the left
 68        let half_step: usize = domain_size.div_round(message_modulus << 1);
 69
 70        lut_full.rotate(-(half_step as i64));
 71
 72        let mut tmp_bytes: Vec<u8> = alloc_aligned(lut_full.n() * size_of::<i64>());
 73        lut_full.normalize(self.basek, 0, &mut tmp_bytes);
 74
 75        if self.extension_factor() > 1 {
 76            (0..self.extension_factor()).for_each(|i| {
 77                module.switch_degree(&mut self.data[i], 0, &lut_full, 0);
 78                if i < self.extension_factor() {
 79                    lut_full.rotate(-1);
 80                }
 81            });
 82        } else {
 83            module.vec_znx_copy(&mut self.data[0], 0, &lut_full, 0);
 84        }
 85    }
 86
 87    pub fn set_raw<D>(&mut self, module: &Module<FFT64>, lut: &ScalarZnx<D>)
 88    where
 89        D: AsRef<[u8]>,
 90    {
 91        let domain_size: usize = self.domain_size();
 92
 93        let size: usize = self.k.div_ceil(self.basek);
 94
 95        let mut lut_full: VecZnx<Vec<u8>> = VecZnx::new::<i64>(domain_size, 1, size);
 96
 97        lut_full.at_mut(0, 0).copy_from_slice(lut.raw());
 98
 99        if self.extension_factor() > 1 {
100            (0..self.extension_factor()).for_each(|i| {
101                module.switch_degree(&mut self.data[i], 0, &lut_full, 0);
102                if i < self.extension_factor() {
103                    lut_full.rotate(-1);
104                }
105            });
106        } else {
107            module.vec_znx_copy(&mut self.data[0], 0, &lut_full, 0);
108        }
109    }
110
111    #[allow(dead_code)]
112    pub(crate) fn rotate(&mut self, k: i64) {
113        let extension_factor: usize = self.extension_factor();
114        let two_n: usize = 2 * self.data[0].n();
115        let two_n_ext: usize = two_n * extension_factor;
116
117        let k_pos: usize = ((k + two_n_ext as i64) % two_n_ext as i64) as usize;
118
119        let k_hi: usize = k_pos / extension_factor;
120        let k_lo: usize = k_pos % extension_factor;
121
122        (0..extension_factor - k_lo).for_each(|i| {
123            self.data[i].rotate(k_hi as i64);
124        });
125
126        (extension_factor - k_lo..extension_factor).for_each(|i| {
127            self.data[i].rotate(k_hi as i64 + 1);
128        });
129
130        self.data.rotate_right(k_lo as usize);
131    }
132}
133
134pub(crate) trait DivRound {
135    fn div_round(self, rhs: Self) -> Self;
136}
137
138impl DivRound for usize {
139    #[inline]
140    fn div_round(self, rhs: Self) -> Self {
141        (self + rhs / 2) / rhs
142    }
143}