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