1use backend::{
  2    Backend, FFT64, MatZnxDft, MatZnxDftAlloc, MatZnxDftOps, Module, ScalarZnx, Scratch, VecZnxAlloc, VecZnxDftAlloc, VecZnxOps,
  3    ZnxInfos, ZnxZero,
  4};
  5use sampling::source::Source;
  6
  7use crate::{GLWECiphertext, GLWECiphertextFourier, GLWESecret, GetRow, Infos, ScratchCore, SetRow, div_ceil};
  8
  9pub struct GGLWECiphertext<C, B: Backend> {
 10    pub(crate) data: MatZnxDft<C, B>,
 11    pub(crate) basek: usize,
 12    pub(crate) k: usize,
 13    pub(crate) digits: usize,
 14}
 15
 16impl<B: Backend> GGLWECiphertext<Vec<u8>, B> {
 17    pub fn alloc(
 18        module: &Module<B>,
 19        basek: usize,
 20        k: usize,
 21        rows: usize,
 22        digits: usize,
 23        rank_in: usize,
 24        rank_out: usize,
 25    ) -> Self {
 26        let size: usize = div_ceil(k, basek);
 27        debug_assert!(
 28            size > digits,
 29            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
 30            size,
 31            digits
 32        );
 33
 34        assert!(
 35            rows * digits <= size,
 36            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
 37            rows,
 38            digits,
 39            size
 40        );
 41
 42        Self {
 43            data: module.new_mat_znx_dft(rows, rank_in, rank_out + 1, size),
 44            basek: basek,
 45            k,
 46            digits,
 47        }
 48    }
 49
 50    pub fn bytes_of(
 51        module: &Module<FFT64>,
 52        basek: usize,
 53        k: usize,
 54        rows: usize,
 55        digits: usize,
 56        rank_in: usize,
 57        rank_out: usize,
 58    ) -> usize {
 59        let size: usize = div_ceil(k, basek);
 60        debug_assert!(
 61            size > digits,
 62            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
 63            size,
 64            digits
 65        );
 66
 67        assert!(
 68            rows * digits <= size,
 69            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
 70            rows,
 71            digits,
 72            size
 73        );
 74
 75        module.bytes_of_mat_znx_dft(rows, rank_in, rank_out + 1, size)
 76    }
 77}
 78
 79impl<T, B: Backend> Infos for GGLWECiphertext<T, B> {
 80    type Inner = MatZnxDft<T, B>;
 81
 82    fn inner(&self) -> &Self::Inner {
 83        &self.data
 84    }
 85
 86    fn basek(&self) -> usize {
 87        self.basek
 88    }
 89
 90    fn k(&self) -> usize {
 91        self.k
 92    }
 93}
 94
 95impl<T, B: Backend> GGLWECiphertext<T, B> {
 96    pub fn rank(&self) -> usize {
 97        self.data.cols_out() - 1
 98    }
 99
100    pub fn digits(&self) -> usize {
101        self.digits
102    }
103
104    pub fn rank_in(&self) -> usize {
105        self.data.cols_in()
106    }
107
108    pub fn rank_out(&self) -> usize {
109        self.data.cols_out() - 1
110    }
111}
112
113impl GGLWECiphertext<Vec<u8>, FFT64> {
114    pub fn generate_from_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
115        let size = div_ceil(k, basek);
116        GLWECiphertext::encrypt_sk_scratch_space(module, basek, k)
117            + module.bytes_of_vec_znx(rank + 1, size)
118            + module.bytes_of_vec_znx(1, size)
119            + module.bytes_of_vec_znx_dft(rank + 1, size)
120    }
121
122    pub fn generate_from_pk_scratch_space(_module: &Module<FFT64>, _basek: usize, _k: usize, _rank: usize) -> usize {
123        unimplemented!()
124    }
125}
126
127impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GGLWECiphertext<DataSelf, FFT64> {
128    pub fn encrypt_sk<DataPt: AsRef<[u8]>, DataSk: AsRef<[u8]>>(
129        &mut self,
130        module: &Module<FFT64>,
131        pt: &ScalarZnx<DataPt>,
132        sk: &GLWESecret<DataSk, FFT64>,
133        source_xa: &mut Source,
134        source_xe: &mut Source,
135        sigma: f64,
136        scratch: &mut Scratch,
137    ) {
138        #[cfg(debug_assertions)]
139        {
140            assert_eq!(self.rank_in(), pt.cols());
141            assert_eq!(self.rank_out(), sk.rank());
142            assert_eq!(self.n(), module.n());
143            assert_eq!(sk.n(), module.n());
144            assert_eq!(pt.n(), module.n());
145            assert!(
146                scratch.available()
147                    >= GGLWECiphertext::generate_from_sk_scratch_space(module, self.basek(), self.k(), self.rank()),
148                "scratch.available: {} < GGLWECiphertext::generate_from_sk_scratch_space(module, self.rank()={}, \
149                 self.size()={}): {}",
150                scratch.available(),
151                self.rank(),
152                self.size(),
153                GGLWECiphertext::generate_from_sk_scratch_space(module, self.basek(), self.k(), self.rank())
154            );
155            assert!(
156                self.rows() * self.digits() * self.basek() <= self.k(),
157                "self.rows() : {} * self.digits() : {} * self.basek() : {} = {} >= self.k() = {}",
158                self.rows(),
159                self.digits(),
160                self.basek(),
161                self.rows() * self.digits() * self.basek(),
162                self.k()
163            );
164        }
165
166        let rows: usize = self.rows();
167        let digits: usize = self.digits();
168        let basek: usize = self.basek();
169        let k: usize = self.k();
170        let rank_in: usize = self.rank_in();
171        let rank_out: usize = self.rank_out();
172
173        let (mut tmp_pt, scrach_1) = scratch.tmp_glwe_pt(module, basek, k);
174        let (mut tmp_ct, scrach_2) = scrach_1.tmp_glwe_ct(module, basek, k, rank_out);
175        let (mut tmp_ct_dft, scratch_3) = scrach_2.tmp_glwe_fourier(module, basek, k, rank_out);
176
177        // For each input column (i.e. rank) produces a GGLWE ciphertext of rank_out+1 columns
178        //
179        // Example for ksk rank 2 to rank 3:
180        //
181        // (-(a0*s0 + a1*s1 + a2*s2) + s0', a0, a1, a2)
182        // (-(b0*s0 + b1*s1 + b2*s2) + s0', b0, b1, b2)
183        //
184        // Example ksk rank 2 to rank 1
185        //
186        // (-(a*s) + s0, a)
187        // (-(b*s) + s1, b)
188        (0..rank_in).for_each(|col_i| {
189            (0..rows).for_each(|row_i| {
190                // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
191                tmp_pt.data.zero(); // zeroes for next iteration
192                module.vec_znx_add_scalar_inplace(
193                    &mut tmp_pt.data,
194                    0,
195                    (digits - 1) + row_i * digits,
196                    pt,
197                    col_i,
198                );
199                module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scratch_3);
200
201                // rlwe encrypt of vec_znx_pt into vec_znx_ct
202                tmp_ct.encrypt_sk(module, &tmp_pt, sk, source_xa, source_xe, sigma, scratch_3);
203
204                // Switch vec_znx_ct into DFT domain
205                tmp_ct.dft(module, &mut tmp_ct_dft);
206
207                // Stores vec_znx_dft_ct into thw i-th row of the MatZnxDft
208                self.set_row(module, row_i, col_i, &tmp_ct_dft);
209            });
210        });
211    }
212}
213
214impl<C: AsRef<[u8]>> GetRow<FFT64> for GGLWECiphertext<C, FFT64> {
215    fn get_row<R: AsMut<[u8]> + AsRef<[u8]>>(
216        &self,
217        module: &Module<FFT64>,
218        row_i: usize,
219        col_j: usize,
220        res: &mut GLWECiphertextFourier<R, FFT64>,
221    ) {
222        module.vmp_extract_row(&mut res.data, &self.data, row_i, col_j);
223    }
224}
225
226impl<C: AsMut<[u8]> + AsRef<[u8]>> SetRow<FFT64> for GGLWECiphertext<C, FFT64> {
227    fn set_row<R: AsRef<[u8]>>(
228        &mut self,
229        module: &Module<FFT64>,
230        row_i: usize,
231        col_j: usize,
232        a: &GLWECiphertextFourier<R, FFT64>,
233    ) {
234        module.vmp_prepare_row(&mut self.data, row_i, col_j, &a.data);
235    }
236}