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