1use backend::{
  2    Backend, FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps, ScalarZnxToRef, Scratch,
  3    ZnxView, ZnxViewMut,
  4};
  5use sampling::source::Source;
  6
  7use crate::{Distribution, FourierGLWESecret, GGSWCiphertext, Infos, LWESecret};
  8
  9pub struct BlindRotationKeyCGGI<D, B: Backend> {
 10    pub(crate) data: Vec<GGSWCiphertext<D, B>>,
 11    pub(crate) dist: Distribution,
 12    pub(crate) x_pow_a: Option<Vec<ScalarZnxDft<Vec<u8>, B>>>,
 13}
 14
 15// pub struct BlindRotationKeyFHEW<B: Backend> {
 16//    pub(crate) data: Vec<GGSWCiphertext<Vec<u8>, B>>,
 17//    pub(crate) auto: Vec<GLWEAutomorphismKey<Vec<u8>, B>>,
 18//}
 19
 20impl BlindRotationKeyCGGI<Vec<u8>, FFT64> {
 21    pub fn allocate(module: &Module<FFT64>, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
 22        let mut data: Vec<GGSWCiphertext<Vec<u8>, FFT64>> = Vec::with_capacity(n_lwe);
 23        (0..n_lwe).for_each(|_| data.push(GGSWCiphertext::alloc(module, basek, k, rows, 1, rank)));
 24        Self {
 25            data,
 26            dist: Distribution::NONE,
 27            x_pow_a: None::<Vec<ScalarZnxDft<Vec<u8>, FFT64>>>,
 28        }
 29    }
 30
 31    pub fn generate_from_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
 32        GGSWCiphertext::encrypt_sk_scratch_space(module, basek, k, rank)
 33    }
 34}
 35
 36impl<D: AsRef<[u8]>> BlindRotationKeyCGGI<D, FFT64> {
 37    #[allow(dead_code)]
 38    pub(crate) fn n(&self) -> usize {
 39        self.data[0].n()
 40    }
 41
 42    #[allow(dead_code)]
 43    pub(crate) fn rows(&self) -> usize {
 44        self.data[0].rows()
 45    }
 46
 47    #[allow(dead_code)]
 48    pub(crate) fn k(&self) -> usize {
 49        self.data[0].k()
 50    }
 51
 52    #[allow(dead_code)]
 53    pub(crate) fn size(&self) -> usize {
 54        self.data[0].size()
 55    }
 56
 57    #[allow(dead_code)]
 58    pub(crate) fn rank(&self) -> usize {
 59        self.data[0].rank()
 60    }
 61
 62    pub(crate) fn basek(&self) -> usize {
 63        self.data[0].basek()
 64    }
 65
 66    pub(crate) fn block_size(&self) -> usize {
 67        match self.dist {
 68            Distribution::BinaryBlock(value) => value,
 69            _ => 1,
 70        }
 71    }
 72}
 73
 74impl<D: AsRef<[u8]> + AsMut<[u8]>> BlindRotationKeyCGGI<D, FFT64> {
 75    pub fn generate_from_sk<DataSkGLWE, DataSkLWE>(
 76        &mut self,
 77        module: &Module<FFT64>,
 78        sk_glwe: &FourierGLWESecret<DataSkGLWE, FFT64>,
 79        sk_lwe: &LWESecret<DataSkLWE>,
 80        source_xa: &mut Source,
 81        source_xe: &mut Source,
 82        sigma: f64,
 83        scratch: &mut Scratch,
 84    ) where
 85        DataSkGLWE: AsRef<[u8]>,
 86        DataSkLWE: AsRef<[u8]>,
 87    {
 88        #[cfg(debug_assertions)]
 89        {
 90            assert_eq!(self.data.len(), sk_lwe.n());
 91            assert_eq!(sk_glwe.n(), module.n());
 92            assert_eq!(sk_glwe.rank(), self.data[0].rank());
 93            match sk_lwe.dist {
 94                Distribution::BinaryBlock(_)
 95                | Distribution::BinaryFixed(_)
 96                | Distribution::BinaryProb(_)
 97                | Distribution::ZERO => {}
 98                _ => panic!(
 99                    "invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb (or ZERO for debugging)"
100                ),
101            }
102        }
103
104        self.dist = sk_lwe.dist;
105
106        let mut pt: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
107        let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data.to_ref();
108
109        self.data.iter_mut().enumerate().for_each(|(i, ggsw)| {
110            pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
111            ggsw.encrypt_sk(module, &pt, sk_glwe, source_xa, source_xe, sigma, scratch);
112        });
113
114        match sk_lwe.dist {
115            Distribution::BinaryBlock(_) => {
116                let mut x_pow_a: Vec<ScalarZnxDft<Vec<u8>, FFT64>> = Vec::with_capacity(module.n() << 1);
117                let mut buf: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
118                (0..module.n() << 1).for_each(|i| {
119                    let mut res: ScalarZnxDft<Vec<u8>, FFT64> = module.new_scalar_znx_dft(1);
120                    set_xai_plus_y(module, i, 0, &mut res, &mut buf);
121                    x_pow_a.push(res);
122                });
123                self.x_pow_a = Some(x_pow_a);
124            }
125            _ => {}
126        }
127    }
128}
129
130pub fn set_xai_plus_y<A, B>(module: &Module<FFT64>, ai: usize, y: i64, res: &mut ScalarZnxDft<A, FFT64>, buf: &mut ScalarZnx<B>)
131where
132    A: AsRef<[u8]> + AsMut<[u8]>,
133    B: AsRef<[u8]> + AsMut<[u8]>,
134{
135    let n: usize = module.n();
136
137    {
138        let raw: &mut [i64] = buf.at_mut(0, 0);
139        if ai < n {
140            raw[ai] = 1;
141        } else {
142            raw[(ai - n) & (n - 1)] = -1;
143        }
144        raw[0] += y;
145    }
146
147    module.svp_prepare(res, 0, buf, 0);
148
149    {
150        let raw: &mut [i64] = buf.at_mut(0, 0);
151
152        if ai < n {
153            raw[ai] = 0;
154        } else {
155            raw[(ai - n) & (n - 1)] = 0;
156        }
157        raw[0] = 0;
158    }
159}