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