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