1use backend::hal::{
  2    api::{ScratchAvailable, TakeVecZnxDft, ZnxZero},
  3    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::{
  7    AutomorphismKey, AutomorphismKeyExec, GLWECiphertext, GLWEKeyswitchFamily, GLWESwitchingKey, GLWESwitchingKeyExec, Infos,
  8};
  9
 10impl AutomorphismKey<Vec<u8>> {
 11    pub fn keyswitch_scratch_space<B: Backend>(
 12        module: &Module<B>,
 13        n: usize,
 14        basek: usize,
 15        k_out: usize,
 16        k_in: usize,
 17        k_ksk: usize,
 18        digits: usize,
 19        rank: usize,
 20    ) -> usize
 21    where
 22        Module<B>: GLWEKeyswitchFamily<B>,
 23    {
 24        GLWESwitchingKey::keyswitch_scratch_space(module, n, basek, k_out, k_in, k_ksk, digits, rank, rank)
 25    }
 26
 27    pub fn keyswitch_inplace_scratch_space<B: Backend>(
 28        module: &Module<B>,
 29        n: usize,
 30        basek: usize,
 31        k_out: usize,
 32        k_ksk: usize,
 33        digits: usize,
 34        rank: usize,
 35    ) -> usize
 36    where
 37        Module<B>: GLWEKeyswitchFamily<B>,
 38    {
 39        GLWESwitchingKey::keyswitch_inplace_scratch_space(module, n, basek, k_out, k_ksk, digits, rank)
 40    }
 41}
 42
 43impl<DataSelf: DataMut> AutomorphismKey<DataSelf> {
 44    pub fn keyswitch<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
 45        &mut self,
 46        module: &Module<B>,
 47        lhs: &AutomorphismKey<DataLhs>,
 48        rhs: &GLWESwitchingKeyExec<DataRhs, B>,
 49        scratch: &mut Scratch<B>,
 50    ) where
 51        Module<B>: GLWEKeyswitchFamily<B>,
 52        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
 53    {
 54        self.key.keyswitch(module, &lhs.key, rhs, scratch);
 55    }
 56
 57    pub fn keyswitch_inplace<DataRhs: DataRef, B: Backend>(
 58        &mut self,
 59        module: &Module<B>,
 60        rhs: &AutomorphismKeyExec<DataRhs, B>,
 61        scratch: &mut Scratch<B>,
 62    ) where
 63        Module<B>: GLWEKeyswitchFamily<B>,
 64        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
 65    {
 66        self.key.keyswitch_inplace(module, &rhs.key, scratch);
 67    }
 68}
 69
 70impl GLWESwitchingKey<Vec<u8>> {
 71    pub fn keyswitch_scratch_space<B: Backend>(
 72        module: &Module<B>,
 73        n: usize,
 74        basek: usize,
 75        k_out: usize,
 76        k_in: usize,
 77        k_ksk: usize,
 78        digits: usize,
 79        rank_in: usize,
 80        rank_out: usize,
 81    ) -> usize
 82    where
 83        Module<B>: GLWEKeyswitchFamily<B>,
 84    {
 85        GLWECiphertext::keyswitch_scratch_space(
 86            module, n, basek, k_out, k_in, k_ksk, digits, rank_in, rank_out,
 87        )
 88    }
 89
 90    pub fn keyswitch_inplace_scratch_space<B: Backend>(
 91        module: &Module<B>,
 92        n: usize,
 93        basek: usize,
 94        k_out: usize,
 95        k_ksk: usize,
 96        digits: usize,
 97        rank: usize,
 98    ) -> usize
 99    where
100        Module<B>: GLWEKeyswitchFamily<B>,
101    {
102        GLWECiphertext::keyswitch_inplace_scratch_space(module, n, basek, k_out, k_ksk, digits, rank)
103    }
104}
105
106impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
107    pub fn keyswitch<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
108        &mut self,
109        module: &Module<B>,
110        lhs: &GLWESwitchingKey<DataLhs>,
111        rhs: &GLWESwitchingKeyExec<DataRhs, B>,
112        scratch: &mut Scratch<B>,
113    ) where
114        Module<B>: GLWEKeyswitchFamily<B>,
115        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
116    {
117        #[cfg(debug_assertions)]
118        {
119            assert_eq!(
120                self.rank_in(),
121                lhs.rank_in(),
122                "ksk_out input rank: {} != ksk_in input rank: {}",
123                self.rank_in(),
124                lhs.rank_in()
125            );
126            assert_eq!(
127                lhs.rank_out(),
128                rhs.rank_in(),
129                "ksk_in output rank: {} != ksk_apply input rank: {}",
130                self.rank_out(),
131                rhs.rank_in()
132            );
133            assert_eq!(
134                self.rank_out(),
135                rhs.rank_out(),
136                "ksk_out output rank: {} != ksk_apply output rank: {}",
137                self.rank_out(),
138                rhs.rank_out()
139            );
140        }
141
142        (0..self.rank_in()).for_each(|col_i| {
143            (0..self.rows()).for_each(|row_j| {
144                self.at_mut(row_j, col_i)
145                    .keyswitch(module, &lhs.at(row_j, col_i), rhs, scratch);
146            });
147        });
148
149        (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
150            (0..self.rank_in()).for_each(|col_j| {
151                self.at_mut(row_i, col_j).data.zero();
152            });
153        });
154    }
155
156    pub fn keyswitch_inplace<DataRhs: DataRef, B: Backend>(
157        &mut self,
158        module: &Module<B>,
159        rhs: &GLWESwitchingKeyExec<DataRhs, B>,
160        scratch: &mut Scratch<B>,
161    ) where
162        Module<B>: GLWEKeyswitchFamily<B>,
163        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
164    {
165        #[cfg(debug_assertions)]
166        {
167            assert_eq!(
168                self.rank_out(),
169                rhs.rank_out(),
170                "ksk_out output rank: {} != ksk_apply output rank: {}",
171                self.rank_out(),
172                rhs.rank_out()
173            );
174        }
175
176        (0..self.rank_in()).for_each(|col_i| {
177            (0..self.rows()).for_each(|row_j| {
178                self.at_mut(row_j, col_i)
179                    .keyswitch_inplace(module, rhs, scratch)
180            });
181        });
182    }
183}