1use std::collections::HashMap;
  2
  3use backend::hal::{
  4    api::{
  5        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxBigAutomorphismInplace,
  6        VecZnxBigSubSmallBInplace, VecZnxCopy, VecZnxEncodeVeci64, VecZnxFillUniform, VecZnxNormalizeInplace,
  7        VecZnxRotateInplace, VecZnxRshInplace, VecZnxStd, VecZnxSubABInplace, VecZnxSwithcDegree, ZnxView, ZnxViewMut,
  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,
 19    GLWEPlaintext, GLWESecret, GLWESecretExec, GLWESecretFamily, GLWESwitchingKeyEncryptSkFamily, Infos,
 20    noise::var_noise_gglwe_product,
 21};
 22
 23pub(crate) trait TraceTestModuleFamily<B: Backend> = GLWESecretFamily<B>
 24    + GLWESwitchingKeyEncryptSkFamily<B>
 25    + GLWEKeyswitchFamily<B>
 26    + GLWEDecryptFamily<B>
 27    + GGLWEExecLayoutFamily<B>
 28    + VecZnxStd
 29    + VecZnxSwithcDegree
 30    + VecZnxAddScalarInplace
 31    + VecZnxEncodeVeci64
 32    + VecZnxRotateInplace
 33    + VecZnxBigSubSmallBInplace<B>
 34    + VecZnxBigAutomorphismInplace<B>
 35    + VecZnxCopy
 36    + VecZnxAutomorphism
 37    + VecZnxRshInplace;
 38
 39pub(crate) trait TraceTestScratchFamily<B: Backend> = TakeVecZnxDftImpl<B>
 40    + TakeVecZnxBigImpl<B>
 41    + TakeSvpPPolImpl<B>
 42    + ScratchOwnedAllocImpl<B>
 43    + ScratchOwnedBorrowImpl<B>
 44    + ScratchAvailableImpl<B>
 45    + TakeScalarZnxImpl<B>
 46    + TakeVecZnxImpl<B>;
 47
 48pub(crate) fn test_trace_inplace<B: Backend>(module: &Module<B>, basek: usize, k: usize, sigma: f64, rank: usize)
 49where
 50    Module<B>: TraceTestModuleFamily<B>,
 51    B: TraceTestScratchFamily<B>,
 52{
 53    let n: usize = module.n();
 54    let k_autokey: usize = k + basek;
 55
 56    let digits: usize = 1;
 57    let rows: usize = k.div_ceil(basek * digits);
 58
 59    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k, rank);
 60    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k);
 61    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k);
 62
 63    let mut source_xs: Source = Source::new([0u8; 32]);
 64    let mut source_xe: Source = Source::new([0u8; 32]);
 65    let mut source_xa: Source = Source::new([0u8; 32]);
 66
 67    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 68        GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, ct.k())
 69            | GLWECiphertext::decrypt_scratch_space(module, n, basek, ct.k())
 70            | AutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_autokey, rank)
 71            | GLWECiphertext::trace_inplace_scratch_space(module, n, basek, ct.k(), k_autokey, digits, rank),
 72    );
 73
 74    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 75    sk.fill_ternary_prob(0.5, &mut source_xs);
 76    let sk_dft: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
 77
 78    let mut data_want: Vec<i64> = vec![0i64; n];
 79
 80    data_want
 81        .iter_mut()
 82        .for_each(|x| *x = source_xa.next_i64() & 0xFF);
 83
 84    module.vec_znx_fill_uniform(basek, &mut pt_have.data, 0, k, &mut source_xa);
 85
 86    ct.encrypt_sk(
 87        module,
 88        &pt_have,
 89        &sk_dft,
 90        &mut source_xa,
 91        &mut source_xe,
 92        sigma,
 93        scratch.borrow(),
 94    );
 95
 96    let mut auto_keys: HashMap<i64, AutomorphismKeyExec<Vec<u8>, B>> = HashMap::new();
 97    let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(module);
 98    let mut tmp: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_autokey, rows, digits, rank);
 99    gal_els.iter().for_each(|gal_el| {
100        tmp.encrypt_sk(
101            module,
102            *gal_el,
103            &sk,
104            &mut source_xa,
105            &mut source_xe,
106            sigma,
107            scratch.borrow(),
108        );
109        let atk_exec: AutomorphismKeyExec<Vec<u8>, B> = AutomorphismKeyExec::from(module, &tmp, scratch.borrow());
110        auto_keys.insert(*gal_el, atk_exec);
111    });
112
113    ct.trace_inplace(module, 0, 5, &auto_keys, scratch.borrow());
114    ct.trace_inplace(module, 5, module.log_n(), &auto_keys, scratch.borrow());
115
116    (0..pt_want.size()).for_each(|i| pt_want.data.at_mut(0, i)[0] = pt_have.data.at(0, i)[0]);
117
118    ct.decrypt(module, &mut pt_have, &sk_dft, scratch.borrow());
119
120    module.vec_znx_sub_ab_inplace(&mut pt_want.data, 0, &pt_have.data, 0);
121    module.vec_znx_normalize_inplace(basek, &mut pt_want.data, 0, scratch.borrow());
122
123    let noise_have: f64 = module.vec_znx_std(basek, &pt_want.data, 0).log2();
124
125    let mut noise_want: f64 = var_noise_gglwe_product(
126        n as f64,
127        basek,
128        0.5,
129        0.5,
130        1.0 / 12.0,
131        sigma * sigma,
132        0.0,
133        rank as f64,
134        k,
135        k_autokey,
136    );
137    noise_want += sigma * sigma * (-2.0 * (k) as f64).exp2();
138    noise_want += n as f64 * 1.0 / 12.0 * 0.5 * rank as f64 * (-2.0 * (k) as f64).exp2();
139    noise_want = noise_want.sqrt().log2();
140
141    assert!(
142        (noise_have - noise_want).abs() < 1.0,
143        "{} > {}",
144        noise_have,
145        noise_want
146    );
147}