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