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