1use poulpy_hal::{
  2    api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxFillUniform},
  3    layouts::{Backend, Module, Scratch, ScratchOwned},
  4    source::Source,
  5};
  6
  7use crate::{
  8    GLWECompressedEncryptSk, GLWEEncryptPk, GLWEEncryptSk, GLWEPublicKeyGenerate, GLWESub, ScratchTakeCore,
  9    decryption::GLWEDecrypt,
 10    encryption::SIGMA,
 11    layouts::{
 12        GLWE, GLWEAlloc, GLWELayout, GLWEPlaintext, GLWEPlaintextLayout, GLWEPublicKey, GLWEPublicKeyPrepare,
 13        GLWEPublicKeyPreparedAlloc, GLWESecret, GLWESecretPrepare, GLWESecretPreparedAlloc, LWEInfos,
 14        compressed::GLWECompressed,
 15        prepared::{GLWEPublicKeyPrepared, GLWESecretPrepared},
 16    },
 17};
 18
 19pub fn test_glwe_encrypt_sk<BE: Backend>(module: &Module<BE>)
 20where
 21    Module<BE>: GLWEAlloc
 22        + GLWEEncryptSk<BE>
 23        + GLWEDecrypt<BE>
 24        + GLWESecretPreparedAlloc<BE>
 25        + GLWESecretPrepare<BE>
 26        + VecZnxFillUniform
 27        + GLWESub,
 28    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 29    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 30{
 31    let base2k: usize = 8;
 32    let k_ct: usize = 54;
 33    let k_pt: usize = 30;
 34
 35    for rank in 1_usize..3 {
 36        let n: usize = module.n();
 37
 38        let glwe_infos: GLWELayout = GLWELayout {
 39            n: n.into(),
 40            base2k: base2k.into(),
 41            k: k_ct.into(),
 42            rank: rank.into(),
 43        };
 44
 45        let pt_infos: GLWEPlaintextLayout = GLWEPlaintextLayout {
 46            n: n.into(),
 47            base2k: base2k.into(),
 48            k: k_pt.into(),
 49        };
 50
 51        let mut ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(module, &glwe_infos);
 52        let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(module, &pt_infos);
 53        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(module, &pt_infos);
 54
 55        let mut source_xs: Source = Source::new([0u8; 32]);
 56        let mut source_xe: Source = Source::new([0u8; 32]);
 57        let mut source_xa: Source = Source::new([0u8; 32]);
 58
 59        let mut scratch: ScratchOwned<BE> =
 60            ScratchOwned::alloc(GLWE::encrypt_sk_tmp_bytes(module, &glwe_infos) | GLWE::decrypt_tmp_bytes(module, &glwe_infos));
 61
 62        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(module, &glwe_infos);
 63        sk.fill_ternary_prob(0.5, &mut source_xs);
 64
 65        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
 66        sk_prepared.prepare(module, &sk);
 67
 68        module.vec_znx_fill_uniform(base2k, &mut pt_want.data, 0, &mut source_xa);
 69
 70        ct.encrypt_sk(
 71            module,
 72            &pt_want,
 73            &sk_prepared,
 74            &mut source_xa,
 75            &mut source_xe,
 76            scratch.borrow(),
 77        );
 78
 79        ct.decrypt(module, &mut pt_have, &sk_prepared, scratch.borrow());
 80
 81        module.glwe_sub_inplace(&mut pt_want, &pt_have);
 82
 83        let noise_have: f64 = pt_want.data.std(base2k, 0) * (ct.k().as_u32() as f64).exp2();
 84        let noise_want: f64 = SIGMA;
 85
 86        assert!(noise_have <= noise_want + 0.2);
 87    }
 88}
 89
 90pub fn test_glwe_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
 91where
 92    Module<BE>: GLWEAlloc
 93        + GLWECompressedEncryptSk<BE>
 94        + GLWEDecrypt<BE>
 95        + GLWESecretPreparedAlloc<BE>
 96        + GLWESecretPrepare<BE>
 97        + VecZnxFillUniform
 98        + GLWESub,
 99    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
