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