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