1use backend::hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPMatPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat},
  4};
  5
  6use crate::{AutomorphismKey, GGLWECiphertext, GGLWEExecLayoutFamily, GLWESwitchingKey, GLWETensorKey, Infos};
  7
  8#[derive(PartialEq, Eq)]
  9pub struct GGLWECiphertextExec<D: Data, B: Backend> {
 10    pub(crate) data: VmpPMat<D, B>,
 11    pub(crate) basek: usize,
 12    pub(crate) k: usize,
 13    pub(crate) digits: usize,
 14}
 15
 16impl<B: Backend> GGLWECiphertextExec<Vec<u8>, B> {
 17    pub fn alloc(
 18        module: &Module<B>,
 19        n: usize,
 20        basek: usize,
 21        k: usize,
 22        rows: usize,
 23        digits: usize,
 24        rank_in: usize,
 25        rank_out: usize,
 26    ) -> Self
 27    where
 28        Module<B>: GGLWEExecLayoutFamily<B>,
 29    {
 30        let size: usize = k.div_ceil(basek);
 31        debug_assert!(
 32            size > digits,
 33            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
 34            size,
 35            digits
 36        );
 37
 38        assert!(
 39            rows * digits <= size,
 40            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
 41            rows,
 42            digits,
 43            size
 44        );
 45
 46        Self {
 47            data: module.vmp_pmat_alloc(n, rows, rank_in, rank_out + 1, size),
 48            basek: basek,
 49            k,
 50            digits,
 51        }
 52    }
 53
 54    pub fn bytes_of(
 55        module: &Module<B>,
 56        n: usize,
 57        basek: usize,
 58        k: usize,
 59        rows: usize,
 60        digits: usize,
 61        rank_in: usize,
 62        rank_out: usize,
 63    ) -> usize
 64    where
 65        Module<B>: GGLWEExecLayoutFamily<B>,
 66    {
 67        let size: usize = k.div_ceil(basek);
 68        debug_assert!(
 69            size > digits,
 70            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
 71            size,
 72            digits
 73        );
 74
 75        assert!(
 76            rows * digits <= size,
 77            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
 78            rows,
 79            digits,
 80            size
 81        );
 82
 83        module.vmp_pmat_alloc_bytes(n, rows, rank_in, rank_out + 1, rows)
 84    }
 85}
 86
 87impl<D: Data, B: Backend> Infos for GGLWECiphertextExec<D, B> {
 88    type Inner = VmpPMat<D, B>;
 89
 90    fn inner(&self) -> &Self::Inner {
 91        &self.data
 92    }
 93
 94    fn basek(&self) -> usize {
 95        self.basek
 96    }
 97
 98    fn k(&self) -> usize {
 99        self.k
100    }
101}
102
103impl<D: Data, B: Backend> GGLWECiphertextExec<D, B> {
104    pub fn rank(&self) -> usize {
105        self.data.cols_out() - 1
106    }
107
108    pub fn digits(&self) -> usize {
109        self.digits
110    }
111
112    pub fn rank_in(&self) -> usize {
113        self.data.cols_in()
114    }
115
116    pub fn rank_out(&self) -> usize {
117        self.data.cols_out() - 1
118    }
119}
120
121impl<D: DataMut, B: Backend> GGLWECiphertextExec<D, B> {
122    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GGLWECiphertext<DataOther>, scratch: &mut Scratch<B>)
123    where
124        DataOther: DataRef,
125        Module<B>: GGLWEExecLayoutFamily<B>,
126    {
127        module.vmp_prepare(&mut self.data, &other.data, scratch);
128        self.basek = other.basek;
129        self.k = other.k;
130        self.digits = other.digits;
131    }
132}
133
134#[derive(PartialEq, Eq)]
135pub struct GLWESwitchingKeyExec<D: Data, B: Backend> {
136    pub(crate) key: GGLWECiphertextExec<D, B>,
137    pub(crate) sk_in_n: usize,  // Degree of sk_in
138    pub(crate) sk_out_n: usize, // Degree of sk_out
139}
140
141impl<B: Backend> GLWESwitchingKeyExec<Vec<u8>, B> {
142    pub fn alloc(
143        module: &Module<B>,
144        n: usize,
145        basek: usize,
146        k: usize,
147        rows: usize,
148        digits: usize,
149        rank_in: usize,
150        rank_out: usize,
151    ) -> Self
152    where
153        Module<B>: GGLWEExecLayoutFamily<B>,
154    {
155        GLWESwitchingKeyExec::<Vec<u8>, B> {
156            key: GGLWECiphertextExec::alloc(module, n, basek, k, rows, digits, rank_in, rank_out),
157            sk_in_n: 0,
158            sk_out_n: 0,
159        }
160    }
161
162    pub fn bytes_of(
163        module: &Module<B>,
164        n: usize,
165        basek: usize,
166        k: usize,
167        rows: usize,
168        digits: usize,
169        rank_in: usize,
170        rank_out: usize,
171    ) -> usize
172    where
173        Module<B>: GGLWEExecLayoutFamily<B>,
174    {
175        GGLWECiphertextExec::bytes_of(module, n, basek, k, rows, digits, rank_in, rank_out)
176    }
177
178    pub fn from<DataOther: DataRef>(module: &Module<B>, other: &GLWESwitchingKey<DataOther>, scratch: &mut Scratch<B>) -> Self
179    where
180        Module<B>: GGLWEExecLayoutFamily<B>,
181    {
182        let mut ksk_exec: GLWESwitchingKeyExec<Vec<u8>, B> = Self::alloc(
183            module,
184            other.n(),
185            other.basek(),
186            other.k(),
187            other.rows(),
188            other.digits(),
189            other.rank_in(),
190            other.rank_out(),
191        );
192        ksk_exec.prepare(module, other, scratch);
193        ksk_exec
194    }
195}
196
197impl<D: Data, B: Backend> Infos for GLWESwitchingKeyExec<D, B> {
198    type Inner = VmpPMat<D, B>;
199
200    fn inner(&self) -> &Self::Inner {
201        self.key.inner()
202    }
203
204    fn basek(&self) -> usize {
205        self.key.basek()
206    }
207
208    fn k(&self) -> usize {
209        self.key.k()
210    }
211}
212
213impl<D: Data, B: Backend> GLWESwitchingKeyExec<D, B> {
214    pub fn rank(&self) -> usize {
215        self.key.data.cols_out() - 1
216    }
217
218    pub fn rank_in(&self) -> usize {
219        self.key.data.cols_in()
220    }
221
222    pub fn rank_out(&self) -> usize {
223        self.key.data.cols_out() - 1
224    }
225
226    pub fn digits(&self) -> usize {
227        self.key.digits()
228    }
229
230    pub fn sk_degree_in(&self) -> usize {
231        self.sk_in_n
232    }
233
234    pub fn sk_degree_out(&self) -> usize {
235        self.sk_out_n
236    }
237}
238
239impl<D: DataMut, B: Backend> GLWESwitchingKeyExec<D, B> {
240    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GLWESwitchingKey<DataOther>, scratch: &mut Scratch<B>)
241    where
242        DataOther: DataRef,
243        Module<B>: GGLWEExecLayoutFamily<B>,
244    {
245        self.key.prepare(module, &other.key, scratch);
246        self.sk_in_n = other.sk_in_n;
247        self.sk_out_n = other.sk_out_n;
248    }
249}
250
251#[derive(PartialEq, Eq)]
252pub struct AutomorphismKeyExec<D: Data, B: Backend> {
253    pub(crate) key: GLWESwitchingKeyExec<D, B>,
254    pub(crate) p: i64,
255}
256
257impl<B: Backend> AutomorphismKeyExec<Vec<u8>, B> {
258    pub fn alloc(module: &Module<B>, n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
259    where
260        Module<B>: GGLWEExecLayoutFamily<B>,
261    {
262        AutomorphismKeyExec::<Vec<u8>, B> {
263            key: GLWESwitchingKeyExec::alloc(module, n, basek, k, rows, digits, rank, rank),
264            p: 0,
265        }
266    }
267
268    pub fn bytes_of(module: &Module<B>, n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
269    where
270        Module<B>: GGLWEExecLayoutFamily<B>,
271    {
272        GLWESwitchingKeyExec::bytes_of(module, n, basek, k, rows, digits, rank, rank)
273    }
274
275    pub fn from<DataOther: DataRef>(module: &Module<B>, other: &AutomorphismKey<DataOther>, scratch: &mut Scratch<B>) -> Self
276    where
277        Module<B>: GGLWEExecLayoutFamily<B>,
278    {
279        let mut atk_exec: AutomorphismKeyExec<Vec<u8>, B> = Self::alloc(
280            module,
281            other.n(),
282            other.basek(),
283            other.k(),
284            other.rows(),
285            other.digits(),
286            other.rank(),
287        );
288        atk_exec.prepare(module, other, scratch);
289        atk_exec
290    }
291}
292
293impl<D: DataMut, B: Backend> AutomorphismKeyExec<D, B> {
294    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &AutomorphismKey<DataOther>, scratch: &mut Scratch<B>)
295    where
296        DataOther: DataRef,
297        Module<B>: GGLWEExecLayoutFamily<B>,
298    {
299        self.key.prepare(module, &other.key, scratch);
300        self.p = other.p;
301    }
302}
303
304impl<D: Data, B: Backend> Infos for AutomorphismKeyExec<D, B> {
305    type Inner = VmpPMat<D, B>;
306
307    fn inner(&self) -> &Self::Inner {
308        &self.key.inner()
309    }
310
311    fn basek(&self) -> usize {
312        self.key.basek()
313    }
314
315    fn k(&self) -> usize {
316        self.key.k()
317    }
318}
319
320impl<D: Data, B: Backend> AutomorphismKeyExec<D, B> {
321    pub fn p(&self) -> i64 {
322        self.p
323    }
324
325    pub fn digits(&self) -> usize {
326        self.key.digits()
327    }
328
329    pub fn rank(&self) -> usize {
330        self.key.rank()
331    }
332
333    pub fn rank_in(&self) -> usize {
334        self.key.rank_in()
335    }
336
337    pub fn rank_out(&self) -> usize {
338        self.key.rank_out()
339    }
340}
341
342#[derive(PartialEq, Eq)]
343pub struct GLWETensorKeyExec<D: Data, B: Backend> {
344    pub(crate) keys: Vec<GLWESwitchingKeyExec<D, B>>,
345}
346
347impl<B: Backend> GLWETensorKeyExec<Vec<u8>, B> {
348    pub fn alloc(module: &Module<B>, n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self
349    where
350        Module<B>: GGLWEExecLayoutFamily<B>,
351    {
352        let mut keys: Vec<GLWESwitchingKeyExec<Vec<u8>, B>> = Vec::new();
353        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
354        (0..pairs).for_each(|_| {
355            keys.push(GLWESwitchingKeyExec::alloc(
356                module, n, basek, k, rows, digits, 1, rank,
357            ));
358        });
359        Self { keys }
360    }
361
362    pub fn bytes_of(module: &Module<B>, n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize
363    where
364        Module<B>: GGLWEExecLayoutFamily<B>,
365    {
366        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
367        pairs * GLWESwitchingKeyExec::bytes_of(module, n, basek, k, rows, digits, 1, rank)
368    }
369}
370
371impl<D: Data, B: Backend> Infos for GLWETensorKeyExec<D, B> {
372    type Inner = VmpPMat<D, B>;
373
374    fn inner(&self) -> &Self::Inner {
375        &self.keys[0].inner()
376    }
377
378    fn basek(&self) -> usize {
379        self.keys[0].basek()
380    }
381
382    fn k(&self) -> usize {
383        self.keys[0].k()
384    }
385}
386
387impl<D: Data, B: Backend> GLWETensorKeyExec<D, B> {
388    pub fn rank(&self) -> usize {
389        self.keys[0].rank()
390    }
391
392    pub fn rank_in(&self) -> usize {
393        self.keys[0].rank_in()
394    }
395
396    pub fn rank_out(&self) -> usize {
397        self.keys[0].rank_out()
398    }
399
400    pub fn digits(&self) -> usize {
401        self.keys[0].digits()
402    }
403}
404
405impl<D: DataMut, B: Backend> GLWETensorKeyExec<D, B> {
406    // Returns a mutable reference to GLWESwitchingKey_{s}(s[i] * s[j])
407    pub fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GLWESwitchingKeyExec<D, B> {
408        if i > j {
409            std::mem::swap(&mut i, &mut j);
410        };
411        let rank: usize = self.rank();
412        &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
413    }
414}
415
416impl<D: DataRef, B: Backend> GLWETensorKeyExec<D, B> {
417    // Returns a reference to GLWESwitchingKey_{s}(s[i] * s[j])
418    pub fn at(&self, mut i: usize, mut j: usize) -> &GLWESwitchingKeyExec<D, B> {
419        if i > j {
420            std::mem::swap(&mut i, &mut j);
421        };
422        let rank: usize = self.rank();
423        &self.keys[i * rank + j - (i * (i + 1) / 2)]
424    }
425}
426
427impl<D: DataMut, B: Backend> GLWETensorKeyExec<D, B> {
428    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GLWETensorKey<DataOther>, scratch: &mut Scratch<B>)
429    where
430        DataOther: DataRef,
431        Module<B>: GGLWEExecLayoutFamily<B>,
432    {
433        #[cfg(debug_assertions)]
434        {
435            assert_eq!(self.keys.len(), other.keys.len());
436        }
437        self.keys
438            .iter_mut()
439            .zip(other.keys.iter())
440            .for_each(|(a, b)| {
441                a.prepare(module, b, scratch);
442            });
443    }
444}