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