1use std::vec;
 2
 3use backend::hal::{
 4    api::{
 5        VecZnxCopy, VecZnxDecodeVeci64, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxRotateInplace, VecZnxSwithcDegree,
 6    },
 7    layouts::{Backend, Module},
 8    oep::{ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl},
 9};
10
11use crate::{DivRound, LookUpTable};
12
13pub(crate) fn test_lut_standard<B: Backend>(module: &Module<B>)
14where
15    Module<B>: VecZnxRotateInplace
16        + VecZnxNormalizeInplace<B>
17        + VecZnxNormalizeTmpBytes
18        + VecZnxSwithcDegree
19        + VecZnxCopy
20        + VecZnxDecodeVeci64,
21    B: ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
22{
23    let n: usize = module.n();
24    let basek: 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 = basek + 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(n, basek, 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    module.decode_vec_i64(basek, &lut.data[0], 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(crate) fn test_lut_extended<B: Backend>(module: &Module<B>)
55where
56    Module<B>: VecZnxRotateInplace
57        + VecZnxNormalizeInplace<B>
58        + VecZnxNormalizeTmpBytes
59        + VecZnxSwithcDegree
60        + VecZnxCopy
61        + VecZnxDecodeVeci64,
62    B: ScratchOwnedAllocImpl<B> + ScratchOwnedBorrowImpl<B>,
63{
64    let n: usize = module.n();
65    let basek: usize = 20;
66    let k_lut: usize = 40;
67    let message_modulus: usize = 16;
68    let extension_factor: usize = 4;
69
70    let log_scale: usize = basek + 1;
71
72    let mut f: Vec<i64> = vec![0i64; message_modulus];
73    f.iter_mut()
74        .enumerate()
75        .for_each(|(i, x)| *x = (i as i64) - 8);
76
77    let mut lut: LookUpTable = LookUpTable::alloc(n, basek, k_lut, extension_factor);
78    lut.set(&module, &f, log_scale);
79
80    let half_step: i64 = lut.domain_size().div_round(message_modulus << 1) as i64;
81    lut.rotate(&module, half_step);
82
83    let step: usize = module.n().div_round(message_modulus);
84
85    let mut lut_dec: Vec<i64> = vec![0i64; module.n()];
86
87    (0..extension_factor).for_each(|ext| {
88        module.decode_vec_i64(basek, &lut.data[ext], 0, log_scale, &mut lut_dec);
89        (0..module.n()).step_by(step).for_each(|i| {
90            (0..step).for_each(|_| {
91                assert_eq!(f[i / step] % message_modulus as i64, lut_dec[i]);
92            });
93        });
94    });
95}