1use backend::{FFT64, Module, Scratch, ZnxZero};
  2
  3use crate::{FourierGLWECiphertext, GLWEAutomorphismKey, GLWESwitchingKey, GetRow, Infos, ScratchCore, SetRow};
  4
  5impl GLWEAutomorphismKey<Vec<u8>, FFT64> {
  6    pub fn keyswitch_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        GLWESwitchingKey::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits, rank, rank)
 16    }
 17
 18    pub fn keyswitch_inplace_scratch_space(
 19        module: &Module<FFT64>,
 20        basek: usize,
 21        k_out: usize,
 22        k_ksk: usize,
 23        digits: usize,
 24        rank: usize,
 25    ) -> usize {
 26        GLWESwitchingKey::keyswitch_inplace_scratch_space(module, basek, k_out, k_ksk, digits, rank)
 27    }
 28}
 29
 30impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GLWEAutomorphismKey<DataSelf, FFT64> {
 31    pub fn keyswitch<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
 32        &mut self,
 33        module: &Module<FFT64>,
 34        lhs: &GLWEAutomorphismKey<DataLhs, FFT64>,
 35        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
 36        scratch: &mut Scratch,
 37    ) {
 38        self.key.keyswitch(module, &lhs.key, rhs, scratch);
 39    }
 40
 41    pub fn keyswitch_inplace<DataRhs: AsRef<[u8]>>(
 42        &mut self,
 43        module: &Module<FFT64>,
 44        rhs: &GLWEAutomorphismKey<DataRhs, FFT64>,
 45        scratch: &mut Scratch,
 46    ) {
 47        self.key.keyswitch_inplace(module, &rhs.key, scratch);
 48    }
 49}
 50
 51impl GLWESwitchingKey<Vec<u8>, FFT64> {
 52    pub fn keyswitch_scratch_space(
 53        module: &Module<FFT64>,
 54        basek: usize,
 55        k_out: usize,
 56        k_in: usize,
 57        k_ksk: usize,
 58        digits: usize,
 59        rank_in: usize,
 60        rank_out: usize,
 61    ) -> usize {
 62        let tmp_in: usize = FourierGLWECiphertext::bytes_of(module, basek, k_in, rank_in);
 63        let tmp_out: usize = FourierGLWECiphertext::bytes_of(module, basek, k_out, rank_out);
 64        let ksk: usize =
 65            FourierGLWECiphertext::keyswitch_scratch_space(module, basek, k_out, k_in, k_ksk, digits, rank_in, rank_out);
 66        tmp_in + tmp_out + ksk
 67    }
 68
 69    pub fn keyswitch_inplace_scratch_space(
 70        module: &Module<FFT64>,
 71        basek: usize,
 72        k_out: usize,
 73        k_ksk: usize,
 74        digits: usize,
 75        rank: usize,
 76    ) -> usize {
 77        let tmp: usize = FourierGLWECiphertext::bytes_of(module, basek, k_out, rank);
 78        let ksk: usize = FourierGLWECiphertext::keyswitch_inplace_scratch_space(module, basek, k_out, k_ksk, digits, rank);
 79        tmp + ksk
 80    }
 81}
 82
 83impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GLWESwitchingKey<DataSelf, FFT64> {
 84    pub fn keyswitch<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
 85        &mut self,
 86        module: &Module<FFT64>,
 87        lhs: &GLWESwitchingKey<DataLhs, FFT64>,
 88        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
 89        scratch: &mut Scratch,
 90    ) {
 91        #[cfg(debug_assertions)]
 92        {
 93            assert_eq!(
 94                self.rank_in(),
 95                lhs.rank_in(),
 96                "ksk_out input rank: {} != ksk_in input rank: {}",
 97                self.rank_in(),
 98                lhs.rank_in()
 99            );
100            assert_eq!(
101                lhs.rank_out(),
102                rhs.rank_in(),
103                "ksk_in output rank: {} != ksk_apply input rank: {}",
104                self.rank_out(),
105                rhs.rank_in()
106            );
107            assert_eq!(
108                self.rank_out(),
109                rhs.rank_out(),
110                "ksk_out output rank: {} != ksk_apply output rank: {}",
111                self.rank_out(),
112                rhs.rank_out()
113            );
114        }
115
116        let (mut tmp_in, scratch1) = scratch.tmp_glwe_fourier(module, lhs.basek(), lhs.k(), lhs.rank());
117        let (mut tmp_out, scratch2) = scratch1.tmp_glwe_fourier(module, self.basek(), self.k(), self.rank());
118
119        (0..self.rank_in()).for_each(|col_i| {
120            (0..self.rows()).for_each(|row_j| {
121                lhs.get_row(module, row_j, col_i, &mut tmp_in);
122                tmp_out.keyswitch(module, &tmp_in, rhs, scratch2);
123                self.set_row(module, row_j, col_i, &tmp_out);
124            });
125        });
126
127        tmp_out.data.zero();
128
129        (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
130            (0..self.rank_in()).for_each(|col_j| {
131                self.set_row(module, row_i, col_j, &tmp_out);
132            });
133        });
134    }
135
136    pub fn keyswitch_inplace<DataRhs: AsRef<[u8]>>(
137        &mut self,
138        module: &Module<FFT64>,
139        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
140        scratch: &mut Scratch,
141    ) {
142        #[cfg(debug_assertions)]
143        {
144            assert_eq!(
145                self.rank_out(),
146                rhs.rank_out(),
147                "ksk_out output rank: {} != ksk_apply output rank: {}",
148                self.rank_out(),
149                rhs.rank_out()
150            );
151        }
152
153        let (mut tmp, scratch1) = scratch.tmp_glwe_fourier(module, self.basek(), self.k(), self.rank());
154
155        (0..self.rank_in()).for_each(|col_i| {
156            (0..self.rows()).for_each(|row_j| {
157                self.get_row(module, row_j, col_i, &mut tmp);
158                tmp.keyswitch_inplace(module, rhs, scratch1);
159                self.set_row(module, row_j, col_i, &tmp);
160            });
161        });
162    }
163}