1use backend::hal::{
  2    api::{ScratchAvailable, TakeVecZnxDft, VecZnxAutomorphism, VecZnxAutomorphismInplace, ZnxZero},
  3    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::{
  7    layouts::{GGLWEAutomorphismKey, GLWECiphertext, Infos, prepared::GGLWEAutomorphismKeyExec},
  8    trait_families::GLWEKeyswitchFamily,
  9};
 10
 11impl GGLWEAutomorphismKey<Vec<u8>> {
 12    pub fn automorphism_scratch_space<B: Backend>(
 13        module: &Module<B>,
 14        n: usize,
 15        basek: usize,
 16        k_out: usize,
 17        k_in: usize,
 18        k_ksk: usize,
 19        digits: usize,
 20        rank: usize,
 21    ) -> usize
 22    where
 23        Module<B>: GLWEKeyswitchFamily<B>,
 24    {
 25        GLWECiphertext::keyswitch_scratch_space(module, n, basek, k_out, k_in, k_ksk, digits, rank, rank)
 26    }
 27
 28    pub fn automorphism_inplace_scratch_space<B: Backend>(
 29        module: &Module<B>,
 30        n: usize,
 31        basek: usize,
 32        k_out: usize,
 33        k_ksk: usize,
 34        digits: usize,
 35        rank: usize,
 36    ) -> usize
 37    where
 38        Module<B>: GLWEKeyswitchFamily<B>,
 39    {
 40        GGLWEAutomorphismKey::automorphism_scratch_space(module, n, basek, k_out, k_out, k_ksk, digits, rank)
 41    }
 42}
 43
 44impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
 45    pub fn automorphism<'a, DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
 46        &mut self,
 47        module: &Module<B>,
 48        lhs: &GGLWEAutomorphismKey<DataLhs>,
 49        rhs: &GGLWEAutomorphismKeyExec<DataRhs, B>,
 50        scratch: &mut Scratch<B>,
 51    ) where
 52        Module<B>: GLWEKeyswitchFamily<B> + VecZnxAutomorphism + VecZnxAutomorphismInplace,
 53        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
 54    {
 55        #[cfg(debug_assertions)]
 56        {
 57            assert_eq!(
 58                self.rank_in(),
 59                lhs.rank_in(),
 60                "ksk_out input rank: {} != ksk_in input rank: {}",
 61                self.rank_in(),
 62                lhs.rank_in()
 63            );
 64            assert_eq!(
 65                lhs.rank_out(),
 66                rhs.rank_in(),
 67                "ksk_in output rank: {} != ksk_apply input rank: {}",
 68                self.rank_out(),
 69                rhs.rank_in()
 70            );
 71            assert_eq!(
 72                self.rank_out(),
 73                rhs.rank_out(),
 74                "ksk_out output rank: {} != ksk_apply output rank: {}",
 75                self.rank_out(),
 76                rhs.rank_out()
 77            );
 78            assert!(
 79                self.k() <= lhs.k(),
 80                "output k={} cannot be greater than input k={}",
 81                self.k(),
 82                lhs.k()
 83            )
 84        }
 85
 86        let cols_out: usize = rhs.rank_out() + 1;
 87
 88        let p: i64 = lhs.p();
 89        let p_inv = module.galois_element_inv(p);
 90
 91        (0..self.rank_in()).for_each(|col_i| {
 92            (0..self.rows()).for_each(|row_j| {
 93                let mut res_ct: GLWECiphertext<&mut [u8]> = self.at_mut(row_j, col_i);
 94                let lhs_ct: GLWECiphertext<&[u8]> = lhs.at(row_j, col_i);
 95
 96                // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
 97                (0..cols_out).for_each(|i| {
 98                    module.vec_znx_automorphism(lhs.p(), &mut res_ct.data, i, &lhs_ct.data, i);
 99                });
100
101                // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
102                res_ct.keyswitch_inplace(module, &rhs.key, scratch);
103
104                // Applies back the automorphism X^{-k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) to (-pi^{-1}_{k'+k}(s)a + s, a)
105                (0..cols_out).for_each(|i| {
106                    module.vec_znx_automorphism_inplace(p_inv, &mut res_ct.data, i);
107                });
108            });
109        });
110
111        (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
112            (0..self.rank_in()).for_each(|col_j| {
113                self.at_mut(row_i, col_j).data.zero();
114            });
115        });
116
117        self.p = (lhs.p * rhs.p) % (module.cyclotomic_order() as i64);
118    }
119
120    pub fn automorphism_inplace<DataRhs: DataRef, B: Backend>(
121        &mut self,
122        module: &Module<B>,
123        rhs: &GGLWEAutomorphismKeyExec<DataRhs, B>,
124        scratch: &mut Scratch<B>,
125    ) where
126        Module<B>: GLWEKeyswitchFamily<B> + VecZnxAutomorphism + VecZnxAutomorphismInplace,
127        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
128    {
129        unsafe {
130            let self_ptr: *mut GGLWEAutomorphismKey<DataSelf> = self as *mut GGLWEAutomorphismKey<DataSelf>;
131            self.automorphism(&module, &*self_ptr, rhs, scratch);
132        }
133    }
134}