1use backend::{FFT64, Module, ZnxView};
 2
 3use crate::blind_rotation::lut::{DivRound, LookUpTable};
 4
 5#[test]
 6fn standard() {
 7    let module: Module<FFT64> = Module::<FFT64>::new(32);
 8    let basek: usize = 20;
 9    let k_lut: usize = 40;
10    let message_modulus: usize = 16;
11    let extension_factor: usize = 1;
12
13    let scale: usize = (1 << (basek - 1)) / message_modulus;
14
15    fn lut_fn(x: i64) -> i64 {
16        x - 8
17    }
18
19    let mut lut: LookUpTable = LookUpTable::alloc(&module, basek, k_lut, extension_factor);
20    lut.set(&module, lut_fn, message_modulus);
21
22    let half_step: i64 = lut.domain_size().div_round(message_modulus << 1) as i64;
23    lut.rotate(half_step);
24
25    let step: usize = lut.domain_size().div_round(message_modulus);
26
27    (0..lut.domain_size()).step_by(step).for_each(|i| {
28        (0..step).for_each(|j| {
29            assert_eq!(
30                lut_fn((i / step) as i64) % message_modulus as i64,
31                lut.data[0].raw()[0] / scale as i64
32            );
33            lut.rotate(-1);
34        });
35    });
36}
37
38#[test]
39fn extended() {
40    let module: Module<FFT64> = Module::<FFT64>::new(32);
41    let basek: usize = 20;
42    let k_lut: usize = 40;
43    let message_modulus: usize = 16;
44    let extension_factor: usize = 4;
45
46    let scale: usize = (1 << (basek - 1)) / message_modulus;
47
48    fn lut_fn(x: i64) -> i64 {
49        x - 8
50    }
51
52    let mut lut: LookUpTable = LookUpTable::alloc(&module, basek, k_lut, extension_factor);
53    lut.set(&module, lut_fn, message_modulus);
54
55    let half_step: i64 = lut.domain_size().div_round(message_modulus << 1) as i64;
56    lut.rotate(half_step);
57
58    let step: usize = lut.domain_size().div_round(message_modulus);
59
60    (0..lut.domain_size()).step_by(step).for_each(|i| {
61        (0..step).for_each(|j| {
62            assert_eq!(
63                lut_fn((i / step) as i64) % message_modulus as i64,
64                lut.data[0].raw()[0] / scale as i64
65            );
66            lut.rotate(-1);
67        });
68    });
69}