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