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