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
 90        let p: i64 = a.p();
 91        let p_inv: i64 = self.galois_element_inv(p);
 92
 93        {
 94            let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
 95            let a: &GGLWE<&[u8]> = &a.to_ref();
 96
 97            for row in 0..res.dnum().as_usize() {
 98                for col in 0..cols_out {
 99                    let mut res_tmp: GLWE<&mut [u8]> = res.at_mut(row, col);
100                    let a_ct: GLWE<&[u8]> = a.at(row, col);
101
102                    // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
103                    for i in 0..cols_out {
104                        self.vec_znx_automorphism(p, res_tmp.data_mut(), i, &a_ct.data, i);
105                    }
106
107                    // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
108                    self.glwe_keyswitch_inplace(&mut res_tmp, key, scratch);
109
110                    // Applies back the automorphism X^{-k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) to (-pi^{-1}_{k'+k}(s)a + s, a)
111                    (0..cols_out).for_each(|i| {
112                        self.vec_znx_automorphism_inplace(p_inv, res_tmp.data_mut(), i, scratch);
113                    });
114                }
115            }
116        }
117
118        res.set_p((p * key.p()) % (self.cyclotomic_order() as i64));
119    }
120
121    fn automorphism_key_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
122    where
123        R: GGLWEToMut + SetAutomorphismGaloisElement + GetAutomorphismGaloisElement + GGLWEInfos,
124        K: GGLWEPreparedToRef<BE> + GetAutomorphismGaloisElement + GGLWEInfos,
125        Scratch<BE>: ScratchTakeCore<BE>,
126    {
127        assert_eq!(
128            res.rank(),
129            key.rank(),
130            "key rank: {} != key rank: {}",
131            res.rank(),
132            key.rank()
133        );
134
135        let cols_out: usize = (key.rank_out() + 1).into();
136        let p: i64 = res.p();
137        let p_inv: i64 = self.galois_element_inv(p);
138
139        {
140            let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
141            for row in 0..res.dnum().as_usize() {
142                for col in 0..cols_out {
143                    let mut res_tmp: GLWE<&mut [u8]> = res.at_mut(row, col);
144
145                    // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
146                    for i in 0..cols_out {
147                        self.vec_znx_automorphism_inplace(p_inv, res_tmp.data_mut(), i, scratch);
148                    }
149
150                    // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
151                    self.glwe_keyswitch_inplace(&mut res_tmp, key, scratch);
152
153                    // Applies back the automorphism X^{-k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) to (-pi^{-1}_{k'+k}(s)a + s, a)
154                    for i in 0..cols_out {
155                        self.vec_znx_automorphism_inplace(p_inv, res_tmp.data_mut(), i, scratch);
156                    }
157                }
158            }
159        }
160
161        res.set_p((res.p() * key.p()) % (self.cyclotomic_order() as i64));
162    }
163}