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