1use std::time::Instant;
  2
  3use backend::{
  4    FFT64, MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnxDftOps, Scratch, VecZnxDftOps, VecZnxOps, ZnxView, ZnxViewMut,
  5    ZnxZero,
  6};
  7use itertools::izip;
  8
  9use crate::{
 10    GGSWCiphertext, GLWECiphertext, GLWECiphertextToMut, Infos, LWECiphertext, ScratchCore,
 11    blind_rotation::{key::BlindRotationKeyCGGI, lut::LookUpTable},
 12    lwe::ciphertext::LWECiphertextToRef,
 13};
 14
 15pub fn cggi_blind_rotate_scratch_space(
 16    module: &Module<FFT64>,
 17    basek: usize,
 18    k_lut: usize,
 19    k_brk: usize,
 20    rows: usize,
 21    rank: usize,
 22) -> usize {
 23    let size = k_brk.div_ceil(basek);
 24    GGSWCiphertext::<Vec<u8>, FFT64>::bytes_of(module, basek, k_brk, rows, 1, rank)
 25        + (module.mat_znx_dft_mul_x_pow_minus_one_scratch_space(size, rank + 1)
 26            | GLWECiphertext::external_product_inplace_scratch_space(module, basek, k_lut, k_brk, 1, rank))
 27}
 28
 29pub fn cggi_blind_rotate<DataRes, DataIn>(
 30    module: &Module<FFT64>,
 31    res: &mut GLWECiphertext<DataRes>,
 32    lwe: &LWECiphertext<DataIn>,
 33    lut: &LookUpTable,
 34    brk: &BlindRotationKeyCGGI<FFT64>,
 35    scratch: &mut Scratch,
 36) where
 37    DataRes: AsRef<[u8]> + AsMut<[u8]>,
 38    DataIn: AsRef<[u8]>,
 39{
 40    let basek = res.basek();
 41
 42    let mut lwe_2n: Vec<i64> = vec![0i64; lwe.n() + 1]; // TODO: from scratch space
 43    let mut out_mut: GLWECiphertext<&mut [u8]> = res.to_mut();
 44    let lwe_ref: LWECiphertext<&[u8]> = lwe.to_ref();
 45
 46    let cols: usize = out_mut.rank() + 1;
 47
 48    mod_switch_2n(module, &mut lwe_2n, &lwe_ref);
 49
 50    let a: &[i64] = &lwe_2n[1..];
 51    let b: i64 = lwe_2n[0];
 52
 53    out_mut.data.zero();
 54
 55    // Initialize out to X^{b} * LUT(X)
 56    module.vec_znx_rotate(b, &mut out_mut.data, 0, &lut.data[0], 0);
 57
 58    let block_size: usize = brk.block_size();
 59
 60    // ACC + [sum DFT(X^ai -1) * (DFT(ACC) x BRKi)]
 61
 62    let (mut acc_dft, scratch1) = scratch.tmp_glwe_fourier(module, brk.basek(), out_mut.k(), out_mut.rank());
 63    let (mut acc_add_dft, scratch2) = scratch1.tmp_glwe_fourier(module, brk.basek(), out_mut.k(), out_mut.rank());
 64    let (mut vmp_res, scratch3) = scratch2.tmp_glwe_fourier(module, basek, out_mut.k(), out_mut.rank());
 65    let (mut xai_minus_one, scratch4) = scratch3.tmp_scalar_znx(module, 1);
 66    let (mut xai_minus_one_dft, scratch5) = scratch4.tmp_scalar_znx_dft(module, 1);
 67
 68    let start: Instant = Instant::now();
 69    izip!(
 70        a.chunks_exact(block_size),
 71        brk.data.chunks_exact(block_size)
 72    )
 73    .for_each(|(ai, ski)| {
 74        out_mut.dft(module, &mut acc_dft);
 75        acc_add_dft.data.zero();
 76
 77        izip!(ai.iter(), ski.iter()).for_each(|(aii, skii)| {
 78            // vmp_res = DFT(acc) * BRK[i]
 79            module.vmp_apply(&mut vmp_res.data, &acc_dft.data, &skii.data, scratch5);
 80
 81            // DFT(X^ai -1)
 82            xai_minus_one.zero();
 83            xai_minus_one.at_mut(0, 0)[0] = 1;
 84            module.vec_znx_rotate_inplace(*aii, &mut xai_minus_one, 0);
 85            xai_minus_one.at_mut(0, 0)[0] -= 1;
 86            module.svp_prepare(&mut xai_minus_one_dft, 0, &xai_minus_one, 0);
 87
 88            // DFT(X^ai -1) * (DFT(acc) * BRK[i])
 89            (0..cols).for_each(|i| {
 90                module.svp_apply_inplace(&mut vmp_res.data, i, &xai_minus_one_dft, 0);
 91                module.vec_znx_dft_add_inplace(&mut acc_add_dft.data, i, &vmp_res.data, i);
 92            });
 93        });
 94
 95        (0..cols).for_each(|i| {
 96            module.vec_znx_dft_add_inplace(&mut acc_dft.data, i, &acc_add_dft.data, i);
 97        });
 98
 99        acc_dft.idft(module, &mut out_mut, scratch5);
100    });
101    let duration: std::time::Duration = start.elapsed();
102    println!("external products: {} us", duration.as_micros());
103}
104
105pub(crate) fn mod_switch_2n(module: &Module<FFT64>, res: &mut [i64], lwe: &LWECiphertext<&[u8]>) {
106    let basek: usize = lwe.basek();
107
108    let log2n: usize = module.log_n() + 1;
109
110    res.copy_from_slice(&lwe.data.at(0, 0));
111
112    if basek > log2n {
113        let diff: usize = basek - log2n;
114        res.iter_mut().for_each(|x| {
115            *x = div_signed_by_pow2(x, diff);
116        })
117    } else {
118        let rem: usize = basek - (log2n % basek);
119        let size: usize = log2n.div_ceil(basek);
120        (1..size).for_each(|i| {
121            if i == size - 1 && rem != basek {
122                let k_rem: usize = basek - rem;
123                izip!(lwe.data.at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
124                    *y = (*y << k_rem) + (x >> rem);
125                });
126            } else {
127                izip!(lwe.data.at(0, i).iter(), res.iter_mut()).for_each(|(x, y)| {
128                    *y = (*y << basek) + x;
129                });
130            }
131        })
132    }
133}
134
135#[inline(always)]
136fn div_signed_by_pow2(x: &i64, k: usize) -> i64 {
137    let bias: i64 = (1 << k) - 1;
138    (x + ((x >> 63) & bias)) >> k
139}