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.noise(module, &pt_want, &sk_prepared, scratch.borrow()).std().log2();
 72        let noise_want: f64 = SIGMA.log2() - (ct.k().as_usize() as f64) + 0.5;
 73
 74        assert!(noise_have <= noise_want + 0.2);
 75    }
 76}
 77
 78pub fn test_glwe_compressed_encrypt_sk<BE: Backend>(module: &Module<BE>)
 79where
 80    Module<BE>: GLWECompressedEncryptSk<BE> + GLWENoise<BE> + GLWESecretPreparedFactory<BE> + VecZnxFillUniform + GLWESub,
 81    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
 82    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
 83{
 84    let base2k: usize = 8;
 85    let k_ct: usize = 54;
 86    let k_pt: usize = 30;
 87
 88    for rank in 1_usize..3 {
 89        let n: usize = module.n();
 90
 91        let glwe_infos: GLWELayout = GLWELayout {
 92            n: n.into(),
 93            base2k: base2k.into(),
 94            k: k_ct.into(),
 95            rank: rank.into(),
 96        };
 97
 98        let pt_infos: GLWEPlaintextLayout = GLWEPlaintextLayout {
 99            n: n.into(),
100            base2k: base2k.into(),
101            k: k_pt.into(),
102        };
103
104        let mut ct_compressed: GLWECompressed<Vec<u8>> = GLWECompressed::alloc_from_infos(&glwe_infos);
105
106        let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&pt_infos);
107
108        let mut source_xs: Source = Source::new([0u8; 32]);
109        let mut source_xe: Source = Source::new([0u8; 32]);
110        let mut source_xa: Source = Source::new([0u8; 32]);
111
112        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
113            GLWECompressed::encrypt_sk_tmp_bytes(module, &glwe_infos).max(module.glwe_noise_tmp_bytes(&glwe_infos)),
114        );
115
116        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_infos);
117        sk.fill_ternary_prob(0.5, &mut source_xs);
118
119        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
120        sk_prepared.prepare(module, &sk);
121
122        module.vec_znx_fill_uniform(base2k, &mut pt_want.data, 0, &mut source_xa);
123
124        let seed_xa: [u8; 32] = [1u8; 32];
125
126        ct_compressed.encrypt_sk(module, &pt_want, &sk_prepared, seed_xa, &mut source_xe, scratch.borrow());
127
128        let mut ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_infos);
129        ct.decompress(module, &ct_compressed);
130
131        let noise_have: f64 = ct.noise(module, &pt_want, &sk_prepared, scratch.borrow()).std().log2();
132        let noise_want: f64 = SIGMA.log2() - (ct.k().as_usize() as f64) + 0.5;
133        assert!(noise_have <= noise_want + 0.2);
134    }
135}
136
137pub fn test_glwe_encrypt_zero_sk<BE: Backend>(module: &Module<BE>)
138where
139    Module<BE>: GLWEEncryptSk<BE> + GLWENoise<BE> + GLWESecretPreparedFactory<BE> + VecZnxFillUniform + GLWESub,
140    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
141    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
142{
143    let base2k: usize = 8;
144    let k_ct: usize = 54;
145
146    for rank in 1_usize..3 {
147        let n: usize = module.n();
148
149        let glwe_infos: GLWELayout = GLWELayout {
150            n: n.into(),
151            base2k: base2k.into(),
152            k: k_ct.into(),
153            rank: rank.into(),
154        };
155
156        let pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_infos);
157
158        let mut source_xs: Source = Source::new([0u8; 32]);
159        let mut source_xe: Source = Source::new([1u8; 32]);
160        let mut source_xa: Source = Source::new([0u8; 32]);
161
162        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
163            module
164                .glwe_noise_tmp_bytes(&glwe_infos)
165                .max(GLWE::encrypt_sk_tmp_bytes(module, &glwe_infos)),
166        );
167
168        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_infos);
169        sk.fill_ternary_prob(0.5, &mut source_xs);
170
171        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
172        sk_prepared.prepare(module, &sk);
173
174        let mut ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_infos);
175
176        ct.encrypt_zero_sk(module, &sk_prepared, &mut source_xa, &mut source_xe, scratch.borrow());
177
178        let noise_have: f64 = ct.noise(module, &pt, &sk_prepared, scratch.borrow()).std().log2();
179        let noise_want: f64 = SIGMA.log2() - (ct.k().as_usize() as f64) + 0.5;
180        assert!(noise_have <= noise_want + 0.2);
181    }
182}
183
184pub fn test_glwe_encrypt_pk<BE: Backend>(module: &Module<BE>)
185where
186    Module<BE>: GLWEEncryptPk<BE>
187        + GLWEPublicKeyPreparedFactory<BE>
188        + GLWEPublicKeyGenerate<BE>
189        + GLWENoise<BE>
190        + GLWESecretPreparedFactory<BE>
191        + VecZnxFillUniform
192        + GLWESub,
193    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
194    Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
195{
196    let base2k: usize = 8;
197    let k_ct: usize = 54;
198
199    for rank in 1_usize..3 {
200        let n: usize = module.n();
201
202        let glwe_infos: GLWELayout = GLWELayout {
203            n: n.into(),
204            base2k: base2k.into(),
205            k: k_ct.into(),
206            rank: rank.into(),
207        };
208
209        let mut ct: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_infos);
210        let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_infos);
211
212        let mut source_xs: Source = Source::new([0u8; 32]);
213        let mut source_xe: Source = Source::new([0u8; 32]);
214        let mut source_xa: Source = Source::new([0u8; 32]);
215        let mut source_xu: Source = Source::new([0u8; 32]);
216
217        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
218            module
219                .glwe_noise_tmp_bytes(&glwe_infos)
220                .max(GLWE::encrypt_pk_tmp_bytes(module, &glwe_infos)),
221        );
222
223        let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc_from_infos(&glwe_infos);
224        sk.fill_ternary_prob(0.5, &mut source_xs);
225
226        let mut sk_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
227        sk_prepared.prepare(module, &sk);
228
229        let mut pk: GLWEPublicKey<Vec<u8>> = GLWEPublicKey::alloc_from_infos(&glwe_infos);
230        pk.generate(module, &sk_prepared, &mut source_xa, &mut source_xe);
231
232        module.vec_znx_fill_uniform(base2k, &mut pt_want.data, 0, &mut source_xa);
233
234        let mut pk_prepared: GLWEPublicKeyPrepared<Vec<u8>, BE> = GLWEPublicKeyPrepared::alloc_from_infos(module, &glwe_infos);
235        pk_prepared.prepare(module, &pk);
236
237        ct.encrypt_pk(
238            module,
239            &pt_want,
240            &pk_prepared,
241            &mut source_xu,
242            &mut source_xe,
243            scratch.borrow(),
244        );
245
246        let noise_have: f64 = ct.noise(module, &pt_want, &sk_prepared, scratch.borrow()).std().log2();
247        let noise_want: f64 = ((((rank as f64) + 1.0) * n as f64 * 0.5 * SIGMA * SIGMA).sqrt()).log2() - (k_ct as f64);
248        assert!(noise_have <= noise_want + 0.2);
249    }
250}