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