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::rc::Rc;
  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<Rc<Ring<u64>>> = moduli
 14            .into_iter()
 15            .map(|prime| Rc::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 a_add_b_into_c<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.a_add_b_into_c::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
152    }
153
154    #[inline(always)]
155    pub fn a_add_b_into_b<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.a_add_b_into_b::<REDUCE>(&a.0[i], &mut b.0[i]));
172    }
173
174    #[inline(always)]
175    pub fn a_sub_b_into_c<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.iter().enumerate().for_each(|(i, ring)| {
200            ring.a_sub_b_into_c::<BRANGE, REDUCE>(&a.0[i], &b.0[i], &mut c.0[i])
201        });
202    }
203
204    #[inline(always)]
205    pub fn a_sub_b_into_b<const BRANGE: u8, const REDUCE: REDUCEMOD>(
206        &self,
207        a: &PolyRNS<u64>,
208        b: &mut PolyRNS<u64>,
209    ) {
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.a_sub_b_into_b::<BRANGE, REDUCE>(&a.0[i], &mut b.0[i]));
226    }
227
228    #[inline(always)]
229    pub fn a_sub_b_into_a<const BRANGE: u8, const REDUCE: REDUCEMOD>(
230        &self,
231        b: &PolyRNS<u64>,
232        a: &mut PolyRNS<u64>,
233    ) {
234        debug_assert!(
235            a.level() >= self.level(),
236            "a.level()={} < self.level()={}",
237            a.level(),
238            self.level()
239        );
240        debug_assert!(
241            b.level() >= self.level(),
242            "b.level()={} < self.level()={}",
243            b.level(),
244            self.level()
245        );
246        self.0
247            .iter()
248            .enumerate()
249            .for_each(|(i, ring)| ring.a_sub_b_into_a::<BRANGE, REDUCE>(&b.0[i], &mut a.0[i]));
250    }
251
252    #[inline(always)]
253    pub fn a_neg_into_b<const ARANGE: u8, const REDUCE: REDUCEMOD>(
254        &self,
255        a: &PolyRNS<u64>,
256        b: &mut PolyRNS<u64>,
257    ) {
258        debug_assert!(
259            a.level() >= self.level(),
260            "a.level()={} < self.level()={}",
261            a.level(),
262            self.level()
263        );
264        debug_assert!(
265            b.level() >= self.level(),
266            "b.level()={} < self.level()={}",
267            b.level(),
268            self.level()
269        );
270        self.0
271            .iter()
272            .enumerate()
273            .for_each(|(i, ring)| ring.a_neg_into_b::<ARANGE, REDUCE>(&a.0[i], &mut b.0[i]));
274    }
275
276    #[inline(always)]
277    pub fn a_neg_into_a<const ARANGE: u8, const REDUCE: REDUCEMOD>(&self, a: &mut PolyRNS<u64>) {
278        debug_assert!(
279            a.level() >= self.level(),
280            "a.level()={} < self.level()={}",
281            a.level(),
282            self.level()
283        );
284        self.0
285            .iter()
286            .enumerate()
287            .for_each(|(i, ring)| ring.a_neg_into_a::<ARANGE, REDUCE>(&mut a.0[i]));
288    }
289
290    #[inline(always)]
291    pub fn mul_montgomery_external<const REDUCE: REDUCEMOD>(
292        &self,
293        a: &PolyRNS<Montgomery<u64>>,
294        b: &PolyRNS<u64>,
295        c: &mut PolyRNS<u64>,
296    ) {
297        debug_assert!(
298            a.level() >= self.level(),
299            "a.level()={} < self.level()={}",
300            a.level(),
301            self.level()
302        );
303        debug_assert!(
304            b.level() >= self.level(),
305            "b.level()={} < self.level()={}",
306            b.level(),
307            self.level()
308        );
309        debug_assert!(
310            c.level() >= self.level(),
311            "c.level()={} < self.level()={}",
312            c.level(),
313            self.level()
314        );
315        self.0.iter().enumerate().for_each(|(i, ring)| {
316            ring.a_mul_b_montgomery_into_c::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i])
317        });
318    }
319
320    #[inline(always)]
321    pub fn mul_montgomery_external_inplace<const REDUCE: REDUCEMOD>(
322        &self,
323        a: &PolyRNS<Montgomery<u64>>,
324        b: &mut PolyRNS<u64>,
325    ) {
326        debug_assert!(
327            a.level() >= self.level(),
328            "a.level()={} < self.level()={}",
329            a.level(),
330            self.level()
331        );
332        debug_assert!(
333            b.level() >= self.level(),
334            "b.level()={} < self.level()={}",
335            b.level(),
336            self.level()
337        );
338        self.0
339            .iter()
340            .enumerate()
341            .for_each(|(i, ring)| ring.a_mul_b_montgomery_into_a::<REDUCE>(&a.0[i], &mut b.0[i]));
342    }
343
344    #[inline(always)]
345    pub fn mul_scalar<const REDUCE: REDUCEMOD>(
346        &self,
347        a: &PolyRNS<u64>,
348        b: &u64,
349        c: &mut PolyRNS<u64>,
350    ) {
351        debug_assert!(
352            a.level() >= self.level(),
353            "a.level()={} < self.level()={}",
354            a.level(),
355            self.level()
356        );
357        debug_assert!(
358            c.level() >= self.level(),
359            "b.level()={} < self.level()={}",
360            c.level(),
361            self.level()
362        );
363        self.0
364            .iter()
365            .enumerate()
366            .for_each(|(i, ring)| ring.a_mul_b_scalar_into_c::<REDUCE>(&a.0[i], b, &mut c.0[i]));
367    }
368
369    #[inline(always)]
370    pub fn mul_scalar_inplace<const REDUCE: REDUCEMOD>(&self, b: &u64, a: &mut PolyRNS<u64>) {
371        debug_assert!(
372            a.level() >= self.level(),
373            "b.level()={} < self.level()={}",
374            a.level(),
375            self.level()
376        );
377        self.0
378            .iter()
379            .enumerate()
380            .for_each(|(i, ring)| ring.a_mul_b_scalar_into_a::<REDUCE>(b, &mut a.0[i]));
381    }
382
383    #[inline(always)]
384    pub fn a_sub_b_add_scalar_mul_scalar_barrett_into_e<
385        const BRANGE: u8,
386        const REDUCE: REDUCEMOD,
387    >(
388        &self,
389        a: &PolyRNS<u64>,
390        b: &PolyRNS<u64>,
391        c: &u64,
392        d: &Barrett<u64>,
393        e: &mut PolyRNS<u64>,
394    ) {
395        debug_assert!(
396            a.level() >= self.level(),
397            "a.level()={} < self.level()={}",
398            a.level(),
399            self.level()
400        );
401        debug_assert!(
402            b.level() >= self.level(),
403            "b.level()={} < self.level()={}",
404            b.level(),
405            self.level()
406        );
407        debug_assert!(
408            e.level() >= self.level(),
409            "e.level()={} < self.level()={}",
410            e.level(),
411            self.level()
412        );
413        self.0.iter().enumerate().for_each(|(i, ring)| {
414            ring.a_sub_b_add_c_scalar_mul_d_scalar_barrett_into_e::<BRANGE, REDUCE>(
415                &a.0[i],
416                &b.0[i],
417                c,
418                d,
419                &mut e.0[i],
420            )
421        });
422    }
423
424    #[inline(always)]
425    pub fn b_sub_a_add_c_scalar_mul_d_scalar_barrett_into_a<
426        const BRANGE: u8,
427        const REDUCE: REDUCEMOD,
428    >(
429        &self,
430        b: &PolyRNS<u64>,
431        c: &u64,
432        d: &Barrett<u64>,
433        a: &mut PolyRNS<u64>,
434    ) {
435        debug_assert!(
436            a.level() >= self.level(),
437            "a.level()={} < self.level()={}",
438            a.level(),
439            self.level()
440        );
441        debug_assert!(
442            b.level() >= self.level(),
443            "b.level()={} < self.level()={}",
444            b.level(),
445            self.level()
446        );
447        self.0.iter().enumerate().for_each(|(i, ring)| {
448            ring.b_sub_a_add_c_scalar_mul_d_scalar_barrett_into_a::<BRANGE, REDUCE>(
449                &b.0[i],
450                c,
451                d,
452                &mut a.0[i],
453            )
454        });
455    }
456}