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