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::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<DataPt, DataSk>(
107        &mut self,
108        module: &Module<FFT64>,
109        pt: &ScalarZnx<DataPt>,
110        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
111        source_xa: &mut Source,
112        source_xe: &mut Source,
113        sigma: f64,
114        bound: f64,
115        scratch: &mut Scratch,
116    ) where
117        ScalarZnx<DataPt>: ScalarZnxToRef,
118        ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
119    {
120        self.0.encrypt_sk(
121            module, pt, sk_dft, source_xa, source_xe, sigma, bound, scratch,
122        );
123    }
124}
125
126impl GLWESwitchingKey<Vec<u8>, FFT64> {
127    pub fn keyswitch_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
128        <GGLWECiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_scratch_space(
129            module, res_size, lhs, rhs,
130        )
131    }
132
133    pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize) -> usize {
134        <GGLWECiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_inplace_scratch_space(
135            module, res_size, rhs,
136        )
137    }
138
139    pub fn external_product_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
140        <GGSWCiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_vec_glwe_scratch_space(
141            module, res_size, lhs, rhs,
142        )
143    }
144
145    pub fn external_product_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize) -> usize {
146        <GGSWCiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_glwe_inplace_scratch_space(
147            module, res_size, rhs,
148        )
149    }
150}
151
152impl<DataSelf> GLWESwitchingKey<DataSelf, FFT64>
153where
154    MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
155{
156    pub fn keyswitch<DataLhs, DataRhs>(
157        &mut self,
158        module: &Module<FFT64>,
159        lhs: &GLWESwitchingKey<DataLhs, FFT64>,
160        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
161        scratch: &mut base2k::Scratch,
162    ) where
163        MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
164        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
165    {
166        rhs.0
167            .prod_with_vec_glwe(module, &mut self.0, &lhs.0, scratch);
168    }
169
170    pub fn keyswitch_inplace<DataRhs>(
171        &mut self,
172        module: &Module<FFT64>,
173        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
174        scratch: &mut base2k::Scratch,
175    ) where
176        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
177    {
178        rhs.0
179            .prod_with_vec_glwe_inplace(module, &mut self.0, scratch);
180    }
181
182    pub fn external_product<DataLhs, DataRhs>(
183        &mut self,
184        module: &Module<FFT64>,
185        lhs: &GLWESwitchingKey<DataLhs, FFT64>,
186        rhs: &GGSWCiphertext<DataRhs, FFT64>,
187        scratch: &mut Scratch,
188    ) where
189        MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
190        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
191    {
192        rhs.prod_with_vec_glwe(module, &mut self.0, &lhs.0, scratch);
193    }
194
195    pub fn external_product_inplace<DataRhs>(
196        &mut self,
197        module: &Module<FFT64>,
198        rhs: &GGSWCiphertext<DataRhs, FFT64>,
199        scratch: &mut Scratch,
200    ) where
201        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
202    {
203        rhs.prod_with_vec_glwe_inplace(module, &mut self.0, scratch);
204    }
205}