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;
7use std::rc::Rc;
8
9impl Ring<u64> {
10 pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(
11 &self,
12 polys: &mut Vec<Option<Poly<u64>>>,
13 log_gap: usize,
14 ) {
15 let log_n: usize = self.log_n();
16 let log_nth_root: usize = log_n + 1;
17 let nth_root: usize = 1 << log_nth_root;
18 let log_start: usize = log_n - log_gap;
19 let mut log_end: usize = log_n;
20
21 let mut indices: Vec<usize> = Vec::<usize>::new();
22
23 // Retrives non-empty indexes
24 polys.iter().enumerate().for_each(|(i, poly)| {
25 if !poly.is_none() {
26 indices.push(i);
27 }
28 });
29
30 let gap: usize = max_gap(&indices);
31
32 if !ZEROGARBAGE {
33 if gap > 0 {
34 log_end -= gap.trailing_zeros() as usize;
35 }
36 }
37
38 assert!(
39 log_start < log_end,
40 "invalid input polys: gap between non None value is smaller than 2^log_gap"
41 );
42
43 let n_inv: Barrett<u64> = self
44 .modulus
45 .barrett
46 .prepare(self.modulus.inv(1 << (log_end - log_start)));
47
48 indices.iter().for_each(|i| {
49 if let Some(poly) = polys[*i].as_mut() {
50 if !NTT {
51 self.ntt_inplace::<true>(poly);
52 }
53 self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, poly);
54 }
55 });
56
57 let x_pow2: Vec<Poly<Montgomery<u64>>> = self.gen_x_pow_2::<true, false>(log_n);
58 let mut tmpa: Poly<u64> = self.new_poly();
59 let mut tmpb: Poly<u64> = self.new_poly();
60
61 for i in log_start..log_end {
62 let t: usize = 1 << (log_n - 1 - i);
63
64 let (polys_lo, polys_hi) = polys.split_at_mut(t);
65
66 for j in 0..t {
67 if let Some(poly_hi) = polys_hi[j].as_mut() {
68 self.a_mul_b_montgomery_into_a::<ONCE>(&x_pow2[log_n - i - 1], poly_hi);
69
70 if let Some(poly_lo) = polys_lo[j].as_mut() {
71 self.a_sub_b_into_c::<1, ONCE>(poly_lo, poly_hi, &mut tmpa);
72 self.a_add_b_into_b::<ONCE>(poly_hi, poly_lo);
73 }
74 }
75
76 if let Some(poly_lo) = polys_lo[j].as_mut() {
77 let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root);
78
79 if !polys_hi[j].is_none() {
80 self.a_apply_automorphism_add_b_into_b::<ONCE, true>(
81 &tmpa,
82 gal_el,
83 2 << self.log_n(),
84 poly_lo,
85 );
86 } else {
87 self.a_apply_automorphism_into_b::<true>(
88 poly_lo, gal_el, nth_root, &mut tmpa,
89 );
90 self.a_add_b_into_b::<ONCE>(&tmpa, poly_lo);
91 }
92 } else if let Some(poly_hi) = polys_hi[j].as_mut() {
93 let gal_el: usize = self.galois_element((1 << i) >> 1, i == 0, log_nth_root);
94 self.a_apply_automorphism_into_b::<true>(poly_hi, gal_el, nth_root, &mut tmpa);
95 self.a_sub_b_into_a::<1, ONCE>(&tmpa, poly_hi);
96 std::mem::swap(&mut polys_lo[j], &mut polys_hi[j]);
97 }
98 }
99
100 polys.truncate(t);
101 }
102
103 if !NTT {
104 if let Some(poly) = polys[0].as_mut() {
105 self.intt_inplace::<false>(poly);
106 }
107 }
108 }
109}
110
111// Returns the largest gap between two values in an ordered array of distinct values.
112// Panics if the array is not ordered or values are not distincts.
113fn max_gap(vec: &[usize]) -> usize {
114 let mut gap: usize = usize::MAX;
115 for i in 1..vec.len() {
116 let (l, r) = (vec[i - 1], vec[i]);
117 assert!(
118 r > l,
119 "invalid input vec: not sorted or collision between indices"
120 );
121 gap = min(gap, r - l);
122 if gap == 1 {
123 break;
124 }
125 }
126 gap
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}