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