1use poulpy_hal::{
  2    api::VecZnxAutomorphism,
  3    layouts::{Backend, DataMut, GaloisElement, Module, Scratch},
  4};
  5
  6use crate::{
  7    ScratchTakeCore,
  8    automorphism::glwe_ct::GLWEAutomorphism,
  9    layouts::{
 10        AutomorphismKey, GGLWE, GGLWEInfos, GGLWEPreparedToRef, GGLWEToMut, GGLWEToRef, GLWE,
 11        prepared::{GetAutomorphismGaloisElement, SetAutomorphismGaloisElement},
 12    },
 13};
 14
 15impl AutomorphismKey<Vec<u8>> {
 16    pub fn automorphism_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
 17    where
 18        R: GGLWEInfos,
 19        A: GGLWEInfos,
 20        K: GGLWEInfos,
 21        M: AutomorphismKeyAutomorphism<BE>,
 22    {
 23        module.automorphism_key_automorphism_tmp_bytes(res_infos, a_infos, key_infos)
 24    }
 25}
 26
 27impl<DataSelf: DataMut> AutomorphismKey<DataSelf> {
 28    pub fn automorphism<A, K, M, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
 29    where
 30        A: GGLWEToRef + GetAutomorphismGaloisElement + GGLWEInfos,
 31        K: GGLWEPreparedToRef<BE> + GetAutomorphismGaloisElement + GGLWEInfos,
 32        Scratch<BE>: ScratchTakeCore<BE>,
 33        M: AutomorphismKeyAutomorphism<BE>,
 34    {
 35        module.automorphism_key_automorphism(self, a, key, scratch);
 36    }
 37
 38    pub fn automorphism_inplace<K, M, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
 39    where
 40        K: GGLWEPreparedToRef<BE> + GetAutomorphismGaloisElement + GGLWEInfos,
 41        Scratch<BE>: ScratchTakeCore<BE>,
 42        M: AutomorphismKeyAutomorphism<BE>,
 43    {
 44        module.automorphism_key_automorphism_inplace(self, key, scratch);
 45    }
 46}
 47
 48impl<BE: Backend> AutomorphismKeyAutomorphism<BE> for Module<BE> where
 49    Self: GaloisElement + GLWEAutomorphism<BE> + VecZnxAutomorphism
 50{
 51}
 52
 53pub trait AutomorphismKeyAutomorphism<BE: Backend>
 54where
 55    Self: GaloisElement + GLWEAutomorphism<BE> + VecZnxAutomorphism,
 56{
 57    fn automorphism_key_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
 58    where
 59        R: GGLWEInfos,
 60        A: GGLWEInfos,
 61        K: GGLWEInfos,
 62    {
 63        self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
 64    }
 65
 66    fn automorphism_key_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
 67    where
 68        R: GGLWEToMut + SetAutomorphismGaloisElement + GGLWEInfos,
 69        A: GGLWEToRef + GetAutomorphismGaloisElement + GGLWEInfos,
 70        K: GGLWEPreparedToRef<BE> + GetAutomorphismGaloisElement + GGLWEInfos,
 71        Scratch<BE>: ScratchTakeCore<BE>,
 72    {
 73        assert!(
 74            res.dnum().as_u32() <= a.dnum().as_u32(),
 75            "res dnum: {} > a dnum: {}",
 76            res.dnum(),
 77            a.dnum()
 78        );
 79
 80        assert_eq!(
 81            res.dsize(),
 82            a.dsize(),
 83            "res dnum: {} != a dnum: {}",
 84            res.dsize(),
 85            a.dsize()
 86        );
 87
 88        let cols_out: usize = (key.rank_out() + 1).into();
 89        let cols_in: usize = key.rank_in().into();
 90
 91        let p: i64 = a.p();
 92        let p_inv: i64 = self.galois_element_inv(p);
 93
 94        {
 95            let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
 96            let a: &GGLWE<&[u8]> = &a.to_ref();
 97
 98            for row in 0..res.dnum().as_usize() {
 99                for col in 0..cols_in {
100                    let mut res_tmp: GLWE<&mut [u8]> = res.at_mut(row, col);
101                    let a_ct: GLWE<&[u8]> = a.at(row, col);
102
103                    // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
104                    for i in 0..cols_out {
105                        self.vec_znx_automorphism(p, res_tmp.data_mut(), i, &a_ct.data, i);
106                    }
107
108                    // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
109                    self.glwe_keyswitch_inplace(&mut res_tmp, key, scratch);
110
111                    // Applies back the automorphism X^{-k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) to (-pi^{-1}_{k'+k}(s)a + s, a)
112                    (0..cols_out).for_each(|i| {
113                        self.vec_znx_automorphism_inplace(p_inv, res_tmp.data_mut(), i, scratch);
114                    });
115                }
116            }
117        }
118
119        res.set_p((p * key.p()) % (self.cyclotomic_order() as i64));
120    }
121
122    fn automorphism_key_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
123    where
124        R: GGLWEToMut + SetAutomorphismGaloisElement + GetAutomorphismGaloisElement + GGLWEInfos,
125        K: GGLWEPreparedToRef<BE> + GetAutomorphismGaloisElement + GGLWEInfos,
126        Scratch<BE>: ScratchTakeCore<BE>,
127    {
128        assert_eq!(
129            res.rank(),
130            key.rank(),
131            "key rank: {} != key rank: {}",
132            res.rank(),
133            key.rank()
134        );
135
136        let cols_out: usize = (key.rank_out() + 1).into();
137        let cols_in: usize = key.rank_in().into();
138        let p: i64 = res.p();
139        let p_inv: i64 = self.galois_element_inv(p);
140
141        {
142            let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
143            for row in 0..res.dnum().as_usize() {
144                for col in 0..cols_in {
145                    let mut res_tmp: GLWE<&mut [u8]> = res.at_mut(row, col);
146
147                    // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
148                    for i in 0..cols_out {
149                        self.vec_znx_automorphism_inplace(p, res_tmp.data_mut(), i, scratch);
150                    }
151
152                    // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
153                    self.glwe_keyswitch_inplace(&mut res_tmp, key, scratch);
154
155                    // Applies back the automorphism X^{-k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) to (-pi^{-1}_{k'+k}(s)a + s, a)
156                    for i in 0..cols_out {
157                        self.vec_znx_automorphism_inplace(p_inv, res_tmp.data_mut(), i, scratch);
158                    }
159                }
160            }
161        }
162
163        res.set_p((res.p() * key.p()) % (self.cyclotomic_order() as i64));
164    }
165}