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