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            assert!(
 66                self.k() <= lhs.k(),
 67                "output k={} cannot be greater than input k={}",
 68                self.k(),
 69                lhs.k()
 70            )
 71        }
 72
 73        let cols_out: usize = rhs.rank_out() + 1;
 74
 75        (0..self.rank_in()).for_each(|col_i| {
 76            (0..self.rows()).for_each(|row_j| {
 77                let (mut tmp_idft_data, scratct1) = scratch.tmp_vec_znx_big(module, cols_out, self.size());
 78
 79                {
 80                    let (mut tmp_dft, scratch2) = scratct1.tmp_fourier_glwe_ct(module, lhs.basek(), lhs.k(), lhs.rank());
 81
 82                    // Extracts relevant row
 83                    lhs.get_row(module, row_j, col_i, &mut tmp_dft);
 84
 85                    // Get a VecZnxBig from scratch space
 86
 87                    // Switches input outside of DFT
 88                    (0..cols_out).for_each(|i| {
 89                        module.vec_znx_idft(&mut tmp_idft_data, i, &tmp_dft.data, i, scratch2);
 90                    });
 91                }
 92
 93                // Consumes to small vec znx
 94                let mut tmp_idft_small_data: VecZnx<&mut [u8]> = tmp_idft_data.to_vec_znx_small();
 95
 96                // Reverts the automorphis key from (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
 97                (0..cols_out).for_each(|i| {
 98                    module.vec_znx_automorphism_inplace(lhs.p(), &mut tmp_idft_small_data, i);
 99                });
100
101                // Wraps into ciphertext
102                let mut tmp_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
103                    data: tmp_idft_small_data,
104                    basek: self.basek(),
105                    k: self.k(),
106                };
107
108                // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
109                tmp_idft.keyswitch_inplace(module, &rhs.key, scratct1);
110
111                {
112                    let (mut tmp_dft, _) = scratct1.tmp_fourier_glwe_ct(module, self.basek(), self.k(), self.rank());
113
114                    // Applies back the automorphism X^{k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) -> (-pi^{-1}_{k'+k}(s)a + s, a)
115                    // and switches back to DFT domain
116                    (0..self.rank_out() + 1).for_each(|i| {
117                        module.vec_znx_automorphism_inplace(lhs.p(), &mut tmp_idft.data, i);
118                        module.vec_znx_dft(1, 0, &mut tmp_dft.data, i, &tmp_idft.data, i);
119                    });
120
121                    // Sets back the relevant row
122                    self.set_row(module, row_j, col_i, &tmp_dft);
123                }
124            });
125        });
126
127        let (mut tmp_dft, _) = scratch.tmp_fourier_glwe_ct(module, self.basek(), self.k(), self.rank());
128        tmp_dft.data.zero();
129
130        (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
131            (0..self.rank_in()).for_each(|col_j| {
132                self.set_row(module, row_i, col_j, &tmp_dft);
133            });
134        });
135
136        self.p = (lhs.p * rhs.p) % (module.cyclotomic_order() as i64);
137    }
138
139    pub fn automorphism_inplace<DataRhs: AsRef<[u8]>>(
140        &mut self,
141        module: &Module<FFT64>,
142        rhs: &GLWEAutomorphismKey<DataRhs, FFT64>,
143        scratch: &mut Scratch,
144    ) {
145        unsafe {
146            let self_ptr: *mut GLWEAutomorphismKey<DataSelf, FFT64> = self as *mut GLWEAutomorphismKey<DataSelf, FFT64>;
147            self.automorphism(&module, &*self_ptr, rhs, scratch);
148        }
149    }
150}