1use poulpy_hal::layouts::{Backend, DataMut, 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
 54pub trait GGSWKeySwitch<BE: Backend>
 55where
 56    Self: GLWEKeyswitch<BE> + GGSWExpandRows<BE>,
 57{
 58    fn ggsw_keyswitch_tmp_bytes<R, A, K, T>(&self, res_infos: &R, a_infos: &A, key_infos: &K, tsk_infos: &T) -> usize
 59    where
 60        R: GGSWInfos,
 61        A: GGSWInfos,
 62        K: GGLWEInfos,
 63        T: GGLWEInfos,
 64    {
 65        assert_eq!(key_infos.rank_in(), key_infos.rank_out());
 66        assert_eq!(tsk_infos.rank_in(), tsk_infos.rank_out());
 67        assert_eq!(key_infos.rank_in(), tsk_infos.rank_in());
 68
 69        let rank: usize = key_infos.rank_out().into();
 70
 71        let size_out: usize = res_infos.k().div_ceil(res_infos.base2k()) as usize;
 72        let res_znx: usize = VecZnx::bytes_of(self.n(), rank + 1, size_out);
 73        let ci_dft: usize = self.bytes_of_vec_znx_dft(rank + 1, size_out);
 74        let ks: usize = self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos);
 75        let expand_rows: usize = self.ggsw_expand_rows_tmp_bytes(res_infos, tsk_infos);
 76        let res_dft: usize = self.bytes_of_vec_znx_dft(rank + 1, size_out);
 77
 78        if a_infos.base2k() == tsk_infos.base2k() {
 79            res_znx + ci_dft + (ks | expand_rows | res_dft)
 80        } else {
 81            let a_conv: usize = VecZnx::bytes_of(
 82                self.n(),
 83                1,
 84                res_infos.k().div_ceil(tsk_infos.base2k()) as usize,
 85            ) + self.vec_znx_normalize_tmp_bytes();
 86            res_znx + ci_dft + (a_conv | ks | expand_rows | res_dft)
 87        }
 88    }
 89
 90    fn ggsw_keyswitch<R, A, K, T>(&self, res: &mut R, a: &A, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
 91    where
 92        R: GGSWToMut,
 93        A: GGSWToRef,
 94        K: GLWESwitchingKeyPreparedToRef<BE>,
 95        T: TensorKeyPreparedToRef<BE>,
 96        Scratch<BE>: ScratchTakeCore<BE>,
 97    {
 98        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
 99        let a: &GGSW<&[u8]> = &a.to_ref();
100
101        assert_eq!(res.ggsw_layout(), a.ggsw_layout());
102
103        for row in 0..a.dnum().into() {
104            // Key-switch column 0, i.e.
105            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0s0' + a1s1' + a2s2') + M[i], a0, a1, a2)
106            self.glwe_keyswitch(&mut res.at_mut(row, 0), &a.at(row, 0), key, scratch);
107        }
108
109        self.ggsw_expand_row(res, tsk, scratch);
110    }
111
112    fn ggsw_keyswitch_inplace<R, K, T>(&self, res: &mut R, key: &K, tsk: &T, scratch: &mut Scratch<BE>)
113    where
114        R: GGSWToMut,
115        K: GLWESwitchingKeyPreparedToRef<BE>,
116        T: TensorKeyPreparedToRef<BE>,
117        Scratch<BE>: ScratchTakeCore<BE>,
118    {
119        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
120
121        for row in 0..res.dnum().into() {
122            // Key-switch column 0, i.e.
123            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0s0' + a1s1' + a2s2') + M[i], a0, a1, a2)
124            self.glwe_keyswitch_inplace(&mut res.at_mut(row, 0), key, scratch);
125        }
126
127        self.ggsw_expand_row(res, tsk, scratch);
128    }
129}