1use std::collections::HashMap;
  2
  3use poulpy_hal::{
  4    api::{
  5        ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
  6        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxBigAddInplace,
  7        VecZnxBigAddSmallInplace, VecZnxBigAllocBytes, VecZnxBigAutomorphismInplace, VecZnxBigNormalize,
  8        VecZnxBigNormalizeTmpBytes, VecZnxBigSubSmallNegateInplace, VecZnxCopy, VecZnxDftAllocBytes, VecZnxDftApply,
  9        VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes,
 10        VecZnxRotateInplace, VecZnxRshInplace, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing, VmpApplyDftToDft,
 11        VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPrepare,
 12    },
 13    layouts::{Backend, Module, ScratchOwned, ZnxView, ZnxViewMut},
 14    oep::{
 15        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
 16        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 17    },
 18    source::Source,
 19};
 20
 21use crate::{
 22    encryption::SIGMA,
 23    layouts::{
 24        GGLWEAutomorphismKey, GGLWEAutomorphismKeyLayout, GLWECiphertext, GLWECiphertextLayout, GLWEPlaintext, GLWESecret,
 25        LWEInfos,
 26        prepared::{GGLWEAutomorphismKeyPrepared, GLWESecretPrepared, PrepareAlloc},
 27    },
 28    noise::var_noise_gglwe_product,
 29};
 30
 31pub fn test_glwe_trace_inplace<B>(module: &Module<B>)
 32where
 33    Module<B>: VecZnxDftAllocBytes
 34        + VecZnxAutomorphism
 35        + VecZnxBigAutomorphismInplace<B>
 36        + VecZnxBigSubSmallNegateInplace<B>
 37        + VecZnxRshInplace<B>
 38        + VecZnxRotateInplace<B>
 39        + VecZnxBigNormalize<B>
 40        + VecZnxDftApply<B>
 41        + SvpApplyDftToDftInplace<B>
 42        + VecZnxIdftApplyConsume<B>
 43        + VecZnxFillUniform
 44        + VecZnxSubInplace
 45        + VecZnxAddInplace
 46        + VecZnxNormalizeInplace<B>
 47        + VecZnxAddNormal
 48        + VecZnxNormalize<B>
 49        + VecZnxSub
 50        + SvpPrepare<B>
 51        + SvpPPolAllocBytes
 52        + SvpPPolAlloc<B>
 53        + VecZnxBigAllocBytes
 54        + VecZnxBigAddInplace<B>
 55        + VecZnxBigAddSmallInplace<B>
 56        + VecZnxNormalizeTmpBytes
 57        + VecZnxAddScalarInplace
 58        + VmpPMatAlloc<B>
 59        + VmpPrepare<B>
 60        + VmpApplyDftToDftTmpBytes
 61        + VmpApplyDftToDft<B>
 62        + VmpApplyDftToDftAdd<B>
 63        + VecZnxBigNormalizeTmpBytes
 64        + VecZnxSwitchRing
 65        + VecZnxCopy,
 66    B: Backend
 67        + TakeVecZnxDftImpl<B>
 68        + TakeVecZnxBigImpl<B>
 69        + TakeSvpPPolImpl<B>
 70        + ScratchOwnedAllocImpl<B>
 71        + ScratchOwnedBorrowImpl<B>
 72        + ScratchAvailableImpl<B>
 73        + TakeScalarZnxImpl<B>
 74        + TakeVecZnxImpl<B>,
 75{
 76    let base2k: usize = 8;
 77    let k: usize = 54;
 78
 79    for rank in 1_usize..3 {
 80        let n: usize = module.n();
 81        let k_autokey: usize = k + base2k;
 82
 83        let digits: usize = 1;
 84        let rows: usize = k.div_ceil(base2k * digits);
 85
 86        let glwe_out_infos: GLWECiphertextLayout = GLWECiphertextLayout {
 87            n: n.into(),
 88            base2k: base2k.into(),
 89            k: k.into(),
 90            rank: rank.into(),
 91        };
 92
 93        let key_infos: GGLWEAutomorphismKeyLayout = GGLWEAutomorphismKeyLayout {
 94            n: n.into(),
 95            base2k: base2k.into(),
 96            k: k_autokey.into(),
 97            rank: rank.into(),
 98            digits: digits.into(),
 99            rows: rows.into(),
100        };
101
102        let mut glwe_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&glwe_out_infos);
103        let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&glwe_out_infos);
104        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&glwe_out_infos);
105
106        let mut source_xs: Source = Source::new([0u8; 32]);
107        let mut source_xe: Source = Source::new([0u8; 32]);
108        let mut source_xa: Source = Source::new([0u8; 32]);
109
110        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
111            GLWECiphertext::encrypt_sk_scratch_space(module, &glwe_out_infos)
112                | GLWECiphertext::decrypt_scratch_space(module, &glwe_out_infos)
113                | GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, &key_infos)
114                | GLWECiphertext::trace_inplace_scratch_space(module, &glwe_out_infos, &key_infos),
115        );
116
117        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&glwe_out_infos);
118        sk.fill_ternary_prob(0.5, &mut source_xs);
119        let sk_dft: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
120
121        let mut data_want: Vec<i64> = vec![0i64; n];
122
123        data_want
124            .iter_mut()
125            .for_each(|x| *x = source_xa.next_i64() & 0xFF);
126
127        module.vec_znx_fill_uniform(base2k, &mut pt_have.data, 0, &mut source_xa);
128
129        glwe_out.encrypt_sk(
130            module,
131            &pt_have,
132            &sk_dft,
133            &mut source_xa,
134            &mut source_xe,
135            scratch.borrow(),
136        );
137
138        let mut auto_keys: HashMap<i64, GGLWEAutomorphismKeyPrepared<Vec<u8>, B>> = HashMap::new();
139        let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(module);
140        let mut tmp: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(&key_infos);
141        gal_els.iter().for_each(|gal_el| {
142            tmp.encrypt_sk(
143                module,
144                *gal_el,
145                &sk,
146                &mut source_xa,
147                &mut source_xe,
148                scratch.borrow(),
149            );
150            let atk_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> = tmp.prepare_alloc(module, scratch.borrow());
151            auto_keys.insert(*gal_el, atk_prepared);
152        });
153
154        glwe_out.trace_inplace(module, 0, 5, &auto_keys, scratch.borrow());
155        glwe_out.trace_inplace(module, 5, module.log_n(), &auto_keys, scratch.borrow());
156
157        (0..pt_want.size()).for_each(|i| pt_want.data.at_mut(0, i)[0] = pt_have.data.at(0, i)[0]);
158
159        glwe_out.decrypt(module, &mut pt_have, &sk_dft, scratch.borrow());
160
161        module.vec_znx_sub_inplace(&mut pt_want.data, 0, &pt_have.data, 0);
162        module.vec_znx_normalize_inplace(base2k, &mut pt_want.data, 0, scratch.borrow());
163
164        let noise_have: f64 = pt_want.std().log2();
165
166        let mut noise_want: f64 = var_noise_gglwe_product(
167            n as f64,
168            base2k,
169            0.5,
170            0.5,
171            1.0 / 12.0,
172            SIGMA * SIGMA,
173            0.0,
174            rank as f64,
175            k,
176            k_autokey,
177        );
178        noise_want += SIGMA * SIGMA * (-2.0 * (k) as f64).exp2();
179        noise_want += n as f64 * 1.0 / 12.0 * 0.5 * rank as f64 * (-2.0 * (k) as f64).exp2();
180        noise_want = noise_want.sqrt().log2();
181
182        assert!(
183            (noise_have - noise_want).abs() < 1.0,
184            "{noise_have} > {noise_want}"
185        );
186    }
187}