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