1use backend::{FFT64, Module, ScratchOwned, 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, extend_factor: usize) -> Self {
11        let size: usize = k.div_ceil(basek);
12        let mut data: Vec<VecZnx<Vec<u8>>> = Vec::with_capacity(extend_factor);
13        (0..extend_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 set(&mut self, module: &Module<FFT64>, f: fn(i64) -> i64, message_modulus: usize) {
24        let basek: usize = self.basek;
25
26        // Get the number minimum limb to store the message modulus
27        let limbs: usize = message_modulus.div_ceil(1 << basek);
28
29        // Scaling factor
30        let scale: i64 = (1 << (basek * limbs - 1)).div_round(message_modulus) as i64;
31
32        // Updates function
33        let f_scaled = |x: i64| (f(x) % message_modulus as i64) * scale;
34
35        // If LUT size > module.n()
36        let domain_size: usize = self.data[0].n() * self.extension_factor();
37
38        let size: usize = self.k.div_ceil(self.basek);
39
40        // Equivalent to AUTO([f(0), f(1), ..., f(n-1)], -1)
41        let mut lut_full: VecZnx<Vec<u8>> = VecZnx::new::<i64>(domain_size, 1, size);
42        {
43            let lut_at: &mut [i64] = lut_full.at_mut(0, limbs - 1);
44
45            let start: usize = 0;
46            let end: usize = (domain_size).div_round(message_modulus);
47
48            let y: i64 = f_scaled(0);
49            (start..end).for_each(|i| {
50                lut_at[i] = y;
51            });
52
53            (1..message_modulus).for_each(|x| {
54                let start: usize = (x * domain_size).div_round(message_modulus);
55                let end: usize = ((x + 1) * domain_size).div_round(message_modulus);
56                let y: i64 = f_scaled(x as i64);
57                (start..end).for_each(|i| {
58                    lut_at[domain_size - i] = -y;
59                })
60            });
61        }
62
63        // Rotates half the step to the left
64        let half_step: usize = domain_size.div_round(message_modulus << 1);
65        module.vec_znx_rotate_inplace(-(half_step as i64), &mut lut_full, 0);
66
67        let mut tmp_bytes: Vec<u8> = alloc_aligned(lut_full.n() * size_of::<i64>());
68        lut_full.normalize(self.basek, 0, &mut tmp_bytes);
69
70        if self.extension_factor() > 1 {
71            let mut scratch: ScratchOwned = ScratchOwned::new(module.bytes_of_vec_znx(1, size));
72            module.vec_znx_split(&mut self.data, 0, &lut_full, 0, scratch.borrow());
73        } else {
74            module.vec_znx_copy(&mut self.data[0], 0, &lut_full, 0);
75        }
76    }
77}
78
79pub trait DivRound {
80    fn div_round(self, rhs: Self) -> Self;
81}
82
83impl DivRound for usize {
84    #[inline]
85    fn div_round(self, rhs: Self) -> Self {
86        (self + rhs / 2) / rhs
87    }
88}