1use poulpy_hal::layouts::{Backend, DataMut, Module, Scratch, VecZnx};
  2
  3use crate::{
  4    GGSWExpandRows, ScratchTakeCore,
  5    keyswitching::glwe_ct::GLWEKeyswitch,
  6    layouts::{
  7        GGLWEInfos, GGSW, GGSWInfos, GGSWToMut, GGSWToRef,
  8        prepared::{GLWESwitchingKeyPreparedToRef, TensorKeyPreparedToRef},
  9    },
 10};
 11
 12impl GGSW<Vec<u8>> {
 13    pub fn keyswitch_tmp_bytes<R, A, K, T, M, BE: Backend>(
 14        module: &M,
 15        res_infos: &R,
 16        a_infos: &A,
 17        key_infos: &K,
 18        tsk_infos: &T,
 19    ) -> usize
 20    where
 21        R: GGSWInfos,
 22        A: GGSWInfos,
 23        K: GGLWEInfos,
 24        T: GGLWEInfos,
 25        M: GGSWKeyswitch<BE>,
 26    {
 27        module.ggsw_keyswitch_tmp_bytes(res_infos, a_infos, key_infos, tsk_infos)
 28    }
 29}
 30
 31impl<D: DataMut> GGSW<D> {
 32    pub fn keyswitch<M, A, K, T, BE: Backend>(&mut self, module: &M, a: &A, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
 33    where
 34        A: GGSWToRef,
 35        K: GLWESwitchingKeyPreparedToRef<BE>,
 36        T: TensorKeyPreparedToRef<BE>,
 37        Scratch<BE>: ScratchTakeCore<BE>,
 38        M: GGSWKeyswitch<BE>,
 39    {
 40        module.ggsw_keyswitch(self, a, key, tsk, scratch);
 41    }
 42
 43    pub fn keyswitch_inplace<M, K, T, BE: Backend>(&mut self, module: &M, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
 44    where
 45        K: GLWESwitchingKeyPreparedToRef<BE>,
 46        T: TensorKeyPreparedToRef<BE>,
 47        Scratch<BE>: ScratchTakeCore<BE>,
 48        M: GGSWKeyswitch<BE>,
 49    {
 50        module.ggsw_keyswitch_inplace(self, key, tsk, scratch);
 51    }
 52}
 53
 54impl<BE: Backend> GGSWKeyswitch<BE> for Module<BE> where Self: GLWEKeyswitch<BE> + GGSWExpandRows<BE> {}
 55
 56pub trait GGSWKeyswitch<BE: Backend>
 57where
 58    Self: GLWEKeyswitch<BE> + GGSWExpandRows<BE>,
 59{
 60    fn ggsw_keyswitch_tmp_bytes<R, A, K, T>(&self, res_infos: &R, a_infos: &A, key_infos: &K, tsk_infos: &T) -> usize
 61    where
 62        R: GGSWInfos,
 63        A: GGSWInfos,
 64        K: GGLWEInfos,
 65        T: GGLWEInfos,
 66    {
 67        assert_eq!(key_infos.rank_in(), key_infos.rank_out());
 68        assert_eq!(tsk_infos.rank_in(), tsk_infos.rank_out());
 69        assert_eq!(key_infos.rank_in(), tsk_infos.rank_in());
 70
 71        let rank: usize = key_infos.rank_out().into();
 72
 73        let size_out: usize = res_infos.k().div_ceil(res_infos.base2k()) as usize;
 74        let res_znx: usize = VecZnx::bytes_of(self.n(), rank + 1, size_out);
 75        let ci_dft: usize = self.bytes_of_vec_znx_dft(rank + 1, size_out);
 76        let ks: usize = self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos);
 77        let expand_rows: usize = self.ggsw_expand_rows_tmp_bytes(res_infos, tsk_infos);
 78        let res_dft: usize = self.bytes_of_vec_znx_dft(rank + 1, size_out);
 79
 80        if a_infos.base2k() == tsk_infos.base2k() {
 81            res_znx + ci_dft + (ks | expand_rows | res_dft)
 82        } else {
 83            let a_conv: usize = VecZnx::bytes_of(
 84                self.n(),
 85                1,
 86                res_infos.k().div_ceil(tsk_infos.base2k()) as usize,
 87            ) + self.vec_znx_normalize_tmp_bytes();
 88            res_znx + ci_dft + (a_conv | ks | expand_rows | res_dft)
 89        }
 90    }
 91
 92    fn ggsw_keyswitch<R, A, K, T>(&self, res: &mut R, a: &A, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
 93    where
 94        R: GGSWToMut,
 95        A: GGSWToRef,
 96        K: GLWESwitchingKeyPreparedToRef<BE>,
 97        T: TensorKeyPreparedToRef<BE>,
 98        Scratch<BE>: ScratchTakeCore<BE>,
 99    {
100        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
101        let a: &GGSW<&[u8]> = &a.to_ref();
102
103        assert_eq!(res.ggsw_layout(), a.ggsw_layout());
104
105        for row in 0..a.dnum().into() {
106            // Key-switch column 0, i.e.
107            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0s0' + a1s1' + a2s2') + M[i], a0, a1, a2)
108            self.glwe_keyswitch(&mut res.at_mut(row, 0), &a.at(row, 0), key, scratch);
109        }
110
111        self.ggsw_expand_row(res, tsk, scratch);
112    }
113
114    fn ggsw_keyswitch_inplace<R, K, T>(&self, res: &mut R, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
115    where
116        R: GGSWToMut,
117        K: GLWESwitchingKeyPreparedToRef<BE>,
118        T: TensorKeyPreparedToRef<BE>,
119        Scratch<BE>: ScratchTakeCore<BE>,
120    {
121        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
122
123        for row in 0..res.dnum().into() {
124            // Key-switch column 0, i.e.
125            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0s0' + a1s1' + a2s2') + M[i], a0, a1, a2)
126            self.glwe_keyswitch_inplace(&mut res.at_mut(row, 0), key, scratch);
127        }
128
129        self.ggsw_expand_row(res, tsk, scratch);
130    }
131}