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