1use backend::{Backend, FFT64, MatZnxDft, MatZnxDftOps, Module};
 2
 3use crate::{FourierGLWECiphertext, GGLWECiphertext, GetRow, Infos, SetRow};
 4
 5pub struct GLWESwitchingKey<Data, B: Backend>(pub(crate) GGLWECiphertext<Data, B>);
 6
 7impl GLWESwitchingKey<Vec<u8>, FFT64> {
 8    pub fn alloc(
 9        module: &Module<FFT64>,
10        basek: usize,
11        k: usize,
12        rows: usize,
13        digits: usize,
14        rank_in: usize,
15        rank_out: usize,
16    ) -> Self {
17        GLWESwitchingKey(GGLWECiphertext::alloc(
18            module, basek, k, rows, digits, rank_in, rank_out,
19        ))
20    }
21
22    pub fn bytes_of(
23        module: &Module<FFT64>,
24        basek: usize,
25        k: usize,
26        rows: usize,
27        digits: usize,
28        rank_in: usize,
29        rank_out: usize,
30    ) -> usize {
31        GGLWECiphertext::<Vec<u8>, FFT64>::bytes_of(module, basek, k, rows, digits, rank_in, rank_out)
32    }
33}
34
35impl<T, B: Backend> Infos for GLWESwitchingKey<T, B> {
36    type Inner = MatZnxDft<T, B>;
37
38    fn inner(&self) -> &Self::Inner {
39        self.0.inner()
40    }
41
42    fn basek(&self) -> usize {
43        self.0.basek()
44    }
45
46    fn k(&self) -> usize {
47        self.0.k()
48    }
49}
50
51impl<T, B: Backend> GLWESwitchingKey<T, B> {
52    pub fn rank(&self) -> usize {
53        self.0.data.cols_out() - 1
54    }
55
56    pub fn rank_in(&self) -> usize {
57        self.0.data.cols_in()
58    }
59
60    pub fn rank_out(&self) -> usize {
61        self.0.data.cols_out() - 1
62    }
63
64    pub fn digits(&self) -> usize {
65        self.0.digits()
66    }
67}
68
69impl<C: AsRef<[u8]>> GetRow<FFT64> for GLWESwitchingKey<C, FFT64> {
70    fn get_row<R: AsMut<[u8]> + AsRef<[u8]>>(
71        &self,
72        module: &Module<FFT64>,
73        row_i: usize,
74        col_j: usize,
75        res: &mut FourierGLWECiphertext<R, FFT64>,
76    ) {
77        module.mat_znx_dft_get_row(&mut res.data, &self.0.data, row_i, col_j);
78    }
79}
80
81impl<C: AsMut<[u8]> + AsRef<[u8]>> SetRow<FFT64> for GLWESwitchingKey<C, FFT64> {
82    fn set_row<R: AsRef<[u8]>>(
83        &mut self,
84        module: &Module<FFT64>,
85        row_i: usize,
86        col_j: usize,
87        a: &FourierGLWECiphertext<R, FFT64>,
88    ) {
89        module.mat_znx_dft_set_row(&mut self.0.data, row_i, col_j, &a.data);
90    }
91}