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