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    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!(noise_have <= noise_want + 0.2);
129}
130
131pub(crate) fn test_encrypt_zero_sk<B: Backend>(module: &Module<B>, basek: usize, k_ct: usize, sigma: f64, rank: usize)
132where
133    Module<B>: EncryptionTestModuleFamily<B> + GLWEEncryptSkFamily<B>,
134    B: EncryptionTestScratchFamily<B>,
135{
136    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_ct);
137
138    let mut source_xs: Source = Source::new([0u8; 32]);
139    let mut source_xe: Source = Source::new([1u8; 32]);
140    let mut source_xa: Source = Source::new([0u8; 32]);
141
142    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
143    sk.fill_ternary_prob(0.5, &mut source_xs);
144    let sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
145
146    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_ct, rank);
147
148    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
149        GLWECiphertext::decrypt_scratch_space(module, basek, k_ct)
150            | GLWECiphertext::encrypt_sk_scratch_space(module, basek, k_ct),
151    );
152
153    ct.encrypt_zero_sk(
154        module,
155        &sk_exec,
156        &mut source_xa,
157        &mut source_xe,
158        sigma,
159        scratch.borrow(),
160    );
161    ct.decrypt(module, &mut pt, &sk_exec, scratch.borrow());
162
163    assert!((sigma - module.vec_znx_std(basek, &pt.data, 0) * (k_ct as f64).exp2()) <= 0.2);
164}
165
166pub(crate) fn test_encrypt_pk<B: Backend>(module: &Module<B>, basek: usize, k_ct: usize, k_pk: usize, sigma: f64, rank: usize)
167where
168    Module<B>: EncryptionTestModuleFamily<B>
169        + GLWEEncryptPkFamily<B>
170        + GLWEEncryptSkFamily<B>
171        + VecZnxDftAlloc<B>
172        + VecZnxFillUniform
173        + VecZnxSubABInplace,
174    B: EncryptionTestScratchFamily<B>,
175{
176    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(module, basek, k_ct, rank);
177    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_ct);
178    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module, basek, k_ct);
179
180    let mut source_xs: Source = Source::new([0u8; 32]);
181    let mut source_xe: Source = Source::new([0u8; 32]);
182    let mut source_xa: Source = Source::new([0u8; 32]);
183    let mut source_xu: Source = Source::new([0u8; 32]);
184
185    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
186    sk.fill_ternary_prob(0.5, &mut source_xs);
187    let sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
188
189    let mut pk: GLWEPublicKey<Vec<u8>> = GLWEPublicKey::alloc(module, basek, k_pk, rank);
190    pk.generate_from_sk(module, &sk_exec, &mut source_xa, &mut source_xe, sigma);
191
192    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
193        GLWECiphertext::encrypt_sk_scratch_space(module, basek, ct.k())
194            | GLWECiphertext::decrypt_scratch_space(module, basek, ct.k())
195            | GLWECiphertext::encrypt_pk_scratch_space(module, basek, pk.k()),
196    );
197
198    module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_ct, &mut source_xa);
199
200    let pk_exec: GLWEPublicKeyExec<Vec<u8>, B> = GLWEPublicKeyExec::from(module, &pk, scratch.borrow());
201
202    ct.encrypt_pk(
203        module,
204        &pt_want,
205        &pk_exec,
206        &mut source_xu,
207        &mut source_xe,
208        sigma,
209        scratch.borrow(),
210    );
211
212    ct.decrypt(module, &mut pt_have, &sk_exec, scratch.borrow());
213
214    pt_want.sub_inplace_ab(module, &pt_have);
215
216    let noise_have: f64 = module.vec_znx_std(basek, &pt_want.data, 0).log2();
217    let noise_want: f64 = ((((rank as f64) + 1.0) * module.n() as f64 * 0.5 * sigma * sigma).sqrt()).log2() - (k_ct as f64);
218
219    assert!(
220        noise_have <= noise_want + 0.2,
221        "{} {}",
222        noise_have,
223        noise_want
224    );
225}