1use std::vec;
 2
 3use poulpy_hal::api::ModuleN;
 4
 5use crate::bin_fhe::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().enumerate().for_each(|(i, x)| *x = (i as i64) - 8);
20
21    let lut_infos: LookUpTableLayout = LookUpTableLayout {
22        n: module.n().into(),
23        extension_factor,
24        k: k_lut.into(),
25        base2k: base2k.into(),
26    };
27
28    let mut lut: LookupTable = LookupTable::alloc(&lut_infos);
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(base2k, 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<M>(module: &M)
47where
48    M: LookupTableFactory + ModuleN,
49{
50    let base2k: usize = 20;
51    let k_lut: usize = 40;
52    let message_modulus: usize = 16;
53    let extension_factor: usize = 4;
54
55    let log_scale: usize = base2k + 1;
56
57    let mut f: Vec<i64> = vec![0i64; message_modulus];
58    f.iter_mut().enumerate().for_each(|(i, x)| *x = (i as i64) - 8);
59
60    let lut_infos: LookUpTableLayout = LookUpTableLayout {
61        n: module.n().into(),
62        extension_factor,
63        k: k_lut.into(),
64        base2k: base2k.into(),
65    };
66
67    let mut lut: LookupTable = LookupTable::alloc(&lut_infos);
68    lut.set(module, &f, log_scale);
69
70    let half_step: i64 = lut.domain_size().div_round(message_modulus << 1) as i64;
71    lut.rotate(module, half_step);
72
73    let step: usize = module.n().div_round(message_modulus);
74
75    let mut lut_dec: Vec<i64> = vec![0i64; module.n()];
76
77    (0..extension_factor).for_each(|ext| {
78        lut.data[ext].decode_vec_i64(base2k, 0, log_scale, &mut lut_dec);
79        (0..module.n()).step_by(step).for_each(|i| {
80            (0..step).for_each(|_| {
81                assert_eq!(f[i / step] % message_modulus as i64, lut_dec[i]);
82            });
83        });
84    });
85}