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