100    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
101{
102    let base2k: usize = 8;
103    let k_ct: usize = 54;
104    let k_pt: usize = 30;
105
106    for rank in 1_usize..3 {
107        // println!("rank: {}", rank);
108        let n: usize = module.n();
109
110        let glwe_infos: GLWELayout = GLWELayout {
111            n: n.into(),
112            base2k: base2k.into(),
113            k: k_ct.into(),
114            rank: rank.into(),
115        };
116
117        let pt_infos: GLWEPlaintextLayout = GLWEPlaintextLayout {
118            n: n.into(),
119            base2k: base2k.into(),
120            k: k_pt.into(),
121        };
122
123        let mut ct_compressed: GLWECompressed<Vec<u8>> = GLWECompressed::alloc_from_infos(module, &glwe_infos);
124
125        let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(module, &pt_infos);
126        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(module, &pt_infos);
127
128        let mut source_xs: Source = Source::new([0u8; 32]);
129        let mut source_xe: Source = Source::new([0u8; 32]);
130        let mut source_xa: Source = Source::new([0u8; 32]);
131
132        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
133            GLWECompressed::encrypt_sk_tmp_bytes(module, &glwe_infos) | GLWE::decrypt_tmp_bytes(module, &glwe_infos),
134        );
135
136        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(module, &glwe_infos);
137        sk.fill_ternary_prob(0.5, &mut source_xs);
138
139        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
140        sk_prepared.prepare(module, &sk);
141
142        module.vec_znx_fill_uniform(base2k, &mut pt_want.data, 0, &mut source_xa);
143
144        let seed_xa: [u8; 32] = [1u8; 32];
145
146        ct_compressed.encrypt_sk(
147            module,
148            &pt_want,
149            &sk_prepared,
150            seed_xa,
151            &mut source_xe,
152            scratch.borrow(),
153        );
154
155        let mut ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(module, &glwe_infos);
156        ct.decompress(module, &ct_compressed);
157
158        ct.decrypt(module, &mut pt_have, &sk_prepared, scratch.borrow());
159
160        module.glwe_sub_inplace(&mut pt_want, &pt_have);
161
162        let noise_have: f64 = pt_want.data.std(base2k, 0) * (ct.k().as_u32() as f64).exp2();
163        let noise_want: f64 = SIGMA;
164
165        assert!(
166            noise_have <= noise_want + 0.2,
167            "{noise_have} <= {}",
168            noise_want + 0.2
169        );
170    }
171}
172
173pub fn test_glwe_encrypt_zero_sk<BE: Backend>(module: &Module<BE>)
174where
175    Module<BE>: GLWEAlloc
176        + GLWEEncryptSk<BE>
177        + GLWEDecrypt<BE>
178        + GLWESecretPreparedAlloc<BE>
179        + GLWESecretPrepare<BE>
180        + VecZnxFillUniform
181        + GLWESub,
182    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
183    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
184{
185    let base2k: usize = 8;
186    let k_ct: usize = 54;
187
188    for rank in 1_usize..3 {
189        let n: usize = module.n();
190
191        let glwe_infos: GLWELayout = GLWELayout {
192            n: n.into(),
193            base2k: base2k.into(),
194            k: k_ct.into(),
195            rank: rank.into(),
196        };
197
198        let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(module, &glwe_infos);
199
200        let mut source_xs: Source = Source::new([0u8; 32]);
201        let mut source_xe: Source = Source::new([1u8; 32]);
202        let mut source_xa: Source = Source::new([0u8; 32]);
203
204        let mut scratch: ScratchOwned<BE> =
205            ScratchOwned::alloc(GLWE::decrypt_tmp_bytes(module, &glwe_infos) | GLWE::encrypt_sk_tmp_bytes(module, &glwe_infos));
206
207        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(module, &glwe_infos);
208        sk.fill_ternary_prob(0.5, &mut source_xs);
209
210        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
211        sk_prepared.prepare(module, &sk);
212
213        let mut ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(module, &glwe_infos);
214
215        ct.encrypt_zero_sk(
216            module,
217            &sk_prepared,
218            &mut source_xa,
219            &mut source_xe,
220            scratch.borrow(),
221        );
222        ct.decrypt(module, &mut pt, &sk_prepared, scratch.borrow());
223
224        assert!((SIGMA - pt.data.std(base2k, 0) * (k_ct as f64).exp2()) <= 0.2);
225    }
226}
227
228pub fn test_glwe_encrypt_pk<BE: Backend>(module: &Module<BE>)
229where
230    Module<BE>: GLWEAlloc
231        + GLWEEncryptPk<BE>
232        + GLWEPublicKeyPrepare<BE>
233        + GLWEPublicKeyPreparedAlloc<BE>
234        + GLWEPublicKeyGenerate<BE>
235        + GLWEDecrypt<BE>
236        + GLWESecretPreparedAlloc<BE>
237        + GLWESecretPrepare<BE>
238        + VecZnxFillUniform
239        + GLWESub,
240    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
241    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
242{
243    let base2k: usize = 8;
244    let k_ct: usize = 54;
245
246    for rank in 1_usize..3 {
247        let n: usize = module.n();
248
249        let glwe_infos: GLWELayout = GLWELayout {
250            n: n.into(),
251            base2k: base2k.into(),
252            k: k_ct.into(),
253            rank: rank.into(),
254        };
255
256        let mut ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(module, &glwe_infos);
257        let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(module, &glwe_infos);
258        let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(module, &glwe_infos);
259
260        let mut source_xs: Source = Source::new([0u8; 32]);
261        let mut source_xe: Source = Source::new([0u8; 32]);
262        let mut source_xa: Source = Source::new([0u8; 32]);
263        let mut source_xu: Source = Source::new([0u8; 32]);
264
265        let mut scratch: ScratchOwned<BE> =
266            ScratchOwned::alloc(GLWE::decrypt_tmp_bytes(module, &glwe_infos) | GLWE::encrypt_pk_tmp_bytes(module, &glwe_infos));
267
268        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(module, &glwe_infos);
269        sk.fill_ternary_prob(0.5, &mut source_xs);
270
271        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
272        sk_prepared.prepare(module, &sk);
273
274        let mut pk: GLWEPublicKey<Vec<u8>> = GLWEPublicKey::alloc_from_infos(module, &glwe_infos);
275        pk.generate(module, &sk_prepared, &mut source_xa, &mut source_xe);
276
277        module.vec_znx_fill_uniform(base2k, &mut pt_want.data, 0, &mut source_xa);
278
279        let mut pk_prepared: GLWEPublicKeyPrepared<Vec<u8>, BE> = GLWEPublicKeyPrepared::alloc_from_infos(module, &glwe_infos);
280        pk_prepared.prepare(module, &pk);
281
282        ct.encrypt_pk(
283            module,
284            &pt_want,
285            &pk_prepared,
286            &mut source_xu,
287            &mut source_xe,
288            scratch.borrow(),
289        );
290
291        ct.decrypt(module, &mut pt_have, &sk_prepared, scratch.borrow());
292
293        module.glwe_sub_inplace(&mut pt_want, &pt_have);
294
295        let noise_have: f64 = pt_want.data.std(base2k, 0).log2();
296        let noise_want: f64 = ((((rank as f64) + 1.0) * n as f64 * 0.5 * SIGMA * SIGMA).sqrt()).log2() - (k_ct as f64);
297
298        assert!(
299            noise_have <= noise_want + 0.2,
300            "{noise_have} <= {}",
301            noise_want + 0.2
302        );
303    }
304}