1use poulpy_core::layouts::{Base2K, Degree, TorusPrecision};
  2use poulpy_hal::{
  3    api::{
  4        ScratchOwnedAlloc, ScratchOwnedBorrow, TakeSlice, VecZnxCopy, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes,
  5        VecZnxRotateInplace, VecZnxRotateInplaceTmpBytes, VecZnxSwitchRing,
  6    },
  7    layouts::{Backend, Module, Scratch, ScratchOwned, VecZnx, ZnxInfos, ZnxViewMut},
  8    reference::{vec_znx::vec_znx_rotate_inplace, znx::ZnxRef},
  9};
 10
 11#[derive(Debug, Clone, Copy)]
 12pub enum LookUpTableRotationDirection {
 13    Left,
 14    Right,
 15}
 16
 17pub struct LookUpTableLayout {
 18    pub n: Degree,
 19    pub extension_factor: usize,
 20    pub k: TorusPrecision,
 21    pub base2k: Base2K,
 22}
 23
 24pub trait LookupTableInfos {
 25    fn n(&self) -> Degree;
 26    fn extension_factor(&self) -> usize;
 27    fn k(&self) -> TorusPrecision;
 28    fn base2k(&self) -> Base2K;
 29    fn size(&self) -> usize;
 30}
 31
 32impl LookupTableInfos for LookUpTableLayout {
 33    fn base2k(&self) -> Base2K {
 34        self.base2k
 35    }
 36
 37    fn extension_factor(&self) -> usize {
 38        self.extension_factor
 39    }
 40
 41    fn k(&self) -> TorusPrecision {
 42        self.k
 43    }
 44
 45    fn size(&self) -> usize {
 46        self.k().as_usize().div_ceil(self.base2k().as_usize())
 47    }
 48
 49    fn n(&self) -> Degree {
 50        self.n
 51    }
 52}
 53
 54pub struct LookupTable {
 55    pub(crate) data: Vec<VecZnx<Vec<u8>>>,
 56    pub(crate) rot_dir: LookUpTableRotationDirection,
 57    pub(crate) base2k: Base2K,
 58    pub(crate) k: TorusPrecision,
 59    pub(crate) drift: usize,
 60}
 61
 62impl LookupTableInfos for LookupTable {
 63    fn base2k(&self) -> Base2K {
 64        self.base2k
 65    }
 66
 67    fn extension_factor(&self) -> usize {
 68        self.data.len()
 69    }
 70
 71    fn k(&self) -> TorusPrecision {
 72        self.k
 73    }
 74
 75    fn n(&self) -> Degree {
 76        self.data[0].n().into()
 77    }
 78
 79    fn size(&self) -> usize {
 80        self.data[0].size()
 81    }
 82}
 83
 84pub trait LookupTableFactory {
 85    fn lookup_table_set(&self, res: &mut LookupTable, f: &[i64], k: usize);
 86    fn lookup_table_rotate(&self, k: i64, res: &mut LookupTable);
 87}
 88
 89impl LookupTable {
 90    pub fn alloc<A>(infos: &A) -> Self
 91    where
 92        A: LookupTableInfos,
 93    {
 94        #[cfg(debug_assertions)]
 95        {
 96            assert!(
 97                infos.extension_factor() & (infos.extension_factor() - 1) == 0,
 98                "extension_factor must be a power of two but is: {}",
 99                infos.extension_factor()
100            );
101        }
102        Self {
103            data: (0..infos.extension_factor())
104                .map(|_| VecZnx::alloc(infos.n().into(), 1, infos.size()))
105                .collect(),
106            base2k: infos.base2k(),
107            k: infos.k(),
108            drift: 0,
109            rot_dir: LookUpTableRotationDirection::Left,
110        }
111    }
112
113    pub fn log_extension_factor(&self) -> usize {
114        (usize::BITS - (self.extension_factor() - 1).leading_zeros()) as _
115    }
116
117    pub fn extension_factor(&self) -> usize {
118        self.data.len()
119    }
120
121    pub fn domain_size(&self) -> usize {
122        self.data.len() * self.data[0].n()
123    }
124
125    pub fn rotation_direction(&self) -> LookUpTableRotationDirection {
126        self.rot_dir
127    }
128
129    // By default X^{-dec(lwe)} is computed during the blind rotation.
130    // Setting [reverse_rotation] to true will reverse the sign of
131    // rotation of the LUT by instead evaluating X^{dec(lwe)} during
132    // the blind rotation.
133    pub fn set_rotation_direction(&mut self, rot_dir: LookUpTableRotationDirection) {
134        self.rot_dir = rot_dir
135    }
136
137    pub fn set<M>(&mut self, module: &M, f: &[i64], k: usize)
138    where
139        M: LookupTableFactory,
140    {
141        module.lookup_table_set(self, f, k);
142    }
143
144    pub(crate) fn rotate<M>(&mut self, module: &M, k: i64)
145    where
146        M: LookupTableFactory,
147    {
148        module.lookup_table_rotate(k, self);
149    }
150}
151
152pub(crate) trait DivRound {
153    fn div_round(self, rhs: Self) -> Self;
154}
155
156impl DivRound for usize {
157    #[inline]
158    fn div_round(self, rhs: Self) -> Self {
159        (self + rhs / 2) / rhs
160    }
161}
162
163#[allow(dead_code)]
164fn max_bit_size(vec: &[i64]) -> u32 {
165    vec.iter()
166        .map(|&v| {
167            if v == 0 {
168                0
169            } else {
170                v.unsigned_abs().ilog2() + 1
171            }
172        })
173        .max()
174        .unwrap_or(0)
175}
176
177impl<BE: Backend> LookupTableFactory for Module<BE>
178where
179    Self: VecZnxRotateInplace<BE>
180        + VecZnxNormalizeInplace<BE>
181        + VecZnxNormalizeTmpBytes
182        + VecZnxSwitchRing
183        + VecZnxCopy
184        + VecZnxRotateInplaceTmpBytes
185        + VecZnxRotateInplace<BE>
186        + VecZnxRotateInplaceTmpBytes,
187    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
188    Scratch<BE>: TakeSlice,
189{
190    fn lookup_table_set(&self, res: &mut LookupTable, f: &[i64], k: usize) {
191        assert!(f.len() <= self.n());
192
193        let base2k: usize = res.base2k.into();
194
195        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
196            self.vec_znx_normalize_tmp_bytes()
197                .max(res.domain_size() << 3),
198        );
199
200        // Get the number minimum limb to store the message modulus
201        let limbs: usize = k.div_ceil(base2k);
202
203        #[cfg(debug_assertions)]
204        {
205            assert!(f.len() <= self.n());
206            assert!(
207                (max_bit_size(f) + (k % base2k) as u32) < i64::BITS,
208                "overflow: max(|f|) << (k%base2k) > i64::BITS"
209            );
210            assert!(limbs <= res.data[0].size());
211        }
212
213        // Scaling factor
214        let mut scale = 1;
215        if !k.is_multiple_of(base2k) {
216            scale <<= base2k - (k % base2k);
217        }
218
219        // #elements in lookup table
220        let f_len: usize = f.len();
221
222        // If LUT size > TakeScalarZnx
223        let domain_size: usize = res.domain_size();
224
225        let size: usize = res.k.div_ceil(res.base2k) as usize;
226
227        // Equivalent to AUTO([f(0), -f(n-1), -f(n-2), ..., -f(1)], -1)
228        let mut lut_full: VecZnx<Vec<u8>> = VecZnx::alloc(domain_size, 1, size);
229
230        let lut_at: &mut [i64] = lut_full.at_mut(0, limbs - 1);
231
232        let step: usize = domain_size.div_round(f_len);
233
234        f.iter().enumerate().for_each(|(i, fi)| {
235            let start: usize = i * step;
236            let end: usize = start + step;
237            lut_at[start..end].fill(fi * scale);
238        });
239
240        let drift: usize = step >> 1;
241
242        // Rotates half the step to the left
243
244        if res.extension_factor() > 1 {
245            let (tmp, _) = scratch.borrow().take_slice(lut_full.n());
246
247            for i in 0..res.extension_factor() {
248                self.vec_znx_switch_ring(&mut res.data[i], 0, &lut_full, 0);
249                if i < res.extension_factor() {
250                    vec_znx_rotate_inplace::<_, ZnxRef>(-1, &mut lut_full, 0, tmp);
251                }
252            }
253        } else {
254            self.vec_znx_copy(&mut res.data[0], 0, &lut_full, 0);
255        }
256
257        for a in res.data.iter_mut() {
258            self.vec_znx_normalize_inplace(res.base2k.into(), a, 0, scratch.borrow());
259        }
260
261        res.rotate(self, -(drift as i64));
262
263        res.drift = drift
264    }
265
266    fn lookup_table_rotate(&self, k: i64, res: &mut LookupTable) {
267        let extension_factor: usize = res.extension_factor();
268        let two_n: usize = 2 * res.data[0].n();
269        let two_n_ext: usize = two_n * extension_factor;
270
271        let mut scratch: ScratchOwned<_> = ScratchOwned::alloc(self.vec_znx_rotate_inplace_tmp_bytes());
272
273        let k_pos: usize = ((k + two_n_ext as i64) % two_n_ext as i64) as usize;
274
275        let k_hi: usize = k_pos / extension_factor;
276        let k_lo: usize = k_pos % extension_factor;
277
278        (0..extension_factor - k_lo).for_each(|i| {
279            self.vec_znx_rotate_inplace(k_hi as i64, &mut res.data[i], 0, scratch.borrow());
280        });
281
282        (extension_factor - k_lo..extension_factor).for_each(|i| {
283            self.vec_znx_rotate_inplace(k_hi as i64 + 1, &mut res.data[i], 0, scratch.borrow());
284        });
285
286        res.data.rotate_right(k_lo);
287    }
288}