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::Barrett;
5use crate::scalar::ScalarRNS;
6use crate::modulus::REDUCEMOD;
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.iter().enumerate().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((0..level).map(|i| {self.0[i].modulus.barrett.prepare(self.0[i].modulus.q - self.0[i].modulus.inv(q_scale))}).collect())
33 }
34
35 pub fn from_bigint_inplace(&self, coeffs: &[BigInt], step:usize, a: &mut PolyRNS<u64>){
36 let level = self.level();
37 assert!(level <= a.level(), "invalid level: level={} > a.level()={}", level, a.level());
38 (0..level).for_each(|i|{self.0[i].from_bigint(coeffs, step, a.at_mut(i))});
39 }
40
41 pub fn to_bigint_inplace(&self, a: &PolyRNS<u64>, step: usize, coeffs: &mut [BigInt]){
42 assert!(step <= a.n(), "invalid step: step={} > a.n()={}", step, a.n());
43 assert!(coeffs.len() <= a.n() / step, "invalid coeffs: coeffs.len()={} > a.n()/step={}", coeffs.len(), a.n()/step);
44
45 let mut inv_crt: Vec<BigInt> = vec![BigInt::default(); self.level()+1];
46 let q_big: BigInt = self.modulus();
47 let q_big_half: BigInt = &q_big>>1;
48
49 inv_crt.iter_mut().enumerate().for_each(|(i, a)|{
50 let qi_big = BigInt::from(self.0[i].modulus.q);
51 *a = &q_big / &qi_big;
52 *a *= a.modinv(&qi_big).unwrap();
53 });
54
55 (0..self.n()).step_by(step).enumerate().for_each(|(i, j)|{
56 coeffs[j] = BigInt::from(a.at(0).0[i]) * &inv_crt[0];
57 (1..self.level()+1).for_each(|k|{
58 coeffs[j] += BigInt::from(a.at(k).0[i] * &inv_crt[k]);
59 });
60 coeffs[j] %= &q_big;
61 if &coeffs[j] >= &q_big_half{
62 coeffs[j] -= &q_big;
63 }
64 });
65 }
66}
67
68impl RingRNS<'_, u64>{
69 pub fn ntt_inplace<const LAZY:bool>(&self, a: &mut PolyRNS<u64>){
70 self.0.iter().enumerate().for_each(|(i, ring)| ring.ntt_inplace::<LAZY>(&mut a.0[i]));
71 }
72
73 pub fn intt_inplace<const LAZY:bool>(&self, a: &mut PolyRNS<u64>){
74 self.0.iter().enumerate().for_each(|(i, ring)| ring.intt_inplace::<LAZY>(&mut a.0[i]));
75 }
76
77 pub fn ntt<const LAZY:bool>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
78 self.0.iter().enumerate().for_each(|(i, ring)| ring.ntt::<LAZY>(&a.0[i], &mut b.0[i]));
79 }
80
81 pub fn intt<const LAZY:bool>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
82 self.0.iter().enumerate().for_each(|(i, ring)| ring.intt::<LAZY>(&a.0[i], &mut b.0[i]));
83 }
84}
85
86impl RingRNS<'_, u64>{
87
88 #[inline(always)]
89 pub fn add<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &PolyRNS<u64>, c: &mut PolyRNS<u64>){
90 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
91 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
92 debug_assert!(c.level() >= self.level(), "c.level()={} < self.level()={}", c.level(), self.level());
93 self.0.iter().enumerate().for_each(|(i, ring)| ring.add::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
94 }
95
96 #[inline(always)]
97 pub fn add_inplace<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
98 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
99 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
100 self.0.iter().enumerate().for_each(|(i, ring)| ring.add_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
101 }
102
103 #[inline(always)]
104 pub fn sub<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &PolyRNS<u64>, c: &mut PolyRNS<u64>){
105 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
106 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
107 debug_assert!(c.level() >= self.level(), "c.level()={} < self.level()={}", c.level(), self.level());
108 self.0.iter().enumerate().for_each(|(i, ring)| ring.sub::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
109 }
110
111 #[inline(always)]
112 pub fn sub_inplace<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
113 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
114 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
115 self.0.iter().enumerate().for_each(|(i, ring)| ring.sub_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
116 }
117
118 #[inline(always)]
119 pub fn neg<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
120 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
121 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
122 self.0.iter().enumerate().for_each(|(i, ring)| ring.neg::<REDUCE>(&a.0[i], &mut b.0[i]));
123 }
124
125 #[inline(always)]
126 pub fn neg_inplace<const REDUCE: REDUCEMOD>(&self, a: &mut PolyRNS<u64>){
127 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
128 self.0.iter().enumerate().for_each(|(i, ring)| ring.neg_inplace::<REDUCE>(&mut a.0[i]));
129 }
130
131 #[inline(always)]
132 pub fn mul_montgomery_external<const REDUCE:REDUCEMOD>(&self, a:&PolyRNS<Montgomery<u64>>, b:&PolyRNS<u64>, c: &mut PolyRNS<u64>){
133 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
134 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
135 debug_assert!(c.level() >= self.level(), "c.level()={} < self.level()={}", c.level(), self.level());
136 self.0.iter().enumerate().for_each(|(i, ring)| ring.mul_montgomery_external::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
137 }
138
139 #[inline(always)]
140 pub fn mul_montgomery_external_inplace<const REDUCE:REDUCEMOD>(&self, a:&PolyRNS<Montgomery<u64>>, b:&mut PolyRNS<u64>){
141 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
142 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
143 self.0.iter().enumerate().for_each(|(i, ring)| ring.mul_montgomery_external_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
144 }
145
146 #[inline(always)]
147 pub fn mul_scalar<const REDUCE:REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &u64, c: &mut PolyRNS<u64>){
148 debug_assert!(a.level() >= self.level(), "a.level()={} < self.level()={}", a.level(), self.level());
149 debug_assert!(c.level() >= self.level(), "b.level()={} < self.level()={}", c.level(), self.level());
150 self.0.iter().enumerate().for_each(|(i, ring)| ring.mul_scalar::<REDUCE>(&a.0[i], b, &mut c.0[i]));
151 }
152
153 #[inline(always)]
154 pub fn mul_scalar_inplace<const REDUCE:REDUCEMOD>(&self, a: &u64, b: &mut PolyRNS<u64>){
155 debug_assert!(b.level() >= self.level(), "b.level()={} < self.level()={}", b.level(), self.level());
156 self.0.iter().enumerate().for_each(|(i, ring)| ring.mul_scalar_inplace::<REDUCE>(a, &mut b.0[i]));
157 }
158}