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