1use base2k::module::{Module, FFT64};
 2use base2k::scalar::Scalar;
 3use base2k::vector::Vector;
 4use itertools::izip;
 5use sampling::source::Source;
 6
 7fn main() {
 8    let n: usize = 16;
 9    let log_base2k: usize = 40;
10    let prec: usize = 54;
11    let log_scale: usize = 18;
12    let module: Module = Module::new::<FFT64>(n);
13
14    let mut carry: Vec<u8> = vec![0; module.vec_znx_big_normalize_tmp_bytes()];
15
16    let seed: [u8; 32] = [0; 32];
17    let mut source: Source = Source::new(seed);
18
19    let mut res: Vector = Vector::new(n, log_base2k, prec);
20
21    // s <- Z_{-1, 0, 1}[X]/(X^{N}+1)
22    let mut s: Scalar = Scalar::new(n);
23    s.fill_ternary_prob(0.5, &mut source);
24
25    // Buffer to store s in the DFT domain
26    let mut s_ppol: base2k::module::SVPPOL = module.svp_new_ppol();
27
28    // s_ppol <- DFT(s)
29    module.svp_prepare(&mut s_ppol, &s);
30
31    // a <- Z_{2^prec}[X]/(X^{N}+1)
32    let mut a: Vector = Vector::new(n, log_base2k, prec);
33    a.fill_uniform(&mut source);
34
35    // Scratch space for DFT values
36    let mut buf_dft: base2k::module::VECZNXDFT = module.new_vec_znx_dft(a.limbs());
37
38    // Applies buf_dft <- s * a
39    module.svp_apply_dft(&mut buf_dft, &s_ppol, &a);
40
41    // Alias scratch space
42    let mut buf_big: base2k::module::VECZNXBIG = buf_dft.as_vec_znx_big();
43
44    // buf_big <- IDFT(buf_dft) (not normalized)
45    module.vec_znx_idft_tmp_a(&mut buf_big, &mut buf_dft, a.limbs());
46
47    let mut m: Vector = Vector::new(n, log_base2k, prec - log_scale);
48    let mut want: Vec<i64> = vec![0; n];
49    want.iter_mut()
50        .for_each(|x| *x = source.next_u64n(16, 15) as i64);
51
52    // m
53    m.set_i64(&want, 4);
54    m.normalize(&mut carry);
55
56    // buf_big <- m - buf_big
57    module.vec_znx_big_sub_small_a_inplace(&mut buf_big, &m);
58
59    // b <- normalize(buf_big) + e
60    let mut b: Vector = Vector::new(n, log_base2k, prec);
61    module.vec_znx_big_normalize(&mut b, &buf_big, &mut carry);
62    b.add_normal(&mut source, 3.2, 19.0);
63
64    //Decrypt
65
66    // buf_big <- a * s
67    module.svp_apply_dft(&mut buf_dft, &s_ppol, &a);
68    module.vec_znx_idft_tmp_a(&mut buf_big, &mut buf_dft, b.limbs());
69
70    // buf_big <- a * s + b
71    module.vec_znx_big_add_small_inplace(&mut buf_big, &b);
72
73    // res <- normalize(buf_big)
74    module.vec_znx_big_normalize(&mut res, &buf_big, &mut carry);
75
76    // have = m * 2^{log_scale} + e
77    let mut have: Vec<i64> = vec![i64::default(); n];
78    res.get_i64(&mut have);
79
80    let scale: f64 = (1 << log_scale) as f64;
81    izip!(want.iter(), have.iter())
82        .enumerate()
83        .for_each(|(i, (a, b))| {
84            println!("{}: {} {}", i, a, (*b as f64) / scale);
85        })
86}