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;
8
9pub fn new_rings(n: usize, moduli: Vec<u64>) -> Vec<Ring<u64>> {
10 assert!(!moduli.is_empty(), "moduli cannot be empty");
11 let rings: Vec<Ring<u64>> = moduli
12 .into_iter()
13 .map(|prime| Ring::new(n, prime, 1))
14 .collect();
15 return rings;
16}
17
18impl<'a> RingRNS<'a, u64> {
19 pub fn new(rings: &'a [Ring<u64>]) -> Self {
20 RingRNS(rings)
21 }
22
23 pub fn modulus(&self) -> BigInt {
24 let mut modulus = BigInt::from(1);
25 self.0
26 .iter()
27 .enumerate()
28 .for_each(|(_, r)| modulus *= BigInt::from(r.modulus.q));
29 modulus
30 }
31
32 pub fn rescaling_constant(&self) -> ScalarRNS<Barrett<u64>> {
33 let level = self.level();
34 let q_scale: u64 = self.0[level].modulus.q;
35 ScalarRNS(
36 (0..level)
37 .map(|i| {
38 self.0[i]
39 .modulus
40 .barrett
41 .prepare(self.0[i].modulus.q - self.0[i].modulus.inv(q_scale))
42 })
43 .collect(),
44 )
45 }
46
47 pub fn from_bigint_inplace(&self, coeffs: &[BigInt], step: usize, a: &mut PolyRNS<u64>) {
48 let level = self.level();
49 assert!(
50 level <= a.level(),
51 "invalid level: level={} > a.level()={}",
52 level,
53 a.level()
54 );
55 (0..level).for_each(|i| self.0[i].from_bigint(coeffs, step, a.at_mut(i)));
56 }
57
58 pub fn to_bigint_inplace(&self, a: &PolyRNS<u64>, step: usize, coeffs: &mut [BigInt]) {
59 assert!(
60 step <= a.n(),
61 "invalid step: step={} > a.n()={}",
62 step,
63 a.n()
64 );
65 assert!(
66 coeffs.len() <= a.n() / step,
67 "invalid coeffs: coeffs.len()={} > a.n()/step={}",
68 coeffs.len(),
69 a.n() / step
70 );
71
72 let mut inv_crt: Vec<BigInt> = vec![BigInt::default(); self.level() + 1];
73 let q_big: BigInt = self.modulus();
74 let q_big_half: BigInt = &q_big >> 1;
75
76 inv_crt.iter_mut().enumerate().for_each(|(i, a)| {
77 let qi_big = BigInt::from(self.0[i].modulus.q);
78 *a = &q_big / &qi_big;
79 *a *= a.modinv(&qi_big).unwrap();
80 });
81
82 (0..self.n()).step_by(step).enumerate().for_each(|(i, j)| {
83 coeffs[j] = BigInt::from(a.at(0).0[i]) * &inv_crt[0];
84 (1..self.level() + 1).for_each(|k| {
85 coeffs[j] += BigInt::from(a.at(k).0[i] * &inv_crt[k]);
86 });
87 coeffs[j] %= &q_big;
88 if &coeffs[j] >= &q_big_half {
89 coeffs[j] -= &q_big;
90 }
91 });
92 }
93}
94
95impl RingRNS<'_, u64> {
96 pub fn ntt_inplace<const LAZY: bool>(&self, a: &mut PolyRNS<u64>) {
97 self.0
98 .iter()
99 .enumerate()
100 .for_each(|(i, ring)| ring.ntt_inplace::<LAZY>(&mut a.0[i]));
101 }
102
103 pub fn intt_inplace<const LAZY: bool>(&self, a: &mut PolyRNS<u64>) {
104 self.0
105 .iter()
106 .enumerate()
107 .for_each(|(i, ring)| ring.intt_inplace::<LAZY>(&mut a.0[i]));
108 }
109
110 pub fn ntt<const LAZY: bool>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>) {
111 self.0
112 .iter()
113 .enumerate()
114 .for_each(|(i, ring)| ring.ntt::<LAZY>(&a.0[i], &mut b.0[i]));
115 }
116
117 pub fn intt<const LAZY: bool>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>) {
118 self.0
119 .iter()
120 .enumerate()
121 .for_each(|(i, ring)| ring.intt::<LAZY>(&a.0[i], &mut b.0[i]));
122 }
123}
124
125impl RingRNS<'_, u64> {
126 #[inline(always)]
127 pub fn add<const REDUCE: REDUCEMOD>(
128 &self,
129 a: &PolyRNS<u64>,
130 b: &PolyRNS<u64>,
131 c: &mut PolyRNS<u64>,
132 ) {
133 debug_assert!(
134 a.level() >= self.level(),
135 "a.level()={} < self.level()={}",
136 a.level(),
137 self.level()
138 );
139 debug_assert!(
140 b.level() >= self.level(),
141 "b.level()={} < self.level()={}",
142 b.level(),
143 self.level()
144 );
145 debug_assert!(
146 c.level() >= self.level(),
147 "c.level()={} < self.level()={}",
148 c.level(),
149 self.level()
150 );
151 self.0
152 .iter()
153 .enumerate()
154 .for_each(|(i, ring)| ring.add::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
155 }
156
157 #[inline(always)]
158 pub fn add_inplace<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>) {
159 debug_assert!(
160 a.level() >= self.level(),
161 "a.level()={} < self.level()={}",
162 a.level(),
163 self.level()
164 );
165 debug_assert!(
166 b.level() >= self.level(),
167 "b.level()={} < self.level()={}",
168 b.level(),
169 self.level()
170 );
171 self.0
172 .iter()
173 .enumerate()
174 .for_each(|(i, ring)| ring.add_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
175 }
176
177 #[inline(always)]
178 pub fn sub<const REDUCE: REDUCEMOD>(
179 &self,
180 a: &PolyRNS<u64>,
181 b: &PolyRNS<u64>,
182 c: &mut PolyRNS<u64>,
183 ) {
184 debug_assert!(
185 a.level() >= self.level(),
186 "a.level()={} < self.level()={}",
187 a.level(),
188 self.level()
189 );
190 debug_assert!(
191 b.level() >= self.level(),
192 "b.level()={} < self.level()={}",
193 b.level(),
194 self.level()
195 );
196 debug_assert!(
197 c.level() >= self.level(),
198 "c.level()={} < self.level()={}",
199 c.level(),
200 self.level()
201 );
202 self.0
203 .iter()
204 .enumerate()
205 .for_each(|(i, ring)| ring.sub::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
206 }
207
208 #[inline(always)]
209 pub fn sub_inplace<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>) {
210 debug_assert!(
211 a.level() >= self.level(),
212 "a.level()={} < self.level()={}",
213 a.level(),
214 self.level()
215 );
216 debug_assert!(
217 b.level() >= self.level(),
218 "b.level()={} < self.level()={}",
219 b.level(),
220 self.level()
221 );
222 self.0
223 .iter()
224 .enumerate()
225 .for_each(|(i, ring)| ring.sub_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
226 }
227
228 #[inline(always)]
229 pub fn neg<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>) {
230 debug_assert!(
231 a.level() >= self.level(),
232 "a.level()={} < self.level()={}",
233 a.level(),
234 self.level()
235 );
236 debug_assert!(
237 b.level() >= self.level(),
238 "b.level()={} < self.level()={}",
239 b.level(),
240 self.level()
241 );
242 self.0
243 .iter()
244 .enumerate()
245 .for_each(|(i, ring)| ring.neg::<REDUCE>(&a.0[i], &mut b.0[i]));
246 }
247
248 #[inline(always)]
249 pub fn neg_inplace<const REDUCE: REDUCEMOD>(&self, a: &mut PolyRNS<u64>) {
250 debug_assert!(
251 a.level() >= self.level(),
252 "a.level()={} < self.level()={}",
253 a.level(),
254 self.level()
255 );
256 self.0
257 .iter()
258 .enumerate()
259 .for_each(|(i, ring)| ring.neg_inplace::<REDUCE>(&mut a.0[i]));
260 }
261
262 #[inline(always)]
263 pub fn mul_montgomery_external<const REDUCE: REDUCEMOD>(
264 &self,
265 a: &PolyRNS<Montgomery<u64>>,
266 b: &PolyRNS<u64>,
267 c: &mut PolyRNS<u64>,
268 ) {
269 debug_assert!(
270 a.level() >= self.level(),
271 "a.level()={} < self.level()={}",
272 a.level(),
273 self.level()
274 );
275 debug_assert!(
276 b.level() >= self.level(),
277 "b.level()={} < self.level()={}",
278 b.level(),
279 self.level()
280 );
281 debug_assert!(
282 c.level() >= self.level(),
283 "c.level()={} < self.level()={}",
284 c.level(),
285 self.level()
286 );
287 self.0.iter().enumerate().for_each(|(i, ring)| {
288 ring.mul_montgomery_external::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i])
289 });
290 }
291
292 #[inline(always)]
293 pub fn mul_montgomery_external_inplace<const REDUCE: REDUCEMOD>(
294 &self,
295 a: &PolyRNS<Montgomery<u64>>,
296 b: &mut PolyRNS<u64>,
297 ) {
298 debug_assert!(
299 a.level() >= self.level(),
300 "a.level()={} < self.level()={}",
301 a.level(),
302 self.level()
303 );
304 debug_assert!(
305 b.level() >= self.level(),
306 "b.level()={} < self.level()={}",
307 b.level(),
308 self.level()
309 );
310 self.0.iter().enumerate().for_each(|(i, ring)| {
311 ring.mul_montgomery_external_inplace::<REDUCE>(&a.0[i], &mut b.0[i])
312 });
313 }
314
315 #[inline(always)]
316 pub fn mul_scalar<const REDUCE: REDUCEMOD>(
317 &self,
318 a: &PolyRNS<u64>,
319 b: &u64,
320 c: &mut PolyRNS<u64>,
321 ) {
322 debug_assert!(
323 a.level() >= self.level(),
324 "a.level()={} < self.level()={}",
325 a.level(),
326 self.level()
327 );
328 debug_assert!(
329 c.level() >= self.level(),
330 "b.level()={} < self.level()={}",
331 c.level(),
332 self.level()
333 );
334 self.0
335 .iter()
336 .enumerate()
337 .for_each(|(i, ring)| ring.mul_scalar::<REDUCE>(&a.0[i], b, &mut c.0[i]));
338 }
339
340 #[inline(always)]
341 pub fn mul_scalar_inplace<const REDUCE: REDUCEMOD>(&self, a: &u64, b: &mut PolyRNS<u64>) {
342 debug_assert!(
343 b.level() >= self.level(),
344 "b.level()={} < self.level()={}",
345 b.level(),
346 self.level()
347 );
348 self.0
349 .iter()
350 .enumerate()
351 .for_each(|(i, ring)| ring.mul_scalar_inplace::<REDUCE>(a, &mut b.0[i]));
352 }
353}