poulpy/math/src/ring/impl_u64/packing.rs
math/src/ring/impl_u64/packing.rsRaw
1use crate::modulus::barrett::Barrett;
2use crate::modulus::montgomery::Montgomery;
3use crate::modulus::ONCE;
4use crate::poly::Poly;
5use crate::ring::Ring;
6use std::cmp::min;
7
8impl Ring<u64> {
9 pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(
10 &self,
11 polys: &mut Vec<Option<&mut Poly<u64>>>,
12 log_gap: usize,
13 ) {
14 let log_n: usize = self.log_n();
15 let log_nth_root: usize = log_n + 1;
16 let nth_root: usize = 1 << log_nth_root;
17 let log_start: usize = log_n - log_gap;
18 let mut log_end: usize = log_n;
19
20 let mut indices: Vec<usize> = Vec::<usize>::new();
21
22 // Retrives non-empty indexes
23 polys.iter().enumerate().for_each(|(i, poly)| {
24 if !poly.is_none() {
25 indices.push(i);
26 }
27 });
28
29 let gap: usize = max_gap(&indices);
30
31 if !ZEROGARBAGE {
32 if gap > 0 {
33 log_end -= gap.trailing_zeros() as usize;
34 }
35 }
36
37 assert!(
38 log_start < log_end,
39 "invalid input polys: gap between non None value is smaller than 2^log_gap"
40 );
41
42 let n_inv: Barrett<u64> = self
43 .modulus
44 .barrett
45 .prepare(self.modulus.inv(1 << (log_end - log_start)));
46
47 indices.iter().for_each(|i| {
48 if let Some(poly) = polys[*i].as_deref_mut() {
49 if !NTT {
50 self.ntt_inplace::<true>(poly);
51 }
52 self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, poly);
53 }
54 });
55
56 let x_pow2: Vec<Poly<Montgomery<u64>>> = self.gen_x_pow_2::<true, false>(log_n);
57 let mut tmpa: Poly<u64> = self.new_poly();
58 let mut tmpb: Poly<u64> = self.new_poly();
59
60 for i in log_start..log_end {
61 let t: usize = 1 << (log_n - 1 - i);
62
63 let (polys_lo, polys_hi) = polys.split_at_mut(t);
64
65 for j in 0..t {
66 if let Some(poly_hi) = polys_hi[j].as_mut() {
67 self.a_mul_b_montgomery_into_a::<ONCE>(&x_pow2[log_n - i - 1], poly_hi);
68
69 if let Some(poly_lo) = polys_lo[j].as_mut() {
70 self.a_sub_b_into_c::<1, ONCE>(poly_lo, poly_hi, &mut tmpa);
71 self.a_add_b_into_b::<ONCE>(poly_hi, poly_lo);
72 }
73 }
74
75 if let Some(poly_lo) = polys_lo[j].as_mut() {
76 let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root);
77
78 if !polys_hi[j].is_none() {
79 self.a_apply_automorphism_add_b_into_b::<ONCE, true>(
80 &tmpa,
81 gal_el,
82 2 << self.log_n(),
83 poly_lo,
84 );
85 } else {
86 self.a_apply_automorphism_into_b::<true>(
87 poly_lo, gal_el, nth_root, &mut tmpa,
88 );
89 self.a_add_b_into_b::<ONCE>(&tmpa, poly_lo);
90 }
91 } else if let Some(poly_hi) = polys_hi[j].as_mut() {
92 let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root);
93 self.a_apply_automorphism_into_b::<true>(poly_hi, gal_el, nth_root, &mut tmpa);
94 self.a_sub_b_into_a::<1, ONCE>(&tmpa, poly_hi);
95 std::mem::swap(&mut polys_lo[j], &mut polys_hi[j]);
96 }
97 }
98
99 polys.truncate(t);
100 }
101
102 if !NTT {
103 if let Some(poly) = polys[0].as_mut() {
104 self.intt_inplace::<false>(poly);
105 }
106 }
107 }
108}
109
110// Returns the largest gap between two values in an ordered array of distinct values.
111// Panics if the array is not ordered or values are not distincts.
112fn max_gap(vec: &[usize]) -> usize {
113 let mut gap: usize = usize::MAX;
114 for i in 1..vec.len() {
115 let (l, r) = (vec[i - 1], vec[i]);
116 assert!(
117 r > l,
118 "invalid input vec: not sorted or collision between indices"
119 );
120 gap = min(gap, r - l);
121 if gap == 1 {
122 break;
123 }
124 }
125 gap
126}
127
128
129pub struct StreamRepacker {
130 accumulators: Vec<Accumulator>,
131 tmp_a: Poly<u64>,
132 tmp_b: Poly<u64>,
133 x_pow_2: Vec<Poly<Montgomery<u64>>>,
134 n_inv: Barrett<u64>,
135 pub results: Vec<Poly<u64>>,
136 counter: usize,
137}
138
139pub struct Accumulator {
140 buf: Poly<u64>,
141 value: bool,
142 control: bool,
143}
144
145impl Accumulator {
146 pub fn new(r: &Ring<u64>) -> Self {
147 Self {
148 buf: r.new_poly(),
149 value: false,
150 control: false,
151 }
152 }
153}
154
155impl StreamRepacker {
156 pub fn new(r: &Ring<u64>) -> Self {
157 let mut accumulators: Vec<Accumulator> = Vec::<Accumulator>::new();
158
159 (0..r.log_n()).for_each(|_| accumulators.push(Accumulator::new(r)));
160
161 Self {
162 accumulators: accumulators,
163 tmp_a: r.new_poly(),
164 tmp_b: r.new_poly(),
165 x_pow_2: r.gen_x_pow_2::<true, false>(r.log_n()),
166 n_inv: r.modulus.barrett.prepare(r.modulus.inv(r.n() as u64)),
167 results: Vec::<Poly<u64>>::new(),
168 counter: 0,
169 }
170 }
171
172 pub fn reset(&mut self) {
173 for i in 0..self.accumulators.len() {
174 self.accumulators[i].value = false;
175 self.accumulators[i].control = false;
176 }
177 self.counter = 0;
178 }
179
180 pub fn add<const NTT: bool>(&mut self, r: &Ring<u64>, a: Option<&Poly<u64>>) {
181 assert!(NTT, "invalid parameterization: const NTT must be true");
182 pack_core::<NTT>(
183 r,
184 a,
185 &mut self.accumulators,
186 &self.n_inv,
187 &self.x_pow_2,
188 &mut self.tmp_a,
189 &mut self.tmp_b,
190 0,
191 );
192 self.counter += 1;
193 if self.counter == r.n() {
194 self.results
195 .push(self.accumulators[r.log_n() - 1].buf.clone());
196 self.reset();
197 }
198 }
199
200 pub fn flush<const NTT: bool>(&mut self, r: &Ring<u64>) {
201 assert!(NTT, "invalid parameterization: const NTT must be true");
202 if self.counter != 0 {
203 while self.counter != r.n() - 1 {
204 self.add::<NTT>(r, None);
205 }
206 }
207 }
208}
209
210fn pack_core<const NTT: bool>(
211 r: &Ring<u64>,
212 a: Option<&Poly<u64>>,
213 accumulators: &mut [Accumulator],
214 n_inv: &Barrett<u64>,
215 x_pow_2: &[Poly<u64>],
216 tmp_a: &mut Poly<u64>,
217 tmp_b: &mut Poly<u64>,
218 i: usize,
219) {
220 if i == r.log_n() {
221 return;
222 }
223
224 let (acc_prev, acc_next) = accumulators.split_at_mut(1);
225
226 if !acc_prev[0].control {
227 let acc_mut_ref: &mut Accumulator = &mut acc_prev[0]; // from split_at_mut
228
229 if let Some(a_ref) = a {
230 acc_mut_ref.buf.copy_from(a_ref);
231 acc_mut_ref.value = true
232 } else {
233 acc_mut_ref.value = false
234 }
235 acc_mut_ref.control = true;
236 } else {
237 combine::<true>(r, &mut acc_prev[0], a, n_inv, x_pow_2, tmp_a, tmp_b, i);
238 acc_prev[0].control = false;
239
240 if acc_prev[0].value {
241 pack_core::<NTT>(
242 r,
243 Some(&acc_prev[0].buf),
244 acc_next,
245 n_inv,
246 x_pow_2,
247 tmp_a,
248 tmp_b,
249 i + 1,
250 );
251 } else {
252 pack_core::<NTT>(r, None, acc_next, n_inv, x_pow_2, tmp_a, tmp_b, i + 1);
253 }
254 }
255}
256
257fn combine<const NTT: bool>(
258 r: &Ring<u64>,
259 acc: &mut Accumulator,
260 b: Option<&Poly<u64>>,
261 n_inv: &Barrett<u64>,
262 x_pow_2: &[Poly<u64>],
263 tmp_a: &mut Poly<u64>,
264 tmp_b: &mut Poly<u64>,
265 i: usize,
266) {
267 let log_n = r.log_n();
268 let log_nth_root = log_n + 1;
269 let nth_root = 1 << log_nth_root;
270 let gal_el: usize = r.galois_element((1 << i) >> 1, i == 0, log_nth_root);
271
272 let a: &mut Poly<u64> = &mut acc.buf;
273
274 if acc.value {
275 if i == 0 {
276 r.a_mul_b_scalar_barrett_into_a::<ONCE>(n_inv, a);
277 }
278
279 if let Some(b) = b {
280 // tmp_a = b * X^t
281 r.a_mul_b_montgomery_into_c::<ONCE>(b, &x_pow_2[log_n - i - 1], tmp_a);
282
283 if i == 0 {
284 r.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, tmp_a);
285 }
286
287 // tmp_b = a - b*X^t
288 r.a_sub_b_into_c::<1, ONCE>(a, tmp_a, tmp_b);
289
290 // a = a + b * X^t
291 r.a_add_b_into_b::<ONCE>(tmp_a, a);
292
293 // a = a + b * X^t + phi(a - b * X^t)
294 r.a_apply_automorphism_add_b_into_b::<ONCE, NTT>(tmp_b, gal_el, nth_root, a);
295 } else {
296 // tmp_a = phi(a)
297 r.a_apply_automorphism_into_b::<NTT>(a, gal_el, nth_root, tmp_a);
298 // a = a + phi(a)
299 r.a_add_b_into_b::<ONCE>(tmp_a, a);
300 }
301 } else {
302 if let Some(b) = b {
303 // tmp_b = b * X^t
304 r.a_mul_b_montgomery_into_c::<ONCE>(b, &x_pow_2[log_n - i - 1], tmp_b);
305
306 if i == 0 {
307 r.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, tmp_b);
308 }
309
310 // tmp_a = phi(b * X^t)
311 r.a_apply_automorphism_into_b::<NTT>(tmp_b, gal_el, nth_root, tmp_a);
312
313 // a = (b* X^t - phi(b* X^t))
314 r.a_sub_b_into_c::<1, ONCE>(tmp_b, tmp_a, a);
315 acc.value = true
316 }
317 }
318}