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