1use std::collections::HashMap;
  2
  3use poulpy_hal::{
  4    api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxFillUniform, VecZnxNormalizeInplace, VecZnxSubInplace},
  5    layouts::{Backend, Module, Scratch, ScratchOwned, ZnxView, ZnxViewMut},
  6    source::Source,
  7};
  8
  9use crate::{
 10    GLWEAutomorphismKeyEncryptSk, GLWEDecrypt, GLWEEncryptSk, ScratchTakeCore,
 11    encryption::SIGMA,
 12    glwe_trace::GLWETrace,
 13    layouts::{
 14        GLWE, GLWEAutomorphismKey, GLWEAutomorphismKeyLayout, GLWEAutomorphismKeyPreparedFactory, GLWELayout, GLWEPlaintext,
 15        GLWESecret, GLWESecretPreparedFactory, LWEInfos,
 16        prepared::{GLWEAutomorphismKeyPrepared, GLWESecretPrepared},
 17    },
 18    noise::var_noise_gglwe_product,
 19};
 20
 21pub fn test_glwe_trace_inplace<BE: Backend>(module: &Module<BE>)
 22where
 23    Module<BE>: GLWETrace<BE>
 24        + GLWEEncryptSk<BE>
 25        + GLWEDecrypt<BE>
 26        + GLWEAutomorphismKeyEncryptSk<BE>
 27        + GLWEAutomorphismKeyPreparedFactory<BE>
 28        + VecZnxFillUniform
 29        + GLWESecretPreparedFactory<BE>
 30        + VecZnxSubInplace
 31        + VecZnxNormalizeInplace<BE>,
 32    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 33    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 34{
 35    let base2k_out: usize = 15;
 36    let base2k_key: usize = 10;
 37    let k: usize = 54;
 38
 39    for rank in 1_usize..3 {
 40        let n: usize = module.n();
 41        let k_autokey: usize = k + base2k_key;
 42
 43        let dsize: usize = 1;
 44        let dnum: usize = k.div_ceil(base2k_key * dsize);
 45
 46        let glwe_out_infos: GLWELayout = GLWELayout {
 47            n: n.into(),
 48            base2k: base2k_out.into(),
 49            k: k.into(),
 50            rank: rank.into(),
 51        };
 52
 53        let key_infos: GLWEAutomorphismKeyLayout = GLWEAutomorphismKeyLayout {
 54            n: n.into(),
 55            base2k: base2k_key.into(),
 56            k: k_autokey.into(),
 57            rank: rank.into(),
 58            dsize: dsize.into(),
 59            dnum: dnum.into(),
 60        };
 61
 62        let mut glwe_out: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_infos);
 63        let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
 64        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
 65
 66        let mut source_xs: Source = Source::new([0u8; 32]);
 67        let mut source_xe: Source = Source::new([0u8; 32]);
 68        let mut source_xa: Source = Source::new([0u8; 32]);
 69
 70        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
 71            GLWE::encrypt_sk_tmp_bytes(module, &glwe_out_infos)
 72                | GLWE::decrypt_tmp_bytes(module, &glwe_out_infos)
 73                | GLWEAutomorphismKey::encrypt_sk_tmp_bytes(module, &key_infos)
 74                | GLWE::trace_tmp_bytes(module, &glwe_out_infos, &glwe_out_infos, &key_infos),
 75        );
 76
 77        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_out_infos);
 78        sk.fill_ternary_prob(0.5, &mut source_xs);
 79
 80        let mut sk_dft: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(module, &sk);
 81        sk_dft.prepare(module, &sk);
 82
 83        let mut data_want: Vec<i64> = vec![0i64; 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(base2k_out, &mut pt_have.data, 0, &mut source_xa);
 90
 91        glwe_out.encrypt_sk(
 92            module,
 93            &pt_have,
 94            &sk_dft,
 95            &mut source_xa,
 96            &mut source_xe,
 97            scratch.borrow(),
 98        );
 99
100        let mut auto_keys: HashMap<i64, GLWEAutomorphismKeyPrepared<Vec<u8>, BE>> = HashMap::new();
101        let gal_els: Vec<i64> = GLWE::trace_galois_elements(module);
102        let mut tmp: GLWEAutomorphismKey<Vec<u8>> = GLWEAutomorphismKey::alloc_from_infos(&key_infos);
103        gal_els.iter().for_each(|gal_el| {
104            tmp.encrypt_sk(
105                module,
106                *gal_el,
107                &sk,
108                &mut source_xa,
109                &mut source_xe,
110                scratch.borrow(),
111            );
112            let mut atk_prepared: GLWEAutomorphismKeyPrepared<Vec<u8>, BE> =
113                GLWEAutomorphismKeyPrepared::alloc_from_infos(module, &tmp);
114            atk_prepared.prepare(module, &tmp, scratch.borrow());
115            auto_keys.insert(*gal_el, atk_prepared);
116        });
117
118        glwe_out.trace_inplace(module, 0, &auto_keys, scratch.borrow());
119
120        (0..pt_want.size()).for_each(|i| pt_want.data.at_mut(0, i)[0] = pt_have.data.at(0, i)[0]);
121
122        glwe_out.decrypt(module, &mut pt_have, &sk_dft, scratch.borrow());
123
124        module.vec_znx_sub_inplace(&mut pt_want.data, 0, &pt_have.data, 0);
125        module.vec_znx_normalize_inplace(
126            pt_want.base2k().as_usize(),
127            &mut pt_want.data,
128            0,
129            scratch.borrow(),
130        );
131
132        let noise_have: f64 = pt_want.stats().std().log2();
133
134        let mut noise_want: f64 = var_noise_gglwe_product(
135            n as f64,
136            base2k_key * dsize,
137            0.5,
138            0.5,
139            1.0 / 12.0,
140            SIGMA * SIGMA,
141            0.0,
142            rank as f64,
143            k,
144            k_autokey,
145        );
146        noise_want += SIGMA * SIGMA * (-2.0 * (k) as f64).exp2();
147        noise_want += n as f64 * 1.0 / 12.0 * 0.5 * rank as f64 * (-2.0 * (k) as f64).exp2();
148        noise_want = noise_want.sqrt().log2();
149
150        assert!(
151            (noise_have - noise_want).abs() < 1.0,
152            "{noise_have} > {noise_want}"
153        );
154    }
155}