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