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