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