1use base2k::{
  2    Backend, FFT64, MatZnxDft, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDftOps, ScalarZnxOps,
  3    ScalarZnxToRef, Scratch, VecZnxDft, VecZnxDftToMut, VecZnxDftToRef,
  4};
  5use sampling::source::Source;
  6
  7use crate::{
  8    elem::{GetRow, Infos, SetRow},
  9    gglwe_ciphertext::GGLWECiphertext,
 10    ggsw_ciphertext::GGSWCiphertext,
 11    glwe_ciphertext_fourier::GLWECiphertextFourier,
 12    keys::{SecretKey, SecretKeyFourier},
 13    keyswitch_key::GLWESwitchingKey,
 14};
 15
 16pub struct AutomorphismKey<Data, B: Backend> {
 17    pub(crate) key: GLWESwitchingKey<Data, B>,
 18    pub(crate) p: i64,
 19}
 20
 21impl AutomorphismKey<Vec<u8>, FFT64> {
 22    pub fn new(module: &Module<FFT64>, basek: usize, p: i64, k: usize, rows: usize, rank: usize) -> Self {
 23        AutomorphismKey {
 24            key: GLWESwitchingKey::new(module, basek, k, rows, rank, rank),
 25            p: p,
 26        }
 27    }
 28}
 29
 30impl<T, B: Backend> Infos for AutomorphismKey<T, B> {
 31    type Inner = MatZnxDft<T, B>;
 32
 33    fn inner(&self) -> &Self::Inner {
 34        &self.key.inner()
 35    }
 36
 37    fn basek(&self) -> usize {
 38        self.key.basek()
 39    }
 40
 41    fn k(&self) -> usize {
 42        self.key.k()
 43    }
 44}
 45
 46impl<T, B: Backend> AutomorphismKey<T, B> {
 47    pub fn p(&self) -> i64 {
 48        self.p
 49    }
 50
 51    pub fn rank(&self) -> usize {
 52        self.key.rank()
 53    }
 54
 55    pub fn rank_in(&self) -> usize {
 56        self.key.rank_in()
 57    }
 58
 59    pub fn rank_out(&self) -> usize {
 60        self.key.rank_out()
 61    }
 62}
 63
 64impl<DataSelf, B: Backend> MatZnxDftToMut<B> for AutomorphismKey<DataSelf, B>
 65where
 66    MatZnxDft<DataSelf, B>: MatZnxDftToMut<B>,
 67{
 68    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
 69        self.key.to_mut()
 70    }
 71}
 72
 73impl<DataSelf, B: Backend> MatZnxDftToRef<B> for AutomorphismKey<DataSelf, B>
 74where
 75    MatZnxDft<DataSelf, B>: MatZnxDftToRef<B>,
 76{
 77    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
 78        self.key.to_ref()
 79    }
 80}
 81
 82impl<C> GetRow<FFT64> for AutomorphismKey<C, FFT64>
 83where
 84    MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
 85{
 86    fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut GLWECiphertextFourier<R, FFT64>)
 87    where
 88        VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
 89    {
 90        module.vmp_extract_row(res, self, row_i, col_j);
 91    }
 92}
 93
 94impl<C> SetRow<FFT64> for AutomorphismKey<C, FFT64>
 95where
 96    MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
 97{
 98    fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &GLWECiphertextFourier<R, FFT64>)
 99    where
