1use backend::{FFT64, FillUniform, Module, ScratchOwned, Stats};
  2use sampling::source::Source;
  3
  4use crate::{FourierGLWECiphertext, FourierGLWESecret, GLWECiphertext, GLWEOps, GLWEPlaintext, GLWEPublicKey, GLWESecret, Infos};
  5
  6#[test]
  7fn encrypt_sk() {
  8    let log_n: usize = 8;
  9    (1..4).for_each(|rank| {
 10        println!("test encrypt_sk rank: {}", rank);
 11        test_encrypt_sk(log_n, 8, 54, 30, 3.2, rank);
 12    });
 13}
 14
 15#[test]
 16fn encrypt_zero_sk() {
 17    let log_n: usize = 8;
 18    (1..4).for_each(|rank| {
 19        println!("test encrypt_zero_sk rank: {}", rank);
 20        test_encrypt_zero_sk(log_n, 8, 64, 3.2, rank);
 21    });
 22}
 23
 24#[test]
 25fn encrypt_pk() {
 26    let log_n: usize = 8;
 27    (1..4).for_each(|rank| {
 28        println!("test encrypt_pk rank: {}", rank);
 29        test_encrypt_pk(log_n, 8, 64, 64, 3.2, rank)
 30    });
 31}
 32
 33fn test_encrypt_sk(log_n: usize, basek: usize, k_ct: usize, k_pt: usize, sigma: f64, rank: usize) {
 34    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 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 = ScratchOwned::new(
 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_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk);
 52
 53    pt_want
 54        .data
 55        .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
 56
 57    ct.encrypt_sk(
 58        &module,
 59        &pt_want,
 60        &sk_dft,
 61        &mut source_xa,
 62        &mut source_xe,
 63        sigma,
 64        scratch.borrow(),
 65    );
 66
 67    ct.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
 68
 69    pt_want.sub_inplace_ab(&module, &pt_have);
 70
 71    let noise_have: f64 = pt_want.data.std(0, basek) * (ct.k() as f64).exp2();
 72    let noise_want: f64 = sigma;
 73
 74    assert!(noise_have <= noise_want + 0.2);
 75}
 76
 77fn test_encrypt_zero_sk(log_n: usize, basek: usize, k_ct: usize, sigma: f64, rank: usize) {
 78    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 79
 80    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
 81
 82    let mut source_xs: Source = Source::new([0u8; 32]);
 83    let mut source_xe: Source = Source::new([1u8; 32]);
 84    let mut source_xa: Source = Source::new([0u8; 32]);
 85
 86    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
 87    sk.fill_ternary_prob(0.5, &mut source_xs);
 88    let sk_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk);
 89
 90    let mut ct_dft: FourierGLWECiphertext<Vec<u8>, FFT64> = FourierGLWECiphertext::alloc(&module, basek, k_ct, rank);
 91
 92    let mut scratch: ScratchOwned = ScratchOwned::new(
 93        FourierGLWECiphertext::decrypt_scratch_space(&module, basek, k_ct)
 94            | FourierGLWECiphertext::encrypt_sk_scratch_space(&module, basek, k_ct, rank),
 95    );
 96
 97    ct_dft.encrypt_zero_sk(
 98        &module,
 99        &sk_dft,
100        &mut source_xa,
101        &mut source_xe,
102        sigma,
103        scratch.borrow(),
104    );
105    ct_dft.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
106
107    assert!((sigma - pt.data.std(0, basek) * (k_ct as f64).exp2()) <= 0.2);
108}
109
110fn test_encrypt_pk(log_n: usize, basek: usize, k_ct: usize, k_pk: usize, sigma: f64, rank: usize) {
111    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
112
113    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct, rank);
114    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
115    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
116
117    let mut source_xs: Source = Source::new([0u8; 32]);
118    let mut source_xe: Source = Source::new([0u8; 32]);
119    let mut source_xa: Source = Source::new([0u8; 32]);
120    let mut source_xu: Source = Source::new([0u8; 32]);
121
122    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
123    sk.fill_ternary_prob(0.5, &mut source_xs);
124    let sk_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk);
125
126    let mut pk: GLWEPublicKey<Vec<u8>, FFT64> = GLWEPublicKey::alloc(&module, basek, k_pk, rank);
127    pk.generate_from_sk(&module, &sk_dft, &mut source_xa, &mut source_xe, sigma);
128
129    let mut scratch: ScratchOwned = ScratchOwned::new(
130        GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct.k())
131            | GLWECiphertext::decrypt_scratch_space(&module, basek, ct.k())
132            | GLWECiphertext::encrypt_pk_scratch_space(&module, basek, pk.k()),
133    );
134
135    pt_want
136        .data
137        .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
138
139    ct.encrypt_pk(
140        &module,
141        &pt_want,
142        &pk,
143        &mut source_xu,
144        &mut source_xe,
145        sigma,
146        scratch.borrow(),
147    );
148
149    ct.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
150
151    pt_want.sub_inplace_ab(&module, &pt_have);
152
153    let noise_have: f64 = pt_want.data.std(0, basek).log2();
154    let noise_want: f64 = ((((rank as f64) + 1.0) * module.n() as f64 * 0.5 * sigma * sigma).sqrt()).log2() - (k_ct as f64);
155
156    assert!(
157        (noise_have - noise_want).abs() < 0.2,
158        "{} {}",
159        noise_have,
160        noise_want
161    );
162}