1use std::collections::HashMap;
  2
  3use backend::hal::{
  4    api::{
  5        MatZnxAlloc, ScalarZnxAlloc, ScalarZnxAllocBytes, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace,
  6        VecZnxAlloc, VecZnxAllocBytes, VecZnxAutomorphism, VecZnxBigSubSmallBInplace, VecZnxEncodeVeci64, VecZnxRotateInplace,
  7        VecZnxStd, VecZnxSwithcDegree,
  8    },
  9    layouts::{Backend, Module, ScratchOwned},
 10    oep::{
 11        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
 12        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 13    },
 14};
 15use sampling::source::Source;
 16
 17use crate::{
 18    AutomorphismKey, AutomorphismKeyExec, GGLWEExecLayoutFamily, GLWECiphertext, GLWEDecryptFamily, GLWEKeyswitchFamily, GLWEOps,
 19    GLWEPacker, GLWEPackingFamily, GLWEPlaintext, GLWESecret, GLWESecretExec, GLWESecretFamily, GLWESwitchingKeyEncryptSkFamily,
 20};
 21
 22pub(crate) trait PackingTestModuleFamily<B: Backend> = GLWEPackingFamily<B>
 23    + GLWESecretFamily<B>
 24    + GLWESwitchingKeyEncryptSkFamily<B>
 25    + GLWEKeyswitchFamily<B>
 26    + GLWEDecryptFamily<B>
 27    + GGLWEExecLayoutFamily<B>
 28    + MatZnxAlloc
 29    + VecZnxAlloc
 30    + ScalarZnxAlloc
 31    + ScalarZnxAllocBytes
 32    + VecZnxAllocBytes
 33    + VecZnxStd
 34    + VecZnxSwithcDegree
 35    + VecZnxAddScalarInplace
 36    + VecZnxEncodeVeci64
 37    + VecZnxRotateInplace
 38    + VecZnxAutomorphism
 39    + VecZnxBigSubSmallBInplace<B>;
 40
 41pub(crate) trait PackingTestScratchFamily<B: Backend> = TakeVecZnxDftImpl<B>
 42    + TakeVecZnxBigImpl<B>
 43    + TakeSvpPPolImpl<B>
 44    + ScratchOwnedAllocImpl<B>
 45    + ScratchOwnedBorrowImpl<B>
 46    + ScratchAvailableImpl<B>
 47    + TakeScalarZnxImpl<B>
 48    + TakeVecZnxImpl<B>;
 49
 50pub(crate) fn test_packing<B: Backend>(module: &Module<B>)
 51where
 52    Module<B>: PackingTestModuleFamily<B>,
 53    B: PackingTestScratchFamily<B>,
 54{
 55    let mut source_xs: Source = Source::new([0u8; 32]);
 56    let mut source_xe: Source = Source::new([0u8; 32]);
 57    let mut source_xa: Source = Source::new([0u8; 32]);
 58
 59    let basek: usize = 18;
 60    let k_ct: usize = 36;
 61    let pt_k: usize = 18;
 62    let rank: usize = 3;
 63    let sigma: f64 = 3.2;
 64    let digits: usize = 1;
 65    let k_ksk: usize = k_ct + basek * digits;
 66
 67    let rows: usize = k_ct.div_ceil(basek * digits);
 68
 69    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 70        GLWECiphertext::encrypt_sk_scratch_space(module, basek, k_ct)
 71            | AutomorphismKey::encrypt_sk_scratch_space(module, basek, k_ksk, rank)
 72            | GLWEPacker::scratch_space(module, basek, k_ct, k_ksk, digits, rank),
 73    );
 74
 75    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
 76    sk.fill_ternary_prob(0.5, &mut source_xs);
 77    let sk_dft: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
 78
 79    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_ct);
 80    let mut data: Vec<i64> = vec![0i64; module.n()];
 81    data.iter_mut().enumerate().for_each(|(i, x)| {
 82        *x = i as i64;
 83    });
 84
 85    module.encode_vec_i64(basek, &mut pt.data, 0, pt_k, &data, 32);
 86
 87    let gal_els: Vec<i64> = GLWEPacker::galois_elements(module);
 88
 89    let mut auto_keys: HashMap<i64, AutomorphismKeyExec<Vec<u8>, B>> = HashMap::new();
 90    let mut tmp: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(module, basek, k_ksk, rows, digits, rank);
 91    gal_els.iter().for_each(|gal_el| {
 92        tmp.encrypt_sk(
 93            module,
 94            *gal_el,
 95            &sk,
 96            &mut source_xa,
 97            &mut source_xe,
 98            sigma,
 99            scratch.borrow(),
100        );
101        let atk_exec: AutomorphismKeyExec<Vec<u8>, B> = AutomorphismKeyExec::from(module, &tmp, scratch.borrow());
102        auto_keys.insert(*gal_el, atk_exec);
103    });
104
105    let log_batch: usize = 0;
106
107    let mut packer: GLWEPacker = GLWEPacker::new(module, log_batch, basek, k_ct, rank);
108
109    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_ct, rank);
110
111    ct.encrypt_sk(
112        module,
113        &pt,
114        &sk_dft,
115        &mut source_xa,
116        &mut source_xe,
117        sigma,
118        scratch.borrow(),
119    );
120
121    let log_n: usize = module.log_n();
122
123    (0..module.n() >> log_batch).for_each(|i| {
124        ct.encrypt_sk(
125            module,
126            &pt,
127            &sk_dft,
128            &mut source_xa,
129            &mut source_xe,
130            sigma,
131            scratch.borrow(),
132        );
133
134        pt.rotate_inplace(module, -(1 << log_batch)); // X^-batch * pt
135
136        if reverse_bits_msb(i, log_n as u32) % 5 == 0 {
137            packer.add(module, Some(&ct), &auto_keys, scratch.borrow());
138        } else {
139            packer.add(
140                module,
141                None::<&GLWECiphertext<Vec<u8>>>,
142                &auto_keys,
143                scratch.borrow(),
144            )
145        }
146    });
147
148    let mut res = GLWECiphertext::alloc(module, basek, k_ct, rank);
149    packer.flush(module, &mut res);
150
151    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_ct);
152    let mut data: Vec<i64> = vec![0i64; module.n()];
153    data.iter_mut().enumerate().for_each(|(i, x)| {
154        if i % 5 == 0 {
155            *x = reverse_bits_msb(i, log_n as u32) as i64;
156        }
157    });
158
159    module.encode_vec_i64(basek, &mut pt_want.data, 0, pt_k, &data, 32);
160
161    res.decrypt(module, &mut pt, &sk_dft, scratch.borrow());
162
163    pt.sub_inplace_ab(module, &pt_want);
164
165    let noise_have: f64 = module.vec_znx_std(basek, &pt.data, 0).log2();
166    // println!("noise_have: {}", noise_have);
167    assert!(
168        noise_have < -((k_ct - basek) as f64),
169        "noise: {}",
170        noise_have
171    );
172}
173
174#[inline(always)]
175fn reverse_bits_msb(x: usize, n: u32) -> usize {
176    x.reverse_bits() >> (usize::BITS - n)
177}