100        VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64>,
101    {
102        module.vmp_prepare_row(self, row_i, col_j, a);
103    }
104}
105
106impl AutomorphismKey<Vec<u8>, FFT64> {
107    pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
108        GGLWECiphertext::encrypt_sk_scratch_space(module, rank, size)
109    }
110
111    pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, rank: usize, pk_size: usize) -> usize {
112        GGLWECiphertext::encrypt_pk_scratch_space(module, rank, pk_size)
113    }
114
115    pub fn keyswitch_scratch_space(
116        module: &Module<FFT64>,
117        out_size: usize,
118        in_size: usize,
119        ksk_size: usize,
120        rank: usize,
121    ) -> usize {
122        GLWESwitchingKey::keyswitch_scratch_space(module, out_size, rank, in_size, rank, ksk_size)
123    }
124
125    pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
126        GLWESwitchingKey::keyswitch_inplace_scratch_space(module, out_size, out_rank, ksk_size)
127    }
128
129    pub fn external_product_scratch_space(
130        module: &Module<FFT64>,
131        out_size: usize,
132        in_size: usize,
133        ggsw_size: usize,
134        rank: usize,
135    ) -> usize {
136        GLWESwitchingKey::external_product_scratch_space(module, out_size, in_size, ggsw_size, rank)
137    }
138
139    pub fn external_product_inplace_scratch_space(
140        module: &Module<FFT64>,
141        out_size: usize,
142        ggsw_size: usize,
143        rank: usize,
144    ) -> usize {
145        GLWESwitchingKey::external_product_inplace_scratch_space(module, out_size, ggsw_size, rank)
146    }
147}
148
149impl<DataSelf> AutomorphismKey<DataSelf, FFT64>
150where
151    MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
152{
153    pub fn encrypt_sk<DataSk>(
154        &mut self,
155        module: &Module<FFT64>,
156        p: i64,
157        sk: &SecretKey<DataSk>,
158        source_xa: &mut Source,
159        source_xe: &mut Source,
160        sigma: f64,
161        scratch: &mut Scratch,
162    ) where
163        ScalarZnx<DataSk>: ScalarZnxToRef,
164    {
165        #[cfg(debug_assertions)]
166        {
167            assert_eq!(self.n(), module.n());
168            assert_eq!(sk.n(), module.n());
169            assert_eq!(self.rank_out(), self.rank_in());
170            assert_eq!(sk.rank(), self.rank());
171        }
172
173        let (sk_out_dft_data, scratch_1) = scratch.tmp_scalar_znx_dft(module, sk.rank());
174
175        let mut sk_out_dft: SecretKeyFourier<&mut [u8], FFT64> = SecretKeyFourier {
176            data: sk_out_dft_data,
177            dist: sk.dist,
178        };
179
180        {
181            (0..self.rank()).for_each(|i| {
182                let (mut sk_inv_auto, _) = scratch_1.tmp_scalar_znx(module, 1);
183                module.scalar_znx_automorphism(module.galois_element_inv(p), &mut sk_inv_auto, 0, sk, i);
184                module.svp_prepare(&mut sk_out_dft, i, &sk_inv_auto, 0);
185            });
186        }
187
188        self.key.encrypt_sk(
189            module,
190            &sk,
191            &sk_out_dft,
192            source_xa,
193            source_xe,
194            sigma,
195            scratch_1,
196        );
197
198        self.p = p;
199    }
200}
201
202impl<DataSelf> AutomorphismKey<DataSelf, FFT64>
203where
204    MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
205{
206    pub fn keyswitch<DataLhs, DataRhs>(
207        &mut self,
208        module: &Module<FFT64>,
209        lhs: &AutomorphismKey<DataLhs, FFT64>,
210        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
211        scratch: &mut base2k::Scratch,
212    ) where
213        MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
214        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
215    {
216        self.key.keyswitch(module, &lhs.key, rhs, scratch);
217    }
218
219    pub fn keyswitch_inplace<DataRhs>(
220        &mut self,
221        module: &Module<FFT64>,
222        rhs: &AutomorphismKey<DataRhs, FFT64>,
223        scratch: &mut base2k::Scratch,
224    ) where
225        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
226    {
227        self.key.keyswitch_inplace(module, &rhs.key, scratch);
228    }
229
230    pub fn external_product<DataLhs, DataRhs>(
231        &mut self,
232        module: &Module<FFT64>,
233        lhs: &AutomorphismKey<DataLhs, FFT64>,
234        rhs: &GGSWCiphertext<DataRhs, FFT64>,
235        scratch: &mut Scratch,
236    ) where
237        MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
238        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
239    {
240        self.key.external_product(module, &lhs.key, rhs, scratch);
241    }
242
243    pub fn external_product_inplace<DataRhs>(
244        &mut self,
245        module: &Module<FFT64>,
246        rhs: &GGSWCiphertext<DataRhs, FFT64>,
247        scratch: &mut Scratch,
248    ) where
249        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
250    {
251        self.key.external_product_inplace(module, rhs, scratch);
252    }
253}