1use backend::{FFT64, Module, Scratch, VecZnx, VecZnxDftOps, VecZnxOps, ZnxZero};
  2
  3use crate::{FourierGLWECiphertext, GLWEAutomorphismKey, GLWECiphertext, GetRow, Infos, ScratchCore, SetRow};
  4
  5impl GLWEAutomorphismKey<Vec<u8>, FFT64> {
  6    pub fn automorphism_scratch_space(
  7        module: &Module<FFT64>,
  8        basek: usize,
  9        k_out: usize,
 10        k_in: usize,
 11        k_ksk: usize,
 12        digits: usize,
 13        rank: usize,
 14    ) -> usize {
 15        let tmp_dft: usize = FourierGLWECiphertext::bytes_of(module, basek, k_in, rank);
 16        let tmp_idft: usize = FourierGLWECiphertext::bytes_of(module, basek, k_out, rank);
 17        let idft: usize = module.vec_znx_idft_tmp_bytes();
 18        let keyswitch: usize = GLWECiphertext::keyswitch_inplace_scratch_space(module, basek, k_out, k_ksk, digits, rank);
 19        tmp_dft + tmp_idft + idft + keyswitch
 20    }
 21
 22    pub fn automorphism_inplace_scratch_space(
 23        module: &Module<FFT64>,
 24        basek: usize,
 25        k_out: usize,
 26        k_ksk: usize,
 27        digits: usize,
 28        rank: usize,
 29    ) -> usize {
 30        GLWEAutomorphismKey::automorphism_scratch_space(module, basek, k_out, k_out, k_ksk, digits, rank)
 31    }
 32}
 33
 34impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GLWEAutomorphismKey<DataSelf, FFT64> {
 35    pub fn automorphism<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
 36        &mut self,
 37        module: &Module<FFT64>,
 38        lhs: &GLWEAutomorphismKey<DataLhs, FFT64>,
 39        rhs: &GLWEAutomorphismKey<DataRhs, FFT64>,
 40        scratch: &mut Scratch,
 41    ) {
 42        #[cfg(debug_assertions)]
 43        {
 44            assert_eq!(
 45                self.rank_in(),
 46                lhs.rank_in(),
 47                "ksk_out input rank: {} != ksk_in input rank: {}",
 48                self.rank_in(),
 49                lhs.rank_in()
 50            );
 51            assert_eq!(
 52                lhs.rank_out(),
 53                rhs.rank_in(),
 54                "ksk_in output rank: {} != ksk_apply input rank: {}",
 55                self.rank_out(),
 56                rhs.rank_in()
 57            );
 58            assert_eq!(
 59                self.rank_out(),
 60                rhs.rank_out(),
 61                "ksk_out output rank: {} != ksk_apply output rank: {}",
 62                self.rank_out(),
 63                rhs.rank_out()
 64            );
 65        }
 66
 67        let cols_out: usize = rhs.rank_out() + 1;
 68
 69        let (mut tmp_dft, scratch1) = scratch.tmp_glwe_fourier(module, lhs.basek(), lhs.k(), lhs.rank());
 70
 71        (0..self.rank_in()).for_each(|col_i| {
 72            (0..self.rows()).for_each(|row_j| {
 73                // Extracts relevant row
 74                lhs.get_row(module, row_j, col_i, &mut tmp_dft);
 75
 76                // Get a VecZnxBig from scratch space
 77                let (mut tmp_idft_data, scratch2) = scratch1.tmp_vec_znx_big(module, cols_out, self.size());
 78
 79                // Switches input outside of DFT
 80                (0..cols_out).for_each(|i| {
 81                    module.vec_znx_idft(&mut tmp_idft_data, i, &tmp_dft.data, i, scratch2);
 82                });
 83
 84                // Consumes to small vec znx
 85                let mut tmp_idft_small_data: VecZnx<&mut [u8]> = tmp_idft_data.to_vec_znx_small();
 86
 87                // Reverts the automorphis key from (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
 88                (0..cols_out).for_each(|i| {
 89                    module.vec_znx_automorphism_inplace(lhs.p(), &mut tmp_idft_small_data, i);
 90                });
 91
 92                // Wraps into ciphertext
 93                let mut tmp_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
 94                    data: tmp_idft_small_data,
 95                    basek: self.basek(),
 96                    k: self.k(),
 97                };
 98
 99                // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
100                tmp_idft.keyswitch_inplace(module, &rhs.key, scratch2);
101
102                // Applies back the automorphism X^{k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) -> (-pi^{-1}_{k'+k}(s)a + s, a)
103                // and switches back to DFT domain
104                (0..self.rank_out() + 1).for_each(|i| {
105                    module.vec_znx_automorphism_inplace(lhs.p(), &mut tmp_idft.data, i);
106                    module.vec_znx_dft(1, 0, &mut tmp_dft.data, i, &tmp_idft.data, i);
107                });
108
109                // Sets back the relevant row
110                self.set_row(module, row_j, col_i, &tmp_dft);
111            });
112        });
113
114        tmp_dft.data.zero();
115
116        (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
117            (0..self.rank_in()).for_each(|col_j| {
118                self.set_row(module, row_i, col_j, &tmp_dft);
119            });
120        });
121
122        self.p = (lhs.p * rhs.p) % (module.cyclotomic_order() as i64);
123    }
124
125    pub fn automorphism_inplace<DataRhs: AsRef<[u8]>>(
126        &mut self,
127        module: &Module<FFT64>,
128        rhs: &GLWEAutomorphismKey<DataRhs, FFT64>,
129        scratch: &mut Scratch,
130    ) {
131        unsafe {
132            let self_ptr: *mut GLWEAutomorphismKey<DataSelf, FFT64> = self as *mut GLWEAutomorphismKey<DataSelf, FFT64>;
133            self.automorphism(&module, &*self_ptr, rhs, scratch);
134        }
135    }
136}