1use backend::hal::{
  2    api::{
  3        ScalarZnxAlloc, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAlloc, VecZnxCopy, VecZnxDftAlloc, VecZnxFillUniform,
  4        VecZnxStd, VecZnxSubABInplace,
  5    },
  6    layouts::{Backend, Module, ScratchOwned},
  7    oep::{
  8        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
  9        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 10    },
 11};
 12use sampling::source::Source;
 13
 14use crate::{
 15    Decompress, GLWECiphertext, GLWECiphertextCompressed, GLWEDecryptFamily, GLWEEncryptPkFamily, GLWEEncryptSkFamily, GLWEOps,
 16    GLWEPlaintext, GLWEPublicKey, GLWEPublicKeyExec, GLWESecret, GLWESecretExec, GLWESecretFamily, Infos,
 17};
 18
 19pub(crate) trait EncryptionTestModuleFamily<B: Backend> =
 20    GLWEDecryptFamily<B> + GLWESecretFamily<B> + VecZnxAlloc + ScalarZnxAlloc + VecZnxStd;
 21
 22pub(crate) trait EncryptionTestScratchFamily<B: Backend> = TakeVecZnxDftImpl<B>
 23    + TakeVecZnxBigImpl<B>
 24    + TakeSvpPPolImpl<B>
 25    + ScratchOwnedAllocImpl<B>
 26    + ScratchOwnedBorrowImpl<B>
 27    + ScratchAvailableImpl<B>
 28    + TakeScalarZnxImpl<B>
 29    + TakeVecZnxImpl<B>;
 30
 31pub(crate) fn test_encrypt_sk<B: Backend>(module: &Module<B>, basek: usize, k_ct: usize, k_pt: usize, sigma: f64, rank: usize)
 32where
 33    Module<B>: EncryptionTestModuleFamily<B> + GLWEEncryptSkFamily<B>,
 34    B: EncryptionTestScratchFamily<B>,
 35{
 36    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_ct, rank);
 37    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_pt);
 38    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_pt);
 39
 40    let mut source_xs: Source = Source::new([0u8; 32]);
 41    let mut source_xe: Source = Source::new([0u8; 32]);
 42    let mut source_xa: Source = Source::new([0u8; 32]);
 43
 44    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 45        GLWECiphertext::encrypt_sk_scratch_space(module, basek, ct.k())
 46            | GLWECiphertext::decrypt_scratch_space(module, basek, ct.k()),
 47    );
 48
 49    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
 50    sk.fill_ternary_prob(0.5, &mut source_xs);
 51    let sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
 52
 53    module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_pt, &mut source_xa);
 54
 55    ct.encrypt_sk(
 56        module,
 57        &pt_want,
 58        &sk_exec,
 59        &mut source_xa,
 60        &mut source_xe,
 61        sigma,
 62        scratch.borrow(),
 63    );
 64
 65    ct.decrypt(module, &mut pt_have, &sk_exec, scratch.borrow());
 66
 67    pt_want.sub_inplace_ab(module, &pt_have);
 68
 69    let noise_have: f64 = module.vec_znx_std(basek, &pt_want.data, 0) * (ct.k() as f64).exp2();
 70    let noise_want: f64 = sigma;
 71
 72    assert!(noise_have <= noise_want + 0.2);
 73}
 74
 75pub(crate) fn test_encrypt_sk_compressed<B: Backend>(
 76    module: &Module<B>,
 77    basek: usize,
 78    k_ct: usize,
 79    k_pt: usize,
 80    sigma: f64,
 81    rank: usize,
 82) where
 83    Module<B>: EncryptionTestModuleFamily<B> + GLWEEncryptSkFamily<B> + VecZnxCopy,
 84    B: EncryptionTestScratchFamily<B>,
 85{
 86    let mut ct_compressed: GLWECiphertextCompressed<Vec<u8>> = GLWECiphertextCompressed::alloc(module, basek, k_ct, rank);
 87
 88    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_pt);
 89    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_pt);
 90
 91    let mut source_xs: Source = Source::new([0u8; 32]);
 92    let mut source_xe: Source = Source::new([0u8; 32]);
 93    let mut source_xa: Source = Source::new([0u8; 32]);
 94
 95    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 96        GLWECiphertextCompressed::encrypt_sk_scratch_space(module, basek, k_ct)
 97            | GLWECiphertext::decrypt_scratch_space(module, basek, k_ct),
 98    );
 99
