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        if res_infos.glwe_layout() == a_infos.glwe_layout() {
 59            self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
 60        } else {
 61            self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos) + GLWE::bytes_of_from_infos(a_infos)
 62        }
 63    }
 64
 65    fn glwe_automorphism_key_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
 66    where
 67        R: GGLWEToMut + SetGaloisElement + GGLWEInfos,
 68        A: GGLWEToRef + GetGaloisElement + GGLWEInfos,
 69        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
 70    {
 71        assert!(
 72            res.dnum().as_u32() <= a.dnum().as_u32(),
 73            "res dnum: {} > a dnum: {}",
 74            res.dnum(),
 75            a.dnum()
 76        );
 77
 78        assert_eq!(res.dsize(), a.dsize(), "res dnum: {} != a dnum: {}", res.dsize(), a.dsize());
 79
 80        assert_eq!(res.base2k(), a.base2k());
 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        let same_layout: bool = res.glwe_layout() == a.glwe_layout();
 89
 90        {
 91            let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
 92            let a: &GGLWE<&[u8]> = &a.to_ref();
 93
 94            for row in 0..res.dnum().as_usize() {
 95                for col in 0..cols_in {
 96                    let mut res_tmp: GLWE<&mut [u8]> = res.at_mut(row, col);
 97                    let a_ct: GLWE<&[u8]> = a.at(row, col);
 98
 99                    if same_layout {
100                        // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
101                        for i in 0..cols_out {
102                            self.vec_znx_automorphism(p, res_tmp.data_mut(), i, &a_ct.data, i);
103                        }
104
105                        // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
106                        self.glwe_keyswitch_inplace(&mut res_tmp, key, scratch);
107                    } else {
108                        let (mut tmp_glwe, scratch_1) = scratch.take_glwe(a);
109
110                        // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
111                        for i in 0..cols_out {
112                            self.vec_znx_automorphism(p, tmp_glwe.data_mut(), i, &a_ct.data, i);
113                        }
114
115                        // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
116                        self.glwe_keyswitch(&mut res_tmp, &tmp_glwe, key, scratch_1);
117                    }
118
119                    // Applies back the automorphism X^{-k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) to (-pi^{-1}_{k'+k}(s)a + s, a)
120                    for i in 0..cols_out {
121                        self.vec_znx_automorphism_inplace(p_inv, res_tmp.data_mut(), i, scratch);
122                    }
123                }
124            }
125        }
126
127        res.set_p((p * key.p()) % self.cyclotomic_order());
128    }
129
130    fn glwe_automorphism_key_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
131    where
132        R: GGLWEToMut + SetGaloisElement + GetGaloisElement + GGLWEInfos,
133        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
134        Scratch<BE>: ScratchTakeCore<BE>,
135    {
136        assert_eq!(res.rank(), key.rank(), "key rank: {} != key rank: {}", res.rank(), key.rank());
137
138        let cols_out: usize = (key.rank_out() + 1).into();
139        let cols_in: usize = key.rank_in().into();
140        let p: i64 = res.p();
141        let p_inv: i64 = self.galois_element_inv(p);
142
143        {
144            let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
145            for row in 0..res.dnum().as_usize() {
146                for col in 0..cols_in {
147                    let mut res_tmp: GLWE<&mut [u8]> = res.at_mut(row, col);
148
149                    // Reverts the automorphism X^{-k}: (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
150                    for i in 0..cols_out {
151                        self.vec_znx_automorphism_inplace(p, res_tmp.data_mut(), i, scratch);
152                    }
153
154                    // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
155                    self.glwe_keyswitch_inplace(&mut res_tmp, key, scratch);
156
157                    // Applies back the automorphism X^{-k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) to (-pi^{-1}_{k'+k}(s)a + s, a)
158                    for i in 0..cols_out {
159                        self.vec_znx_automorphism_inplace(p_inv, res_tmp.data_mut(), i, scratch);
160                    }
161                }
162            }
163        }
164
165        res.set_p((res.p() * key.p()) % self.cyclotomic_order());
166    }
167}
168
169pub trait GLWEAutomorphismKeyAutomorphism<BE: Backend> {
170    fn glwe_automorphism_key_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
171    where
172        R: GGLWEInfos,
173        A: GGLWEInfos,
174        K: GGLWEInfos;
175
176    fn glwe_automorphism_key_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
177    where
178        R: GGLWEToMut + SetGaloisElement + GGLWEInfos,
179        A: GGLWEToRef + GetGaloisElement + GGLWEInfos,
180        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos;
181
182    fn glwe_automorphism_key_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
183    where
184        R: GGLWEToMut + SetGaloisElement + GetGaloisElement + GGLWEInfos,
185        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos;
186}