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