poulpy_core/automorphism/
gglwe_atk.rs

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        GGLWE, GGLWEInfos, GGLWEPreparedToRef, GGLWEToMut, GGLWEToRef, GLWE, GLWEAutomorphismKey, GetGaloisElement,
11        SetGaloisElement,
12    },
13};
14
15impl GLWEAutomorphismKey<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: GLWEAutomorphismKeyAutomorphism<BE>,
22    {
23        module.glwe_automorphism_key_automorphism_tmp_bytes(res_infos, a_infos, key_infos)
24    }
25}
26
27impl<DataSelf: DataMut> GLWEAutomorphismKey<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 + GetGaloisElement + GGLWEInfos,
31        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
32        Scratch<BE>: ScratchTakeCore<BE>,
33        M: GLWEAutomorphismKeyAutomorphism<BE>,
34    {
35        module.glwe_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> + GetGaloisElement + GGLWEInfos,
41        Scratch<BE>: ScratchTakeCore<BE>,
42        M: GLWEAutomorphismKeyAutomorphism<BE>,
43    {
44        module.glwe_automorphism_key_automorphism_inplace(self, key, scratch);
45    }
46}
47
48impl<BE: Backend> GLWEAutomorphismKeyAutomorphism<BE> for Module<BE> where
49    Self: GaloisElement + GLWEAutomorphism<BE> + VecZnxAutomorphism
50{
51}
52
53pub trait GLWEAutomorphismKeyAutomorphism<BE: Backend>
54where
55    Self: GaloisElement + GLWEAutomorphism<BE> + VecZnxAutomorphism,
56{
57    fn glwe_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 glwe_automorphism_key_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
67    where
68        R: GGLWEToMut + SetGaloisElement + GGLWEInfos,
69        A: GGLWEToRef + GetGaloisElement + GGLWEInfos,
70        K: GGLWEPreparedToRef<BE> + GetGaloisElement + 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());
120    }
121
122    fn glwe_automorphism_key_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
123    where
124        R: GGLWEToMut + SetGaloisElement + GetGaloisElement + GGLWEInfos,
125        K: GGLWEPreparedToRef<BE> + GetGaloisElement + 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());
164    }
165}
poulpy_core/automorphism/gglwe_atk.rs

poulpy_core/automorphism/
gglwe_atk.rs
