1use backend::{Backend, FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxToRef, Scratch, ZnxView, ZnxViewMut};
 2use sampling::source::Source;
 3
 4use crate::{Distribution, FourierGLWESecret, GGSWCiphertext, Infos, LWESecret};
 5
 6pub struct BlindRotationKeyCGGI<B: Backend> {
 7    pub(crate) data: Vec<GGSWCiphertext<Vec<u8>, B>>,
 8    pub(crate) dist: Distribution,
 9}
10
11// pub struct BlindRotationKeyFHEW<B: Backend> {
12//    pub(crate) data: Vec<GGSWCiphertext<Vec<u8>, B>>,
13//    pub(crate) auto: Vec<GLWEAutomorphismKey<Vec<u8>, B>>,
14//}
15
16impl BlindRotationKeyCGGI<FFT64> {
17    pub fn allocate(module: &Module<FFT64>, n_lwe: usize, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
18        let mut data: Vec<GGSWCiphertext<Vec<u8>, FFT64>> = Vec::with_capacity(n_lwe);
19        (0..n_lwe).for_each(|_| data.push(GGSWCiphertext::alloc(module, basek, k, rows, 1, rank)));
20        Self {
21            data,
22            dist: Distribution::NONE,
23        }
24    }
25
26    pub fn generate_from_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
27        GGSWCiphertext::encrypt_sk_scratch_space(module, basek, k, rank)
28    }
29
30    pub fn generate_from_sk<DataSkGLWE, DataSkLWE>(
31        &mut self,
32        module: &Module<FFT64>,
33        sk_glwe: &FourierGLWESecret<DataSkGLWE, FFT64>,
34        sk_lwe: &LWESecret<DataSkLWE>,
35        source_xa: &mut Source,
36        source_xe: &mut Source,
37        sigma: f64,
38        scratch: &mut Scratch,
39    ) where
40        DataSkGLWE: AsRef<[u8]>,
41        DataSkLWE: AsRef<[u8]>,
42    {
43        #[cfg(debug_assertions)]
44        {
45            assert_eq!(self.data.len(), sk_lwe.n());
46            assert_eq!(sk_glwe.n(), module.n());
47            assert_eq!(sk_glwe.rank(), self.data[0].rank());
48            match sk_lwe.dist {
49                Distribution::BinaryBlock(_) | Distribution::BinaryFixed(_) | Distribution::BinaryProb(_) => {}
50                _ => panic!("invalid GLWESecret distribution: must be BinaryBlock, BinaryFixed or BinaryProb"),
51            }
52        }
53
54        self.dist = sk_lwe.dist;
55
56        let mut pt: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
57        let sk_ref: ScalarZnx<&[u8]> = sk_lwe.data.to_ref();
58
59        self.data.iter_mut().enumerate().for_each(|(i, ggsw)| {
60            pt.at_mut(0, 0)[0] = sk_ref.at(0, 0)[i];
61            ggsw.encrypt_sk(module, &pt, sk_glwe, source_xa, source_xe, sigma, scratch);
62        })
63    }
64
65    pub(crate) fn block_size(&self) -> usize {
66        match self.dist {
67            Distribution::BinaryBlock(value) => value,
68            _ => 1,
69        }
70    }
71
72    #[allow(dead_code)]
73    pub(crate) fn rows(&self) -> usize {
74        self.data[0].rows()
75    }
76
77    #[allow(dead_code)]
78    pub(crate) fn k(&self) -> usize {
79        self.data[0].k()
80    }
81
82    #[allow(dead_code)]
83    pub(crate) fn rank(&self) -> usize {
84        self.data[0].rank()
85    }
86
87    pub(crate) fn basek(&self) -> usize {
88        self.data[0].basek()
89    }
90}