1use crate::{FourierGLWESecret, GLWEAutomorphismKey, GLWECiphertext, GLWEOps, GLWEPacker, GLWEPlaintext, GLWESecret};
  2use std::collections::HashMap;
  3
  4use backend::{Encoding, FFT64, Module, ScratchOwned, Stats};
  5use sampling::source::Source;
  6
  7#[test]
  8fn apply() {
  9    let log_n: usize = 5;
 10    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 11
 12    let mut source_xs: Source = Source::new([0u8; 32]);
 13    let mut source_xe: Source = Source::new([0u8; 32]);
 14    let mut source_xa: Source = Source::new([0u8; 32]);
 15
 16    let basek: usize = 18;
 17    let k_ct: usize = 36;
 18    let pt_k: usize = 18;
 19    let rank: usize = 3;
 20    let sigma: f64 = 3.2;
 21    let digits: usize = 1;
 22    let k_ksk: usize = k_ct + basek * digits;
 23
 24    let rows: usize = k_ct.div_ceil(basek * digits);
 25
 26    let mut scratch: ScratchOwned = ScratchOwned::new(
 27        GLWECiphertext::encrypt_sk_scratch_space(&module, basek, k_ct)
 28            | GLWECiphertext::decrypt_scratch_space(&module, basek, k_ct)
 29            | GLWEAutomorphismKey::encrypt_sk_scratch_space(&module, basek, k_ksk, rank)
 30            | GLWEPacker::scratch_space(&module, basek, k_ct, k_ksk, digits, rank),
 31    );
 32
 33    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
 34    sk.fill_ternary_prob(0.5, &mut source_xs);
 35    let sk_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk);
 36
 37    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
 38    let mut data: Vec<i64> = vec![0i64; module.n()];
 39    data.iter_mut().enumerate().for_each(|(i, x)| {
 40        *x = i as i64;
 41    });
 42    pt.data.encode_vec_i64(0, basek, pt_k, &data, 32);
 43
 44    let gal_els: Vec<i64> = GLWEPacker::galois_elements(&module);
 45
 46    let mut auto_keys: HashMap<i64, GLWEAutomorphismKey<Vec<u8>, FFT64>> = HashMap::new();
 47    gal_els.iter().for_each(|gal_el| {
 48        let mut key: GLWEAutomorphismKey<Vec<u8>, FFT64> = GLWEAutomorphismKey::alloc(&module, basek, k_ksk, rows, digits, rank);
 49        key.encrypt_sk(
 50            &module,
 51            *gal_el,
 52            &sk,
 53            &mut source_xa,
 54            &mut source_xe,
 55            sigma,
 56            scratch.borrow(),
 57        );
 58        auto_keys.insert(*gal_el, key);
 59    });
 60
 61    let log_batch: usize = 0;
 62
 63    let mut packer: GLWEPacker = GLWEPacker::new(&module, log_batch, basek, k_ct, rank);
 64
 65    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct, rank);
 66
 67    ct.encrypt_sk(
 68        &module,
 69        &pt,
 70        &sk_dft,
 71        &mut source_xa,
 72        &mut source_xe,
 73        sigma,
 74        scratch.borrow(),
 75    );
 76
 77    let mut res: Vec<GLWECiphertext<Vec<u8>>> = Vec::new();
 78
 79    (0..module.n() >> log_batch).for_each(|i| {
 80        ct.encrypt_sk(
 81            &module,
 82            &pt,
 83            &sk_dft,
 84            &mut source_xa,
 85            &mut source_xe,
 86            sigma,
 87            scratch.borrow(),
 88        );
 89
 90        pt.rotate_inplace(&module, -(1 << log_batch)); // X^-batch * pt
 91
 92        if reverse_bits_msb(i, log_n as u32) % 5 == 0 {
 93            packer.add(&module, &mut res, Some(&ct), &auto_keys, scratch.borrow());
 94        } else {
 95            packer.add(
 96                &module,
 97                &mut res,
 98                None::<&GLWECiphertext<Vec<u8>>>,
 99                &auto_keys,
100                scratch.borrow(),
101            )
102        }
103    });
104
105    packer.flush(&module, &mut res, &auto_keys, scratch.borrow());
106    packer.reset();
107
108    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
109
110    res.iter().enumerate().for_each(|(i, res_i)| {
111        let mut data: Vec<i64> = vec![0i64; module.n()];
112        data.iter_mut().enumerate().for_each(|(i, x)| {
113            if i % 5 == 0 {
114                *x = reverse_bits_msb(i, log_n as u32) as i64;
115            }
116        });
117        pt_want.data.encode_vec_i64(0, basek, pt_k, &data, 32);
118
119        res_i.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
120
121        if i & 1 == 0 {
122            pt.sub_inplace_ab(&module, &pt_want);
123        } else {
124            pt.add_inplace(&module, &pt_want);
125        }
126
127        let noise_have = pt.data.std(0, basek).log2();
128        // println!("noise_have: {}", noise_have);
129        assert!(
130            noise_have < -((k_ct - basek) as f64),
131            "noise: {}",
132            noise_have
133        );
134    });
135}
136
137#[inline(always)]
138fn reverse_bits_msb(x: usize, n: u32) -> usize {
139    x.reverse_bits() >> (usize::BITS - n)
140}