1use std::vec;
 2
 3use poulpy_hal::{
 4    api::{VecZnxCopy, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotateInplace, VecZnxSwithcDegree},
 5    layouts::{Backend, Module},
 6    oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl},
 7};
 8
 9use crate::tfhe::blind_rotation::{DivRound, LookUpTable};
10
11pub fn test_lut_standard<B>(module: &Module<B>)
12where
13    Module<B>: VecZnxRotateInplace + VecZnxNormalizeInplace<B> + VecZnxNormalizeTmpBytes + VecZnxSwithcDegree + VecZnxCopy,
14    B: Backend + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
15{
16    let basek: usize = 20;
17    let k_lut: usize = 40;
18    let message_modulus: usize = 16;
19    let extension_factor: usize = 1;
20
21    let log_scale: usize = basek + 1;
22
23    let mut f: Vec<i64> = vec![0i64; message_modulus];
24    f.iter_mut()
25        .enumerate()
26        .for_each(|(i, x)| *x = (i as i64) - 8);
27
28    let mut lut: LookUpTable = LookUpTable::alloc(module, basek, k_lut, extension_factor);
29    lut.set(module, &f, log_scale);
30
31    let half_step: i64 = lut.domain_size().div_round(message_modulus << 1) as i64;
32    lut.rotate(module, half_step);
33
34    let step: usize = lut.domain_size().div_round(message_modulus);
35
36    let mut lut_dec: Vec<i64> = vec![0i64; module.n()];
37    lut.data[0].decode_vec_i64(basek, 0, log_scale, &mut lut_dec);
38
39    (0..lut.domain_size()).step_by(step).for_each(|i| {
40        (0..step).for_each(|_| {
41            assert_eq!(f[i / step] % message_modulus as i64, lut_dec[i]);
42        });
43    });
44}
45
46pub fn test_lut_extended<B>(module: &Module<B>)
47where
48    Module<B>: VecZnxRotateInplace + VecZnxNormalizeInplace<B> + VecZnxNormalizeTmpBytes + VecZnxSwithcDegree + VecZnxCopy,
49    B: Backend + ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
50{
51    let basek: usize = 20;
52    let k_lut: usize = 40;
53    let message_modulus: usize = 16;
54    let extension_factor: usize = 4;
55
56    let log_scale: usize = basek + 1;
57
58    let mut f: Vec<i64> = vec![0i64; message_modulus];
59    f.iter_mut()
60        .enumerate()
61        .for_each(|(i, x)| *x = (i as i64) - 8);
62
63    let mut lut: LookUpTable = LookUpTable::alloc(module, basek, k_lut, extension_factor);
64    lut.set(module, &f, log_scale);
65
66    let half_step: i64 = lut.domain_size().div_round(message_modulus << 1) as i64;
67    lut.rotate(module, half_step);
68
69    let step: usize = module.n().div_round(message_modulus);
70
71    let mut lut_dec: Vec<i64> = vec![0i64; module.n()];
72
73    (0..extension_factor).for_each(|ext| {
74        lut.data[ext].decode_vec_i64(basek, 0, log_scale, &mut lut_dec);
75        (0..module.n()).step_by(step).for_each(|i| {
76            (0..step).for_each(|_| {
77                assert_eq!(f[i / step] % message_modulus as i64, lut_dec[i]);
78            });
79        });
80    });
81}