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