1use std::vec;
 2
 3use backend::hal::{
 4    api::{VecZnxCopy, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotateInplace, VecZnxSwithcDegree, ZnxView},
 5    layouts::{Backend, Module},
 6    oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl},
 7};
 8
 9use crate::{DivRound, LookUpTable};
10
11pub(crate) fn test_lut_standard<B: Backend>(module: &Module<B>)
12where
13    Module<B>: VecZnxRotateInplace + VecZnxNormalizeInplace<B> + VecZnxNormalizeTmpBytes + VecZnxSwithcDegree + VecZnxCopy,
14    B: 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    (0..lut.domain_size()).step_by(step).for_each(|i| {
38        (0..step).for_each(|_| {
39            assert_eq!(
40                f[i / step] % message_modulus as i64,
41                lut.data[0].raw()[0] / (1 << (log_scale % basek)) as i64
42            );
43            lut.rotate(module, -1);
44        });
45    });
46}
47
48pub(crate) fn test_lut_extended<B: Backend>(module: &Module<B>)
49where
50    Module<B>: VecZnxRotateInplace + VecZnxNormalizeInplace<B> + VecZnxNormalizeTmpBytes + VecZnxSwithcDegree + VecZnxCopy,
51    B: ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
52{
53    let n: usize = module.n();
54    let basek: usize = 20;
55    let k_lut: usize = 40;
56    let message_modulus: usize = 16;
57    let extension_factor: usize = 4;
58
59    let log_scale: usize = basek + 1;
60
61    let mut f: Vec<i64> = vec![0i64; message_modulus];
62    f.iter_mut()
63        .enumerate()
64        .for_each(|(i, x)| *x = (i as i64) - 8);
65
66    let mut lut: LookUpTable = LookUpTable::alloc(n, basek, k_lut, extension_factor);
67    lut.set(&module, &f, log_scale);
68
69    let half_step: i64 = lut.domain_size().div_round(message_modulus << 1) as i64;
70    lut.rotate(&module, half_step);
71
72    let step: usize = lut.domain_size().div_round(message_modulus);
73
74    (0..lut.domain_size()).step_by(step).for_each(|i| {
75        (0..step).for_each(|_| {
76            assert_eq!(
77                f[i / step] % message_modulus as i64,
78                lut.data[0].raw()[0] / (1 << (log_scale % basek)) as i64
79            );
80            lut.rotate(&module, -1);
81        });
82    });
83}