1use base2k::{
  2    Backend, FFT64, MatZnxDft, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDft, ScalarZnxDftToRef,
  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    vec_glwe_product::{VecGLWEProduct, VecGLWEProductScratchSpace},
 14};
 15
 16pub struct GLWESwitchingKey<Data, B: Backend>(pub(crate) GGLWECiphertext<Data, B>);
 17
 18impl GLWESwitchingKey<Vec<u8>, FFT64> {
 19    pub fn new(module: &Module<FFT64>, base2k: usize, k: usize, rows: usize, rank_in: usize, rank_out: usize) -> Self {
 20        GLWESwitchingKey(GGLWECiphertext::new(
 21            module, base2k, k, rows, rank_in, rank_out,
 22        ))
 23    }
 24}
 25
 26impl<T, B: Backend> Infos for GLWESwitchingKey<T, B> {
 27    type Inner = MatZnxDft<T, B>;
 28
 29    fn inner(&self) -> &Self::Inner {
 30        &self.0.inner()
 31    }
 32
 33    fn basek(&self) -> usize {
 34        self.0.basek()
 35    }
 36
 37    fn k(&self) -> usize {
 38        self.0.k()
 39    }
 40}
 41
 42impl<DataSelf, B: Backend> MatZnxDftToMut<B> for GLWESwitchingKey<DataSelf, B>
 43where
 44    MatZnxDft<DataSelf, B>: MatZnxDftToMut<B>,
 45{
 46    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
 47        self.0.data.to_mut()
 48    }
 49}
 50
 51impl<DataSelf, B: Backend> MatZnxDftToRef<B> for GLWESwitchingKey<DataSelf, B>
 52where
 53    MatZnxDft<DataSelf, B>: MatZnxDftToRef<B>,
 54{
 55    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
 56        self.0.data.to_ref()
 57    }
 58}
 59
 60impl<C> GetRow<FFT64> for GLWESwitchingKey<C, FFT64>
 61where
 62    MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
 63{
 64    fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut GLWECiphertextFourier<R, FFT64>)
 65    where
 66        VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
 67    {
 68        #[cfg(debug_assertions)]
 69        {
 70            assert_eq!(col_j, 0);
 71        }
 72        module.vmp_extract_row(res, self, row_i, col_j);
 73    }
 74}
 75
 76impl<C> SetRow<FFT64> for GLWESwitchingKey<C, FFT64>
 77where
 78    MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
 79{
 80    fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &GLWECiphertextFourier<R, FFT64>)
 81    where
 82        VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64>,
 83    {
 84        #[cfg(debug_assertions)]
 85        {
 86            assert_eq!(col_j, 0);
 87        }
 88        module.vmp_prepare_row(self, row_i, col_j, a);
 89    }
 90}
 91
 92impl GLWESwitchingKey<Vec<u8>, FFT64> {
 93    pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
 94        GGLWECiphertext::encrypt_sk_scratch_space(module, rank, size)
 95    }
 96
 97    pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, rank: usize, pk_size: usize) -> usize {
 98        GGLWECiphertext::encrypt_pk_scratch_space(module, rank, pk_size)
 99    }
100}
101
102impl<DataSelf> GLWESwitchingKey<DataSelf, FFT64>
103where
104    MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
105{
106    pub fn encrypt_sk<DataSkIn, DataSkOut>(
107        &mut self,
108        module: &Module<FFT64>,
109        sk_in: &SecretKey<DataSkIn>,
110        sk_out_dft: &SecretKeyFourier<DataSkOut, FFT64>,
111        source_xa: &mut Source,
112        source_xe: &mut Source,
113        sigma: f64,
114        scratch: &mut Scratch,
115    ) where
116        ScalarZnx<DataSkIn>: ScalarZnxToRef,
117        ScalarZnxDft<DataSkOut, FFT64>: ScalarZnxDftToRef<FFT64>,
118    {
119        self.0.encrypt_sk(
120            module,
121            &sk_in.data,
122            sk_out_dft,
123            source_xa,
124            source_xe,
125            sigma,
126            scratch,
127        );
128    }
129}
130
131impl GLWESwitchingKey<Vec<u8>, FFT64> {
132    pub fn keyswitch_scratch_space(
133        module: &Module<FFT64>,
134        res_size: usize,
135        lhs: usize,
136        rhs: usize,
137        rank_in: usize,
138        rank_out: usize,
139    ) -> usize {
140        <GGLWECiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_scratch_space(
141            module, res_size, lhs, rhs, rank_in, rank_out,
142        )
143    }
144
145    pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize, rank: usize) -> usize {
146        <GGLWECiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_inplace_scratch_space(
147            module, res_size, rhs, rank,
148        )
149    }
150
151    pub fn external_product_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize, rank: usize) -> usize {
152        <GGSWCiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_scratch_space(
153            module, res_size, lhs, rhs, rank, rank,
154        )
155    }
156
157    pub fn external_product_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize, rank: usize) -> usize {
158        <GGSWCiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_glwe_inplace_scratch_space(
159            module, res_size, rhs, rank,
160        )
161    }
162}
163
164impl<DataSelf> GLWESwitchingKey<DataSelf, FFT64>
165where
166    MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
167{
168    pub fn keyswitch<DataLhs, DataRhs>(
169        &mut self,
170        module: &Module<FFT64>,
171        lhs: &GLWESwitchingKey<DataLhs, FFT64>,
172        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
173        scratch: &mut base2k::Scratch,
174    ) where
175        MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
176        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
177    {
178        rhs.0
179            .prod_with_vec_glwe(module, &mut self.0, &lhs.0, scratch);
180    }
181
182    pub fn keyswitch_inplace<DataRhs>(
183        &mut self,
184        module: &Module<FFT64>,
185        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
186        scratch: &mut base2k::Scratch,
187    ) where
188        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
189    {
190        rhs.0
191            .prod_with_vec_glwe_inplace(module, &mut self.0, scratch);
192    }
193
194    pub fn external_product<DataLhs, DataRhs>(
195        &mut self,
196        module: &Module<FFT64>,
197        lhs: &GLWESwitchingKey<DataLhs, FFT64>,
198        rhs: &GGSWCiphertext<DataRhs, FFT64>,
199        scratch: &mut Scratch,
200    ) where
201        MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
202        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
203    {
204        rhs.prod_with_vec_glwe(module, &mut self.0, &lhs.0, scratch);
205    }
206
207    pub fn external_product_inplace<DataRhs>(
208        &mut self,
209        module: &Module<FFT64>,
210        rhs: &GGSWCiphertext<DataRhs, FFT64>,
211        scratch: &mut Scratch,
212    ) where
213        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
214    {
215        rhs.prod_with_vec_glwe_inplace(module, &mut self.0, scratch);
216    }
217}