1use poulpy_hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxCopy, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotateInplace,
  4        VecZnxRotateInplaceTmpBytes, VecZnxSwitchRing,
  5    },
  6    layouts::{Backend, Module, ScratchOwned, VecZnx, ZnxInfos, ZnxViewMut},
  7};
  8
  9#[derive(Debug, Clone, Copy)]
 10pub enum LookUpTableRotationDirection {
 11    Left,
 12    Right,
 13}
 14
 15pub struct LookUpTable {
 16    pub(crate) data: Vec<VecZnx<Vec<u8>>>,
 17    pub(crate) rot_dir: LookUpTableRotationDirection,
 18    pub(crate) base2k: usize,
 19    pub(crate) k: usize,
 20    pub(crate) drift: usize,
 21}
 22
 23impl LookUpTable {
 24    pub fn alloc<B: Backend>(module: &Module<B>, base2k: usize, k: usize, extension_factor: usize) -> Self {
 25        #[cfg(debug_assertions)]
 26        {
 27            assert!(
 28                extension_factor & (extension_factor - 1) == 0,
 29                "extension_factor must be a power of two but is: {extension_factor}"
 30            );
 31        }
 32        let size: usize = k.div_ceil(base2k);
 33        let mut data: Vec<VecZnx<Vec<u8>>> = Vec::with_capacity(extension_factor);
 34        (0..extension_factor).for_each(|_| {
 35            data.push(VecZnx::alloc(module.n(), 1, size));
 36        });
 37        Self {
 38            data,
 39            base2k,
 40            k,
 41            drift: 0,
 42            rot_dir: LookUpTableRotationDirection::Left,
 43        }
 44    }
 45
 46    pub fn log_extension_factor(&self) -> usize {
 47        (usize::BITS - (self.extension_factor() - 1).leading_zeros()) as _
 48    }
 49
 50    pub fn extension_factor(&self) -> usize {
 51        self.data.len()
 52    }
 53
 54    pub fn domain_size(&self) -> usize {
 55        self.data.len() * self.data[0].n()
 56    }
 57
 58    pub fn rotation_direction(&self) -> LookUpTableRotationDirection {
 59        self.rot_dir
 60    }
 61
 62    // By default X^{-dec(lwe)} is computed during the blind rotation.
 63    // Setting [reverse_rotation] to true will reverse the sign of
 64    // rotation of the LUT by instead evaluating X^{dec(lwe)} during
 65    // the blind rotation.
 66    pub fn set_rotation_direction(&mut self, rot_dir: LookUpTableRotationDirection) {
 67        self.rot_dir = rot_dir
 68    }
 69
 70    pub fn set<B: Backend>(&mut self, module: &Module<B>, f: &[i64], k: usize)
 71    where
 72        Module<B>: VecZnxRotateInplace<B>
 73            + VecZnxNormalizeInplace<B>
 74            + VecZnxNormalizeTmpBytes
 75            + VecZnxSwitchRing
 76            + VecZnxCopy
 77            + VecZnxRotateInplaceTmpBytes,
 78        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
 79    {
 80        assert!(f.len() <= module.n());
 81
 82        let base2k: usize = self.base2k;
 83
 84        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(module.vec_znx_normalize_tmp_bytes());
 85
 86        // Get the number minimum limb to store the message modulus
 87        let limbs: usize = k.div_ceil(base2k);
 88
 89        #[cfg(debug_assertions)]
 90        {
 91            assert!(f.len() <= module.n());
 92            assert!(
 93                (max_bit_size(f) + (k % base2k) as u32) < i64::BITS,
 94                "overflow: max(|f|) << (k%base2k) > i64::BITS"
 95            );
 96            assert!(limbs <= self.data[0].size());
 97        }
 98
 99        // Scaling factor
100        let mut scale = 1;
101        if !k.is_multiple_of(base2k) {
102            scale <<= base2k - (k % base2k);
103        }
104
105        // #elements in lookup table
106        let f_len: usize = f.len();
107
108        // If LUT size > TakeScalarZnx
109        let domain_size: usize = self.domain_size();
110
111        let size: usize = self.k.div_ceil(self.base2k);
112
113        // Equivalent to AUTO([f(0), -f(n-1), -f(n-2), ..., -f(1)], -1)
114        let mut lut_full: VecZnx<Vec<u8>> = VecZnx::alloc(domain_size, 1, size);
115
116        let lut_at: &mut [i64] = lut_full.at_mut(0, limbs - 1);
117
118        let step: usize = domain_size.div_round(f_len);
119
120        f.iter().enumerate().for_each(|(i, fi)| {
121            let start: usize = i * step;
122            let end: usize = start + step;
123            lut_at[start..end].fill(fi * scale);
124        });
125
126        let drift: usize = step >> 1;
127
128        // Rotates half the step to the left
129
130        if self.extension_factor() > 1 {
131            (0..self.extension_factor()).for_each(|i| {
132                module.vec_znx_switch_ring(&mut self.data[i], 0, &lut_full, 0);
133                if i < self.extension_factor() {
134                    module.vec_znx_rotate_inplace(-1, &mut lut_full, 0, scratch.borrow());
135                }
136            });
137        } else {
138            module.vec_znx_copy(&mut self.data[0], 0, &lut_full, 0);
139        }
140
141        self.data.iter_mut().for_each(|a| {
142            module.vec_znx_normalize_inplace(self.base2k, a, 0, scratch.borrow());
143        });
144
145        self.rotate(module, -(drift as i64));
146
147        self.drift = drift
148    }
149
150    #[allow(dead_code)]
151    pub(crate) fn rotate<B: Backend>(&mut self, module: &Module<B>, k: i64)
152    where
153        Module<B>: VecZnxRotateInplace<B> + VecZnxRotateInplaceTmpBytes,
154        ScratchOwned<B>: ScratchOwnedAlloc<B> + ScratchOwnedBorrow<B>,
155    {
156        let extension_factor: usize = self.extension_factor();
157        let two_n: usize = 2 * self.data[0].n();
158        let two_n_ext: usize = two_n * extension_factor;
159
160        let mut scratch: ScratchOwned<_> = ScratchOwned::alloc(module.vec_znx_rotate_inplace_tmp_bytes());
161
162        let k_pos: usize = ((k + two_n_ext as i64) % two_n_ext as i64) as usize;
163
164        let k_hi: usize = k_pos / extension_factor;
165        let k_lo: usize = k_pos % extension_factor;
166
167        (0..extension_factor - k_lo).for_each(|i| {
168            module.vec_znx_rotate_inplace(k_hi as i64, &mut self.data[i], 0, scratch.borrow());
169        });
170
171        (extension_factor - k_lo..extension_factor).for_each(|i| {
172            module.vec_znx_rotate_inplace(k_hi as i64 + 1, &mut self.data[i], 0, scratch.borrow());
173        });
174
175        self.data.rotate_right(k_lo);
176    }
177}
178
179pub(crate) trait DivRound {
180    fn div_round(self, rhs: Self) -> Self;
181}
182
183impl DivRound for usize {
184    #[inline]
185    fn div_round(self, rhs: Self) -> Self {
186        (self + rhs / 2) / rhs
187    }
188}
189
190#[allow(dead_code)]
191fn max_bit_size(vec: &[i64]) -> u32 {
192    vec.iter()
193        .map(|&v| {
194            if v == 0 {
195                0
196            } else {
197                v.unsigned_abs().ilog2() + 1
198            }
199        })
200        .max()
201        .unwrap_or(0)
202}