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