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