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