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