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