poulpy/math/src/ring/impl_u64/ring_rns.rs
math/src/ring/impl_u64/ring_rns.rsRaw
1use crate::modulus::barrett::Barrett;
2use crate::modulus::montgomery::Montgomery;
3use crate::modulus::REDUCEMOD;
4use crate::poly::PolyRNS;
5use crate::ring::{Ring, RingRNS};
6use crate::scalar::ScalarRNS;
7use num_bigint::BigInt;
8use std::sync::Arc;
9
10
11
12impl RingRNS<u64> {
13 pub fn new(n: usize, moduli: Vec<u64>) -> Self {
14 assert!(!moduli.is_empty(), "moduli cannot be empty");
15 let rings: Vec<Arc<Ring<u64>>> = moduli
16 .into_iter()
17 .map(|prime| Arc::new(Ring::new(n, prime, 1)))
18 .collect();
19 return RingRNS(rings);
20 }
21
22 pub fn modulus(&self) -> BigInt {
23 let mut modulus = BigInt::from(1);
24 self.0
25 .iter()
26 .enumerate()
27 .for_each(|(_, r)| modulus *= BigInt::from(r.modulus.q));
28 modulus
29 }
30
31 pub fn rescaling_constant(&self) -> ScalarRNS<Barrett<u64>> {
32 let level = self.level();
33 let q_scale: u64 = self.0[level].modulus.q;
34 ScalarRNS(
35 (0..level)
36 .map(|i| {
37 self.0[i]
38 .modulus
39 .barrett
40 .prepare(self.0[i].modulus.q - self.0[i].modulus.inv(q_scale))
41 })
42 .collect(),
43 )
44 }
45
46 pub fn from_bigint_inplace(&self, coeffs: &[BigInt], step: usize, a: &mut PolyRNS<u64>) {
47 let level = self.level();
48 assert!(
49 level <= a.level(),
50 "invalid level: level={} > a.level()={}",
51 level,
52 a.level()
53 );
54 (0..level).for_each(|i| self.0[i].from_bigint(coeffs, step, a.at_mut(i)));
55 }
56
57 pub fn to_bigint_inplace(&self, a: &PolyRNS<u64>, step: usize, coeffs: &mut [BigInt]) {
58 assert!(
59 step <= a.n(),
60 "invalid step: step={} > a.n()={}",
61 step,
62 a.n()
63 );
64 assert!(
65 coeffs.len() <= a.n() / step,
66 "invalid coeffs: coeffs.len()={} > a.n()/step={}",
67 coeffs.len(),
68 a.n() / step
69 );
70
71 let mut inv_crt: Vec<BigInt> = vec![BigInt::default(); self.level() + 1];
72 let q_big: BigInt = self.modulus();
73 let q_big_half: BigInt = &q_big >> 1;
74
75 inv_crt.iter_mut().enumerate().for_each(|(i, a)| {
76 let qi_big = BigInt::from(self.0[i].modulus.q);
77 *a = &q_big / &qi_big;
78 *a *= a.modinv(&qi_big).unwrap();
79 });
80
81 (0..self.n()).step_by(step).enumerate().for_each(|(i, j)| {
82 coeffs[j] = BigInt::from(a.at(0).0[i]) * &inv_crt[0];
83 (1..self.level() + 1).for_each(|k| {
84 coeffs[j] += BigInt::from(a.at(k).0[i] * &inv_crt[k]);
85 });
86 coeffs[j] %= &q_big;
87 if &coeffs[j] >= &q_big_half {
88 coeffs[j] -= &q_big;
89 }
90 });
91 }
92}
93
94impl RingRNS<u64> {
95 pub fn ntt_inplace<const LAZY: bool>(&self, a: &mut PolyRNS<u64>) {
96 self.0
97 .iter()
98 .enumerate()
99 .for_each(|(i, ring)| ring.ntt_inplace::<LAZY>(&mut a.0[i]));
100 }
101
102 pub fn intt_inplace<const LAZY: bool>(&self, a: &mut PolyRNS<u64>) {
103 self.0
104 .iter()
105 .enumerate()
106 .for_each(|(i, ring)| ring.intt_inplace::<LAZY>(&mut a.0[i]));
107 }
108
109 pub fn ntt<const LAZY: bool>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>) {
110 self.0
111 .iter()
112 .enumerate()
113 .for_each(|(i, ring)| ring.ntt::<LAZY>(&a.0[i], &mut b.0[i]));
114 }
115
116 pub fn intt<const LAZY: bool>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>) {
117 self.0
118 .iter()
119 .enumerate()
120 .for_each(|(i, ring)| ring.intt::<LAZY>(&a.0[i], &mut b.0[i]));
121 }
122}
123
124impl RingRNS<u64> {
125 #[inline(always)]
126 pub fn a_add_b_into_c<const REDUCE: REDUCEMOD>(
127 &self,
128 a: &PolyRNS<u64>,
129 b: &PolyRNS<u64>,
130 c: &mut PolyRNS<u64>,
131 ) {
132 debug_assert!(
133 a.level() >= self.level(),
134 "a.level()={} < self.level()={}",
135 a.level(),
136 self.level()
137 );
138 debug_assert!(
139 b.level() >= self.level(),
140 "b.level()={} < self.level()={}",
141 b.level(),
142 self.level()
143 );
144 debug_assert!(
145 c.level() >= self.level(),
146 "c.level()={} < self.level()={}",
147 c.level(),
148 self.level()
149 );
150 self.0
151 .iter()
152 .enumerate()
153 .for_each(|(i, ring)| ring.a_add_b_into_c::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
154 }
155
156 #[inline(always)]
157 pub fn a_add_b_into_b<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>) {
158 debug_assert!(
159 a.level() >= self.level(),
160 "a.level()={} < self.level()={}",
161 a.level(),
162 self.level()
163 );
164 debug_assert!(
165 b.level() >= self.level(),
166 "b.level()={} < self.level()={}",
167 b.level(),
168 self.level()
169 );
170 self.0
171 .iter()
172 .enumerate()
173 .for_each(|(i, ring)| ring.a_add_b_into_b::<REDUCE>(&a.0[i], &mut b.0[i]));
174 }
175
176 #[inline(always)]
177 pub fn a_sub_b_into_c<const BRANGE: u8, const REDUCE: REDUCEMOD>(
178 &self,
179 a: &PolyRNS<u64>,
180 b: &PolyRNS<u64>,
181 c: &mut PolyRNS<u64>,
182 ) {
183 debug_assert!(
184 a.level() >= self.level(),
185 "a.level()={} < self.level()={}",
186 a.level(),
187 self.level()
188 );
189 debug_assert!(
190 b.level() >= self.level(),
191 "b.level()={} < self.level()={}",
192 b.level(),
193 self.level()
194 );
195 debug_assert!(
196 c.level() >= self.level(),
197 "c.level()={} < self.level()={}",
198 c.level(),
199 self.level()
200 );
201 self.0
202 .iter()
203 .enumerate()
204 .for_each(|(i, ring)| ring.a_sub_b_into_c::<BRANGE, REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
205 }
206
207 #[inline(always)]
208 pub fn a_sub_b_into_b<const BRANGE: u8, const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>) {
209 debug_assert!(
210 a.level() >= self.level(),
211 "a.level()={} < self.level()={}",
212 a.level(),
213 self.level()
214 );
215 debug_assert!(
216 b.level() >= self.level(),
217 "b.level()={} < self.level()={}",
218 b.level(),
219 self.level()
220 );
221 self.0
222 .iter()
223 .enumerate()
224 .for_each(|(i, ring)| ring.a_sub_b_into_b::<BRANGE, REDUCE>(&a.0[i], &mut b.0[i]));
225 }
226
227 #[inline(always)]
228 pub fn a_sub_b_into_a<const BRANGE: u8, const REDUCE: REDUCEMOD>(&self, b: &PolyRNS<u64>, a: &mut PolyRNS<u64>) {
229 debug_assert!(
230 a.level() >= self.level(),
231 "a.level()={} < self.level()={}",
232 a.level(),
233 self.level()
234 );
235 debug_assert!(
236 b.level() >= self.level(),
237 "b.level()={} < self.level()={}",
238 b.level(),
239 self.level()
240 );
241 self.0
242 .iter()
243 .enumerate()
244 .for_each(|(i, ring)| ring.a_sub_b_into_a::<BRANGE, REDUCE>(&b.0[i], &mut a.0[i]));
245 }
246
247 #[inline(always)]
248 pub fn a_neg_into_b<const ARANGE: u8, const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>) {
249 debug_assert!(
250 a.level() >= self.level(),
251 "a.level()={} < self.level()={}",
252 a.level(),
253 self.level()
254 );
255 debug_assert!(
256 b.level() >= self.level(),
257 "b.level()={} < self.level()={}",
258 b.level(),
259 self.level()
260 );
261 self.0
262 .iter()
263 .enumerate()
264 .for_each(|(i, ring)| ring.a_neg_into_b::<ARANGE, REDUCE>(&a.0[i], &mut b.0[i]));
265 }
266
267 #[inline(always)]
268 pub fn a_neg_into_a<const ARANGE: u8, const REDUCE: REDUCEMOD>(&self, a: &mut PolyRNS<u64>) {
269 debug_assert!(
270 a.level() >= self.level(),
271 "a.level()={} < self.level()={}",
272 a.level(),
273 self.level()
274 );
275 self.0
276 .iter()
277 .enumerate()
278 .for_each(|(i, ring)| ring.a_neg_into_a::<ARANGE, REDUCE>(&mut a.0[i]));
279 }
280
281 #[inline(always)]
282 pub fn mul_montgomery_external<const REDUCE: REDUCEMOD>(
283 &self,
284 a: &PolyRNS<Montgomery<u64>>,
285 b: &PolyRNS<u64>,
286 c: &mut PolyRNS<u64>,
287 ) {
288 debug_assert!(
289 a.level() >= self.level(),
290 "a.level()={} < self.level()={}",
291 a.level(),
292 self.level()
293 );
294 debug_assert!(
295 b.level() >= self.level(),
296 "b.level()={} < self.level()={}",
297 b.level(),
298 self.level()
299 );
300 debug_assert!(
301 c.level() >= self.level(),
302 "c.level()={} < self.level()={}",
303 c.level(),
304 self.level()
305 );
306 self.0.iter().enumerate().for_each(|(i, ring)| {
307 ring.a_mul_b_montgomery_into_c::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i])
308 });
309 }
310
311 #[inline(always)]
312 pub fn mul_montgomery_external_inplace<const REDUCE: REDUCEMOD>(
313 &self,
314 a: &PolyRNS<Montgomery<u64>>,
315 b: &mut PolyRNS<u64>,
316 ) {
317 debug_assert!(
318 a.level() >= self.level(),
319 "a.level()={} < self.level()={}",
320 a.level(),
321 self.level()
322 );
323 debug_assert!(
324 b.level() >= self.level(),
325 "b.level()={} < self.level()={}",
326 b.level(),
327 self.level()
328 );
329 self.0.iter().enumerate().for_each(|(i, ring)| {
330 ring.a_mul_b_montgomery_into_a::<REDUCE>(&a.0[i], &mut b.0[i])
331 });
332 }
333
334 #[inline(always)]
335 pub fn mul_scalar<const REDUCE: REDUCEMOD>(
336 &self,
337 a: &PolyRNS<u64>,
338 b: &u64,
339 c: &mut PolyRNS<u64>,
340 ) {
341 debug_assert!(
342 a.level() >= self.level(),
343 "a.level()={} < self.level()={}",
344 a.level(),
345 self.level()
346 );
347 debug_assert!(
348 c.level() >= self.level(),
349 "b.level()={} < self.level()={}",
350 c.level(),
351 self.level()
352 );
353 self.0
354 .iter()
355 .enumerate()
356 .for_each(|(i, ring)| ring.a_mul_b_scalar_into_c::<REDUCE>(&a.0[i], b, &mut c.0[i]));
357 }
358
359 #[inline(always)]
360 pub fn mul_scalar_inplace<const REDUCE: REDUCEMOD>(&self, b: &u64, a: &mut PolyRNS<u64>) {
361 debug_assert!(
362 a.level() >= self.level(),
363 "b.level()={} < self.level()={}",
364 a.level(),
365 self.level()
366 );
367 self.0
368 .iter()
369 .enumerate()
370 .for_each(|(i, ring)| ring.a_mul_b_scalar_into_a::<REDUCE>(b, &mut a.0[i]));
371 }
372
373 #[inline(always)]
374 pub fn a_sub_b_add_scalar_mul_scalar_barrett_into_e<const BRANGE:u8, const REDUCE:REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &PolyRNS<u64>, c: &u64, d: &Barrett<u64>, e: &mut PolyRNS<u64>){
375 debug_assert!(
376 a.level() >= self.level(),
377 "a.level()={} < self.level()={}",
378 a.level(),
379 self.level()
380 );
381 debug_assert!(
382 b.level() >= self.level(),
383 "b.level()={} < self.level()={}",
384 b.level(),
385 self.level()
386 );
387 debug_assert!(
388 e.level() >= self.level(),
389 "e.level()={} < self.level()={}",
390 e.level(),
391 self.level()
392 );
393 self.0
394 .iter()
395 .enumerate()
396 .for_each(|(i, ring)| ring.a_sub_b_add_c_scalar_mul_d_scalar_barrett_into_e::<BRANGE, REDUCE>(&a.0[i], &b.0[i], c, d, &mut e.0[i]));
397 }
398
399 #[inline(always)]
400 pub fn b_sub_a_add_c_scalar_mul_d_scalar_barrett_into_a<const BRANGE:u8, const REDUCE:REDUCEMOD>(&self, b: &PolyRNS<u64>, c: &u64, d: &Barrett<u64>, a: &mut PolyRNS<u64>){
401 debug_assert!(
402 a.level() >= self.level(),
403 "a.level()={} < self.level()={}",
404 a.level(),
405 self.level()
406 );
407 debug_assert!(
408 b.level() >= self.level(),
409 "b.level()={} < self.level()={}",
410 b.level(),
411 self.level()
412 );
413 self.0
414 .iter()
415 .enumerate()
416 .for_each(|(i, ring)| ring.b_sub_a_add_c_scalar_mul_d_scalar_barrett_into_a::<BRANGE, REDUCE>(&b.0[i], c, d, &mut a.0[i]));
417 }
418}