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