1use backend::{Encoding, FFT64, Module, ScratchOwned, ZnxView};
  2use sampling::source::Source;
  3
  4use crate::{
  5    FourierGLWESecret, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos, LWECiphertext, LWESecret,
  6    blind_rotation::{
  7        ccgi::{cggi_blind_rotate, cggi_blind_rotate_scratch_space, negate_and_mod_switch_2n},
  8        key::BlindRotationKeyCGGI,
  9        lut::LookUpTable,
 10    },
 11    lwe::{LWEPlaintext, ciphertext::LWECiphertextToRef},
 12};
 13
 14#[test]
 15fn standard() {
 16    blind_rotatio_test(224, 1, 1);
 17}
 18
 19#[test]
 20fn block_binary() {
 21    blind_rotatio_test(224, 7, 1);
 22}
 23
 24#[test]
 25fn block_binary_extended() {
 26    blind_rotatio_test(224, 7, 2);
 27}
 28
 29fn blind_rotatio_test(n_lwe: usize, block_size: usize, extension_factor: usize) {
 30    let module: Module<FFT64> = Module::<FFT64>::new(512);
 31    let basek: usize = 19;
 32
 33    let k_lwe: usize = 24;
 34    let k_brk: usize = 3 * basek;
 35    let rows_brk: usize = 2; // Ensures first limb is noise-free.
 36    let k_lut: usize = 1 * basek;
 37    let k_res: usize = 2 * basek;
 38    let rank: usize = 1;
 39
 40    let message_modulus: usize = 1 << 4;
 41
 42    let mut source_xs: Source = Source::new([1u8; 32]);
 43    let mut source_xe: Source = Source::new([1u8; 32]);
 44    let mut source_xa: Source = Source::new([1u8; 32]);
 45
 46    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
 47    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 48    let sk_glwe_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk_glwe);
 49
 50    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
 51    sk_lwe.fill_binary_block(block_size, &mut source_xs);
 52
 53    let mut scratch: ScratchOwned = ScratchOwned::new(BlindRotationKeyCGGI::generate_from_sk_scratch_space(
 54        &module, basek, k_brk, rank,
 55    ));
 56
 57    let mut scratch_br: ScratchOwned = ScratchOwned::new(cggi_blind_rotate_scratch_space(
 58        &module,
 59        block_size,
 60        extension_factor,
 61        basek,
 62        k_res,
 63        k_brk,
 64        rows_brk,
 65        rank,
 66    ));
 67
 68    let mut brk: BlindRotationKeyCGGI<FFT64> = BlindRotationKeyCGGI::allocate(&module, n_lwe, basek, k_brk, rows_brk, rank);
 69
 70    brk.generate_from_sk(
 71        &module,
 72        &sk_glwe_dft,
 73        &sk_lwe,
 74        &mut source_xa,
 75        &mut source_xe,
 76        3.2,
 77        scratch.borrow(),
 78    );
 79
 80    let mut lwe: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe);
 81
 82    let mut pt_lwe: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe);
 83
 84    let x: i64 = 2;
 85    let bits: usize = 8;
 86
 87    pt_lwe.data.encode_coeff_i64(0, basek, bits, 0, x, bits);
 88
 89    // println!("{}", pt_lwe.data);
 90
 91    lwe.encrypt_sk(&pt_lwe, &sk_lwe, &mut source_xa, &mut source_xe, 3.2);
 92
 93    lwe.decrypt(&mut pt_lwe, &sk_lwe);
 94
 95    // println!("{}", pt_lwe.data);
 96
 97    let mut f: Vec<i64> = vec![0i64; message_modulus];
 98    f.iter_mut()
 99        .enumerate()
100        .for_each(|(i, x)| *x = 2 * (i as i64) + 1);
101
102    let mut lut: LookUpTable = LookUpTable::alloc(&module, basek, k_lut, extension_factor);
103    lut.set(&module, &f, message_modulus);
104
105    let mut res: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_res, rank);
106
107    cggi_blind_rotate(&module, &mut res, &lwe, &lut, &brk, scratch_br.borrow());
108
109    println!("out_mut.data: {}", res.data);
110
111    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_res);
112
113    res.decrypt(&module, &mut pt_have, &sk_glwe_dft, scratch.borrow());
114
115    println!("pt_have: {}", pt_have.data);
116
117    let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
118
119    negate_and_mod_switch_2n(2 * lut.domain_size(), &mut lwe_2n, &lwe.to_ref());
120
121    let pt_want: i64 = (lwe_2n[0]
122        + lwe_2n[1..]
123            .iter()
124            .zip(sk_lwe.data.at(0, 0))
125            .map(|(x, y)| x * y)
126            .sum::<i64>())
127        % (2 * lut.domain_size()) as i64;
128
129    // println!("pt_want: {}", pt_want);
130
131    lut.rotate(pt_want);
132
133    // lut.data.iter().for_each(|d| {
134    //     println!("{}", d);
135    // });
136
137    // First limb should be exactly equal (test are parameterized such that the noise does not reach
138    // the first limb)
139    assert_eq!(pt_have.data.at(0, 0), lut.data[0].at(0, 0));
140
141    // Then checks the noise
142    // module.vec_znx_sub_ab_inplace(&mut lut.data[0], 0, &pt_have.data, 0);
143    // let noise: f64 = lut.data[0].std(0, basek);
144    // println!("noise: {}", noise);
145    // assert!(noise < 1e-3);
146}