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