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