100    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
101    sk.fill_ternary_prob(0.5, &mut source_xs);
102    let sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
103
104    module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_pt, &mut source_xa);
105
106    let seed_xa: [u8; 32] = [1u8; 32];
107
108    ct_compressed.encrypt_sk(
109        module,
110        &pt_want,
111        &sk_exec,
112        seed_xa,
113        &mut source_xe,
114        sigma,
115        scratch.borrow(),
116    );
117
118    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_ct, rank);
119    ct.decompress(module, &ct_compressed);
120
121    ct.decrypt(module, &mut pt_have, &sk_exec, scratch.borrow());
122
123    pt_want.sub_inplace_ab(module, &pt_have);
124
125    let noise_have: f64 = module.vec_znx_std(basek, &pt_want.data, 0) * (ct.k() as f64).exp2();
126    let noise_want: f64 = sigma;
127
128    assert!(
129        noise_have <= noise_want + 0.2,
130        "{} <= {}",
131        noise_have,
132        noise_want + 0.2
133    );
134}
135
136pub(crate) fn test_encrypt_zero_sk<B: Backend>(module: &Module<B>, basek: usize, k_ct: usize, sigma: f64, rank: usize)
137where
138    Module<B>: EncryptionTestModuleFamily<B> + GLWEEncryptSkFamily<B>,
139    B: EncryptionTestScratchFamily<B>,
140{
141    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_ct);
142
143    let mut source_xs: Source = Source::new([0u8; 32]);
144    let mut source_xe: Source = Source::new([1u8; 32]);
145    let mut source_xa: Source = Source::new([0u8; 32]);
146
147    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
148    sk.fill_ternary_prob(0.5, &mut source_xs);
149    let sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
150
151    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_ct, rank);
152
153    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
154        GLWECiphertext::decrypt_scratch_space(module, basek, k_ct)
155            | GLWECiphertext::encrypt_sk_scratch_space(module, basek, k_ct),
156    );
157
158    ct.encrypt_zero_sk(
159        module,
160        &sk_exec,
161        &mut source_xa,
162        &mut source_xe,
163        sigma,
164        scratch.borrow(),
165    );
166    ct.decrypt(module, &mut pt, &sk_exec, scratch.borrow());
167
168    assert!((sigma - module.vec_znx_std(basek, &pt.data, 0) * (k_ct as f64).exp2()) <= 0.2);
169}
170
171pub(crate) fn test_encrypt_pk<B: Backend>(module: &Module<B>, basek: usize, k_ct: usize, k_pk: usize, sigma: f64, rank: usize)
172where
173    Module<B>: EncryptionTestModuleFamily<B>
174        + GLWEEncryptPkFamily<B>
175        + GLWEEncryptSkFamily<B>
176        + VecZnxDftAlloc<B>
177        + VecZnxFillUniform
178        + VecZnxSubABInplace,
179    B: EncryptionTestScratchFamily<B>,
180{
181    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_ct, rank);
182    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_ct);
183    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_ct);
184
185    let mut source_xs: Source = Source::new([0u8; 32]);
186    let mut source_xe: Source = Source::new([0u8; 32]);
187    let mut source_xa: Source = Source::new([0u8; 32]);
188    let mut source_xu: Source = Source::new([0u8; 32]);
189
190    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
191    sk.fill_ternary_prob(0.5, &mut source_xs);
192    let sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
193
194    let mut pk: GLWEPublicKey<Vec<u8>> = GLWEPublicKey::alloc(module, basek, k_pk, rank);
195    pk.generate_from_sk(module, &sk_exec, &mut source_xa, &mut source_xe, sigma);
196
197    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
198        GLWECiphertext::encrypt_sk_scratch_space(module, basek, ct.k())
199            | GLWECiphertext::decrypt_scratch_space(module, basek, ct.k())
200            | GLWECiphertext::encrypt_pk_scratch_space(module, basek, pk.k()),
201    );
202
203    module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_ct, &mut source_xa);
204
205    let pk_exec: GLWEPublicKeyExec<Vec<u8>, B> = GLWEPublicKeyExec::from(module, &pk, scratch.borrow());
206
207    ct.encrypt_pk(
208        module,
209        &pt_want,
210        &pk_exec,
211        &mut source_xu,
212        &mut source_xe,
213        sigma,
214        scratch.borrow(),
215    );
216
217    ct.decrypt(module, &mut pt_have, &sk_exec, scratch.borrow());
218
219    pt_want.sub_inplace_ab(module, &pt_have);
220
221    let noise_have: f64 = module.vec_znx_std(basek, &pt_want.data, 0).log2();
222    let noise_want: f64 = ((((rank as f64) + 1.0) * module.n() as f64 * 0.5 * sigma * sigma).sqrt()).log2() - (k_ct as f64);
223
224    assert!(
225        noise_have <= noise_want + 0.2,
226        "{} {}",
227        noise_have,
228        noise_want
229    );
230}