1use backend::{Decoding, Encoding, FFT64, Module, ScratchOwned, Stats, VecZnxOps, ZnxZero};
  2use itertools::izip;
  3use sampling::source::Source;
  4
  5use crate::{FourierGLWECiphertext, GLWECiphertext, GLWEPlaintext, GLWEPublicKey, GLWESecret, Infos};
  6
  7#[test]
  8fn encrypt_sk() {
  9    let log_n: usize = 8;
 10    (1..4).for_each(|rank| {
 11        println!("test encrypt_sk rank: {}", rank);
 12        test_encrypt_sk(log_n, 8, 54, 30, 3.2, rank);
 13    });
 14}
 15
 16#[test]
 17fn encrypt_zero_sk() {
 18    let log_n: usize = 8;
 19    (1..4).for_each(|rank| {
 20        println!("test encrypt_zero_sk rank: {}", rank);
 21        test_encrypt_zero_sk(log_n, 8, 64, 3.2, rank);
 22    });
 23}
 24
 25#[test]
 26fn encrypt_pk() {
 27    let log_n: usize = 8;
 28    (1..4).for_each(|rank| {
 29        println!("test encrypt_pk rank: {}", rank);
 30        test_encrypt_pk(log_n, 8, 64, 64, 3.2, rank)
 31    });
 32}
 33
 34fn test_encrypt_sk(log_n: usize, basek: usize, k_ct: usize, k_pt: usize, sigma: f64, rank: usize) {
 35    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 36
 37    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct, rank);
 38    let mut pt: 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>, FFT64> = GLWESecret::alloc(&module, rank);
 50    sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
 51
 52    let mut data_want: Vec<i64> = vec![0i64; module.n()];
 53
 54    data_want
 55        .iter_mut()
 56        .for_each(|x| *x = source_xa.next_i64() & 0xFF);
 57
 58    pt.data.encode_vec_i64(0, basek, k_pt, &data_want, 10);
 59
 60    ct.encrypt_sk(
 61        &module,
 62        &pt,
 63        &sk,
 64        &mut source_xa,
 65        &mut source_xe,
 66        sigma,
 67        scratch.borrow(),
 68    );
 69
 70    pt.data.zero();
 71
 72    ct.decrypt(&module, &mut pt, &sk, scratch.borrow());
 73
 74    let mut data_have: Vec<i64> = vec![0i64; module.n()];
 75
 76    pt.data
 77        .decode_vec_i64(0, basek, pt.size() * basek, &mut data_have);
 78
 79    // TODO: properly assert the decryption noise through std(dec(ct) - pt)
 80    let scale: f64 = (1 << (pt.size() * basek - k_pt)) as f64;
 81    izip!(data_want.iter(), data_have.iter()).for_each(|(a, b)| {
 82        let b_scaled = (*b as f64) / scale;
 83        assert!(
 84            (*a as f64 - b_scaled).abs() < 0.1,
 85            "{} {}",
 86            *a as f64,
 87            b_scaled
 88        )
 89    });
 90}
 91
 92fn test_encrypt_zero_sk(log_n: usize, basek: usize, k_ct: usize, sigma: f64, rank: usize) {
 93    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 94
 95    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
 96
 97    let mut source_xs: Source = Source::new([0u8; 32]);
 98    let mut source_xe: Source = Source::new([1u8; 32]);
 99    let mut source_xa: Source = Source::new([0u8; 32]);
100
101    let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
102    sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
103
104    let mut ct_dft: FourierGLWECiphertext<Vec<u8>, FFT64> = FourierGLWECiphertext::alloc(&module, basek, k_ct, rank);
105
106    let mut scratch: ScratchOwned = ScratchOwned::new(
107        FourierGLWECiphertext::decrypt_scratch_space(&module, basek, k_ct)
108            | FourierGLWECiphertext::encrypt_sk_scratch_space(&module, basek, k_ct, rank),
109    );
110
111    ct_dft.encrypt_zero_sk(
112        &module,
113        &sk,
114        &mut source_xa,
115        &mut source_xe,
116        sigma,
117        scratch.borrow(),
118    );
119    ct_dft.decrypt(&module, &mut pt, &sk, scratch.borrow());
120
121    assert!((sigma - pt.data.std(0, basek) * (k_ct as f64).exp2()) <= 0.2);
122}
123
124fn test_encrypt_pk(log_n: usize, basek: usize, k_ct: usize, k_pk: usize, sigma: f64, rank: usize) {
125    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
126
127    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct, rank);
128    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
129
130    let mut source_xs: Source = Source::new([0u8; 32]);
131    let mut source_xe: Source = Source::new([0u8; 32]);
132    let mut source_xa: Source = Source::new([0u8; 32]);
133    let mut source_xu: Source = Source::new([0u8; 32]);
134
135    let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
136    sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
137
138    let mut pk: GLWEPublicKey<Vec<u8>, FFT64> = GLWEPublicKey::alloc(&module, basek, k_pk, rank);
139    pk.generate_from_sk(&module, &sk, &mut source_xa, &mut source_xe, sigma);
140
141    let mut scratch: ScratchOwned = ScratchOwned::new(
142        GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct.k())
143            | GLWECiphertext::decrypt_scratch_space(&module, basek, ct.k())
144            | GLWECiphertext::encrypt_pk_scratch_space(&module, basek, pk.k()),
145    );
146
147    let mut data_want: Vec<i64> = vec![0i64; module.n()];
148
149    data_want
150        .iter_mut()
151        .for_each(|x| *x = source_xa.next_i64() & 0);
152
153    pt_want.data.encode_vec_i64(0, basek, k_ct, &data_want, 10);
154
155    ct.encrypt_pk(
156        &module,
157        &pt_want,
158        &pk,
159        &mut source_xu,
160        &mut source_xe,
161        sigma,
162        scratch.borrow(),
163    );
164
165    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
166
167    ct.decrypt(&module, &mut pt_have, &sk, scratch.borrow());
168
169    module.vec_znx_sub_ab_inplace(&mut pt_want.data, 0, &pt_have.data, 0);
170
171    let noise_have: f64 = pt_want.data.std(0, basek).log2();
172    let noise_want: f64 = ((((rank as f64) + 1.0) * module.n() as f64 * 0.5 * sigma * sigma).sqrt()).log2() - (k_ct as f64);
173
174    assert!(
175        (noise_have - noise_want).abs() < 0.2,
176        "{} {}",
177        noise_have,
178        noise_want
179    );
180}