1use poulpy_hal::layouts::{Backend, DataMut, Module, Scratch};
  2
  3use crate::{
  4    ScratchTakeCore,
  5    keyswitching::glwe_ct::GLWEKeySwitch,
  6    layouts::{
  7        AutomorphismKey, AutomorphismKeyToRef, GGLWE, GGLWEInfos, GGLWEToMut, GGLWEToRef, GLWESwitchingKey,
  8        GLWESwitchingKeyToRef,
  9        prepared::{GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedToRef},
 10    },
 11};
 12
 13impl AutomorphismKey<Vec<u8>> {
 14    pub fn keyswitch_inplace_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
 15    where
 16        R: GGLWEInfos,
 17        A: GGLWEInfos,
 18        K: GGLWEInfos,
 19        M: GGLWEKeySwitch<BE>,
 20    {
 21        module.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
 22    }
 23}
 24
 25impl<DataSelf: DataMut> AutomorphismKey<DataSelf> {
 26    pub fn keyswitch<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
 27    where
 28        A: AutomorphismKeyToRef,
 29        B: GLWESwitchingKeyPreparedToRef<BE>,
 30        Scratch<BE>: ScratchTakeCore<BE>,
 31        M: GGLWEKeySwitch<BE>,
 32    {
 33        module.gglwe_keyswitch(&mut self.key.key, &a.to_ref().key.key, b, scratch);
 34    }
 35
 36    pub fn keyswitch_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
 37    where
 38        A: GLWESwitchingKeyPreparedToRef<BE>,
 39        Scratch<BE>: ScratchTakeCore<BE>,
 40        M: GGLWEKeySwitch<BE>,
 41    {
 42        module.gglwe_keyswitch_inplace(&mut self.key.key, a, scratch);
 43    }
 44}
 45
 46impl GLWESwitchingKey<Vec<u8>> {
 47    pub fn keyswitch_inplace_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
 48    where
 49        R: GGLWEInfos,
 50        A: GGLWEInfos,
 51        K: GGLWEInfos,
 52        M: GGLWEKeySwitch<BE>,
 53    {
 54        module.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
 55    }
 56}
 57
 58impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
 59    pub fn keyswitch<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
 60    where
 61        A: GLWESwitchingKeyToRef,
 62        B: GLWESwitchingKeyPreparedToRef<BE>,
 63        Scratch<BE>: ScratchTakeCore<BE>,
 64        M: GGLWEKeySwitch<BE>,
 65    {
 66        module.gglwe_keyswitch(&mut self.key, &a.to_ref().key, b, scratch);
 67    }
 68
 69    pub fn keyswitch_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
 70    where
 71        A: GLWESwitchingKeyPreparedToRef<BE>,
 72        Scratch<BE>: ScratchTakeCore<BE>,
 73        M: GGLWEKeySwitch<BE>,
 74    {
 75        module.gglwe_keyswitch_inplace(&mut self.key, a, scratch);
 76    }
 77}
 78
 79impl GGLWE<Vec<u8>> {
 80    pub fn keyswitch_inplace_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
 81    where
 82        R: GGLWEInfos,
 83        A: GGLWEInfos,
 84        K: GGLWEInfos,
 85        M: GGLWEKeySwitch<BE>,
 86    {
 87        module.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
 88    }
 89}
 90
 91impl<DataSelf: DataMut> GGLWE<DataSelf> {
 92    pub fn keyswitch<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
 93    where
 94        A: GGLWEToRef,
 95        B: GLWESwitchingKeyPreparedToRef<BE>,
 96        Scratch<BE>: ScratchTakeCore<BE>,
 97        M: GGLWEKeySwitch<BE>,
 98    {
 99        module.gglwe_keyswitch(self, a, b, scratch);
100    }
101
102    pub fn keyswitch_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
103    where
104        A: GLWESwitchingKeyPreparedToRef<BE>,
105        Scratch<BE>: ScratchTakeCore<BE>,
106        M: GGLWEKeySwitch<BE>,
107    {
108        module.gglwe_keyswitch_inplace(self, a, scratch);
109    }
110}
111
112impl<BE: Backend> GGLWEKeySwitch<BE> for Module<BE> where Self: GLWEKeySwitch<BE> {}
113
114pub trait GGLWEKeySwitch<BE: Backend>
115where
116    Self: GLWEKeySwitch<BE>,
117{
118    fn gglwe_keyswitch_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
119    where
120        R: GGLWEInfos,
121        A: GGLWEInfos,
122        K: GGLWEInfos,
123    {
124        self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
125    }
126
127    fn gglwe_keyswitch<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch<BE>)
128    where
129        R: GGLWEToMut,
130        A: GGLWEToRef,
131        B: GLWESwitchingKeyPreparedToRef<BE>,
132        Scratch<BE>: ScratchTakeCore<BE>,
133    {
134        let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
135        let a: &GGLWE<&[u8]> = &a.to_ref();
136        let b: &GLWESwitchingKeyPrepared<&[u8], BE> = &b.to_ref();
137
138        assert_eq!(
139            res.rank_in(),
140            a.rank_in(),
141            "res input rank: {} != a input rank: {}",
142            res.rank_in(),
143            a.rank_in()
144        );
145        assert_eq!(
146            a.rank_out(),
147            b.rank_in(),
148            "res output rank: {} != b input rank: {}",
149            a.rank_out(),
150            b.rank_in()
151        );
152        assert_eq!(
153            res.rank_out(),
154            b.rank_out(),
155            "res output rank: {} != b output rank: {}",
156            res.rank_out(),
157            b.rank_out()
158        );
159        assert!(
160            res.dnum() <= a.dnum(),
161            "res.dnum()={} > a.dnum()={}",
162            res.dnum(),
163            a.dnum()
164        );
165        assert_eq!(
166            res.dsize(),
167            a.dsize(),
168            "res dsize: {} != a dsize: {}",
169            res.dsize(),
170            a.dsize()
171        );
172
173        for row in 0..res.dnum().into() {
174            for col in 0..res.rank_in().into() {
175                self.glwe_keyswitch(&mut res.at_mut(row, col), &a.at(row, col), b, scratch);
176            }
177        }
178    }
179
180    fn gglwe_keyswitch_inplace<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
181    where
182        R: GGLWEToMut,
183        A: GLWESwitchingKeyPreparedToRef<BE>,
184        Scratch<BE>: ScratchTakeCore<BE>,
185    {
186        let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
187        let a: &GLWESwitchingKeyPrepared<&[u8], BE> = &a.to_ref();
188
189        assert_eq!(
190            res.rank_out(),
191            a.rank_out(),
192            "res output rank: {} != a output rank: {}",
193            res.rank_out(),
194            a.rank_out()
195        );
196
197        for row in 0..res.dnum().into() {
198            for col in 0..res.rank_in().into() {
199                self.glwe_keyswitch_inplace(&mut res.at_mut(row, col), a, scratch);
200            }
201        }
202    }
203}
204
205impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {}