poulpy/math/src/ring/impl_u64/ring_rns.rs
math/src/ring/impl_u64/ring_rns.rsRaw
1use crate::ring::{Ring, RingRNS};
2use crate::poly::PolyRNS;
3use crate::modulus::montgomery::Montgomery;
4use crate::modulus::barrett::BarrettRNS;
5use crate::modulus::REDUCEMOD;
6use num_bigint::BigInt;
7
8pub fn new_rings(n: usize, moduli: Vec<u64>) -> Vec<Ring<u64>>{
9 assert!(!moduli.is_empty(), "moduli cannot be empty");
10 let rings: Vec<Ring<u64>> = moduli
11 .into_iter()
12 .map(|prime| Ring::new(n, prime, 1))
13 .collect();
14 return rings
15}
16
17impl<'a> RingRNS<'a, u64>{
18 pub fn new(rings:&'a [Ring<u64>]) -> Self{
19 RingRNS(rings)
20 }
21
22 pub fn modulus(&self) -> BigInt{
23 let mut modulus = BigInt::from(1);
24 self.0.iter().enumerate().for_each(|(_, r)|modulus *= BigInt::from(r.modulus.q));
25 modulus
26 }
27
28 pub fn rescaling_constant(&self) -> BarrettRNS<u64> {
29 let level = self.level();
30 let q_scale: u64 = self.0[level].modulus.q;
31 BarrettRNS((0..level).map(|i| {self.0[i].modulus.barrett.prepare(self.0[i].modulus.q - self.0[i].modulus.inv(q_scale))}).collect())
32 }
33
34 pub fn from_bigint_inplace(&self, coeffs: &[BigInt], step:usize, a: &mut PolyRNS<u64>){
35 let level = self.level();
36 assert!(level <= a.level(), "invalid level: level={} > a.level()={}", level, a.level());
37 (0..level).for_each(|i|{self.0[i].from_bigint(coeffs, step, a.at_mut(i))});
38 }
39
40 pub fn to_bigint_inplace(&self, a: &PolyRNS<u64>, step: usize, coeffs: &mut [BigInt]){
41 assert!(step <= a.n(), "invalid step: step={} > a.n()={}", step, a.n());
42 assert!(coeffs.len() <= a.n() / step, "invalid coeffs: coeffs.len()={} > a.n()/step={}", coeffs.len(), a.n()/step);
43
44 let mut inv_crt: Vec<BigInt> = vec![BigInt::default(); self.level()+1];
45 let q_big: BigInt = self.modulus();
46 let q_big_half: BigInt = &q_big>>1;
47
48 inv_crt.iter_mut().enumerate().for_each(|(i, a)|{
49 let qi_big = BigInt::from(self.0[i].modulus.q);
50 *a = &q_big / &qi_big;
51 *a *= a.modinv(&qi_big).unwrap();
52 });
53
54 (0..self.n()).step_by(step).enumerate().for_each(|(i, j)|{
55 coeffs[j] = BigInt::from(a.at(0).0[i]) * &inv_crt[0];
56 (1..self.level()+1).for_each(|k|{
57 coeffs[j] += BigInt::from(a.at(k).0[i] * &inv_crt[k]);
58 });
59 coeffs[j] %= &q_big;
60 if &coeffs[j] >= &q_big_half{
61 coeffs[j] -= &q_big;
62 }
63 });
64 }
65}
66
67impl RingRNS<'_, u64>{
68 pub fn ntt_inplace<const LAZY:bool>(&self, a: &mut PolyRNS<u64>){
69 self.0.iter().enumerate().for_each(|(i, ring)| ring.ntt_inplace::<LAZY>(&mut a.0[i]));
70 }
71
72 pub fn intt_inplace<const LAZY:bool>(&self, a: &mut PolyRNS<u64>){
73 self.0.iter().enumerate().for_each(|(i, ring)| ring.intt_inplace::<LAZY>(&mut a.0[i]));
74 }
75
76 pub fn ntt<const LAZY:bool>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
77 self.0.iter().enumerate().for_each(|(i, ring)| ring.ntt::<LAZY>(&a.0[i], &mut b.0[i]));
78 }
79
80 pub fn intt<const LAZY:bool>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
81 self.0.iter().enumerate().for_each(|(i, ring)| ring.intt::<LAZY>(&a.0[i], &mut b.0[i]));
82 }
83}
84
85impl RingRNS<'_, u64>{
86
87 #[inline(always)]
88 pub fn add<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &PolyRNS<u64>, c: &mut PolyRNS<u64>){
89 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
90 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
91 debug_assert!(c.level() >= self.level(), "c.level()={} < self.level()={}", c.level(), self.level());
92 self.0.iter().enumerate().for_each(|(i, ring)| ring.add::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
93 }
94
95 #[inline(always)]
96 pub fn add_inplace<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
97 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
98 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
99 self.0.iter().enumerate().for_each(|(i, ring)| ring.add_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
100 }
101
102 #[inline(always)]
103 pub fn sub<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &PolyRNS<u64>, c: &mut PolyRNS<u64>){
104 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
105 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
106 debug_assert!(c.level() >= self.level(), "c.level()={} < self.level()={}", c.level(), self.level());
107 self.0.iter().enumerate().for_each(|(i, ring)| ring.sub::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
108 }
109
110 #[inline(always)]
111 pub fn sub_inplace<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
112 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
113 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
114 self.0.iter().enumerate().for_each(|(i, ring)| ring.sub_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
115 }
116
117 #[inline(always)]
118 pub fn neg<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
119 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
120 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
121 self.0.iter().enumerate().for_each(|(i, ring)| ring.neg::<REDUCE>(&a.0[i], &mut b.0[i]));
122 }
123
124 #[inline(always)]
125 pub fn neg_inplace<const REDUCE: REDUCEMOD>(&self, a: &mut PolyRNS<u64>){
126 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
127 self.0.iter().enumerate().for_each(|(i, ring)| ring.neg_inplace::<REDUCE>(&mut a.0[i]));
128 }
129
130 #[inline(always)]
131 pub fn mul_montgomery_external<const REDUCE:REDUCEMOD>(&self, a:&PolyRNS<Montgomery<u64>>, b:&PolyRNS<u64>, c: &mut PolyRNS<u64>){
132 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
133 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
134 debug_assert!(c.level() >= self.level(), "c.level()={} < self.level()={}", c.level(), self.level());
135 self.0.iter().enumerate().for_each(|(i, ring)| ring.mul_montgomery_external::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
136 }
137
138 #[inline(always)]
139 pub fn mul_montgomery_external_inplace<const REDUCE:REDUCEMOD>(&self, a:&PolyRNS<Montgomery<u64>>, b:&mut PolyRNS<u64>){
140 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
141 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
142 self.0.iter().enumerate().for_each(|(i, ring)| ring.mul_montgomery_external_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
143 }
144
145 #[inline(always)]
146 pub fn mul_scalar<const REDUCE:REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &u64, c: &mut PolyRNS<u64>){
147 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
148 debug_assert!(c.level() >= self.level(), "b.level()={} < self.level()={}", c.level(), self.level());
149 self.0.iter().enumerate().for_each(|(i, ring)| ring.mul_scalar::<REDUCE>(&a.0[i], b, &mut c.0[i]));
150 }
151
152 #[inline(always)]
153 pub fn mul_scalar_inplace<const REDUCE:REDUCEMOD>(&self, a: &u64, b: &mut PolyRNS<u64>){
154 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
155 self.0.iter().enumerate().for_each(|(i, ring)| ring.mul_scalar_inplace::<REDUCE>(a, &mut b.0[i]));
156 }
157}