poulpy/poulpy-core/src/glwe_trace.rs
poulpy-core/src/glwe_trace.rsRaw
1use std::collections::HashMap;
2
3use poulpy_hal::{
4 api::ModuleLogN,
5 layouts::{Backend, DataMut, GaloisElement, Module, Scratch, VecZnx},
6};
7
8use crate::{
9 GLWEAutomorphism, GLWECopy, GLWEShift, ScratchTakeCore,
10 layouts::{
11 Base2K, GGLWEInfos, GLWE, GLWEInfos, GLWELayout, GLWEToMut, GLWEToRef, LWEInfos,
12 prepared::{AutomorphismKeyPreparedToRef, GetAutomorphismGaloisElement},
13 },
14};
15
16impl GLWE<Vec<u8>> {
17 pub fn trace_galois_elements<M, BE: Backend>(module: &M) -> Vec<i64>
18 where
19 M: GLWETrace<BE>,
20 {
21 module.glwe_trace_galois_elements()
22 }
23
24 pub fn trace_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
25 where
26 R: GLWEInfos,
27 A: GLWEInfos,
28 K: GGLWEInfos,
29 M: GLWETrace<BE>,
30 {
31 module.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos)
32 }
33}
34
35impl<D: DataMut> GLWE<D> {
36 pub fn trace<A, K, M, BE: Backend>(
37 &mut self,
38 module: &M,
39 start: usize,
40 end: usize,
41 a: &A,
42 keys: &HashMap<i64, K>,
43 scratch: &mut Scratch<BE>,
44 ) where
45 A: GLWEToRef,
46 K: AutomorphismKeyPreparedToRef<BE> + GGLWEInfos + GetAutomorphismGaloisElement,
47 Scratch<BE>: ScratchTakeCore<BE>,
48 M: GLWETrace<BE>,
49 {
50 module.glwe_trace(self, start, end, a, keys, scratch);
51 }
52
53 pub fn trace_inplace<K, M, BE: Backend>(
54 &mut self,
55 module: &M,
56 start: usize,
57 end: usize,
58 keys: &HashMap<i64, K>,
59 scratch: &mut Scratch<BE>,
60 ) where
61 K: AutomorphismKeyPreparedToRef<BE> + GGLWEInfos + GetAutomorphismGaloisElement,
62 Scratch<BE>: ScratchTakeCore<BE>,
63 M: GLWETrace<BE>,
64 {
65 module.glwe_trace_inplace(self, start, end, keys, scratch);
66 }
67}
68
69impl<BE: Backend> GLWETrace<BE> for Module<BE> where
70 Self: ModuleLogN + GaloisElement + GLWEAutomorphism<BE> + GLWEShift<BE> + GLWECopy
71{
72}
73
74pub trait GLWETrace<BE: Backend>
75where
76 Self: ModuleLogN + GaloisElement + GLWEAutomorphism<BE> + GLWEShift<BE> + GLWECopy,
77{
78 fn glwe_trace_galois_elements(&self) -> Vec<i64> {
79 (0..self.log_n())
80 .map(|i| {
81 if i == 0 {
82 -1
83 } else {
84 self.galois_element(1 << (i - 1))
85 }
86 })
87 .collect()
88 }
89
90 fn glwe_trace_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
91 where
92 R: GLWEInfos,
93 A: GLWEInfos,
94 K: GGLWEInfos,
95 {
96 let trace: usize = self.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos);
97 if a_infos.base2k() != key_infos.base2k() {
98 let glwe_conv: usize = VecZnx::bytes_of(
99 self.n(),
100 (key_infos.rank_out() + 1).into(),
101 res_infos.k().min(a_infos.k()).div_ceil(key_infos.base2k()) as usize,
102 ) + self.vec_znx_normalize_tmp_bytes();
103 return glwe_conv + trace;
104 }
105
106 trace
107 }
108
109 fn glwe_trace<R, A, K>(&self, res: &mut R, start: usize, end: usize, a: &A, keys: &HashMap<i64, K>, scratch: &mut Scratch<BE>)
110 where
111 R: GLWEToMut,
112 A: GLWEToRef,
113 K: AutomorphismKeyPreparedToRef<BE> + GGLWEInfos + GetAutomorphismGaloisElement,
114 Scratch<BE>: ScratchTakeCore<BE>,
115 {
116 self.glwe_copy(res, a);
117 self.glwe_trace_inplace(res, start, end, keys, scratch);
118 }
119
120 fn glwe_trace_inplace<R, K>(&self, res: &mut R, start: usize, end: usize, keys: &HashMap<i64, K>, scratch: &mut Scratch<BE>)
121 where
122 R: GLWEToMut,
123 K: AutomorphismKeyPreparedToRef<BE> + GGLWEInfos + GetAutomorphismGaloisElement,
124 Scratch<BE>: ScratchTakeCore<BE>,
125 {
126 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
127
128 let basek_ksk: Base2K = keys.get(keys.keys().next().unwrap()).unwrap().base2k();
129
130 #[cfg(debug_assertions)]
131 {
132 assert_eq!(res.n(), self.n() as u32);
133 assert!(start < end);
134 assert!(end <= self.log_n());
135 for key in keys.values() {
136 assert_eq!(key.n(), self.n() as u32);
137 assert_eq!(key.base2k(), basek_ksk);
138 assert_eq!(key.rank_in(), res.rank());
139 assert_eq!(key.rank_out(), res.rank());
140 }
141 }
142
143 if res.base2k() != basek_ksk {
144 let (mut self_conv, scratch_1) = scratch.take_glwe_ct(
145 self,
146 &GLWELayout {
147 n: self.n().into(),
148 base2k: basek_ksk,
149 k: res.k(),
150 rank: res.rank(),
151 },
152 );
153
154 for j in 0..(res.rank() + 1).into() {
155 self.vec_znx_normalize(
156 basek_ksk.into(),
157 &mut self_conv.data,
158 j,
159 basek_ksk.into(),
160 res.data(),
161 j,
162 scratch_1,
163 );
164 }
165
166 for i in start..end {
167 self.glwe_rsh(1, &mut self_conv, scratch_1);
168
169 let p: i64 = if i == 0 {
170 -1
171 } else {
172 self.galois_element(1 << (i - 1))
173 };
174
175 if let Some(key) = keys.get(&p) {
176 self.glwe_automorphism_add_inplace(&mut self_conv, key, scratch_1);
177 } else {
178 panic!("keys[{p}] is empty")
179 }
180 }
181
182 for j in 0..(res.rank() + 1).into() {
183 self.vec_znx_normalize(
184 res.base2k().into(),
185 res.data_mut(),
186 j,
187 basek_ksk.into(),
188 &self_conv.data,
189 j,
190 scratch_1,
191 );
192 }
193 } else {
194 for i in start..end {
195 self.glwe_rsh(1, res, scratch);
196
197 let p: i64 = if i == 0 {
198 -1
199 } else {
200 self.galois_element(1 << (i - 1))
201 };
202
203 if let Some(key) = keys.get(&p) {
204 self.glwe_automorphism_add_inplace(res, key, scratch);
205 } else {
206 panic!("keys[{p}] is empty")
207 }
208 }
209 }
210 }
211}