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