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