arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/math/src/ring/impl_u64/ring.rs
math/src/ring/impl_u64/ring.rsRaw
  1use crate::dft::ntt::Table;
  2use crate::modulus::barrett::Barrett;
  3use crate::modulus::montgomery::Montgomery;
  4use crate::modulus::prime::Prime;
  5use crate::modulus::VectorOperations;
  6use crate::modulus::{BARRETT, REDUCEMOD};
  7use crate::poly::Poly;
  8use crate::ring::Ring;
  9use crate::CHUNK;
 10use num_bigint::BigInt;
 11use num_traits::ToPrimitive;
 12
 13impl Ring<u64> {
 14    pub fn new(n: usize, q_base: u64, q_power: usize) -> Self {
 15        let prime: Prime<u64> = Prime::<u64>::new(q_base, q_power);
 16        Self {
 17            n: n,
 18            modulus: prime.clone(),
 19            dft: Box::new(Table::<u64>::new(prime, (2 * n) as u64)),
 20        }
 21    }
 22
 23    pub fn from_bigint(&self, coeffs: &[BigInt], step: usize, a: &mut Poly<u64>) {
 24        assert!(
 25            step <= a.n(),
 26            "invalid step: step={} > a.n()={}",
 27            step,
 28            a.n()
 29        );
 30        assert!(
 31            coeffs.len() <= a.n() / step,
 32            "invalid coeffs: coeffs.len()={} > a.n()/step={}",
 33            coeffs.len(),
 34            a.n() / step
 35        );
 36        let q_big: BigInt = BigInt::from(self.modulus.q);
 37        a.0.iter_mut()
 38            .step_by(step)
 39            .enumerate()
 40            .for_each(|(i, v)| *v = (&coeffs[i] % &q_big).to_u64().unwrap());
 41    }
 42}
 43
 44impl Ring<u64> {
 45    pub fn ntt_inplace<const LAZY: bool>(&self, poly: &mut Poly<u64>) {
 46        match LAZY {
 47            true => self.dft.forward_inplace_lazy(&mut poly.0),
 48            false => self.dft.forward_inplace(&mut poly.0),
 49        }
 50    }
 51
 52    pub fn intt_inplace<const LAZY: bool>(&self, poly: &mut Poly<u64>) {
 53        match LAZY {
 54            true => self.dft.backward_inplace_lazy(&mut poly.0),
 55            false => self.dft.backward_inplace(&mut poly.0),
 56        }
 57    }
 58
 59    pub fn ntt<const LAZY: bool>(&self, poly_in: &Poly<u64>, poly_out: &mut Poly<u64>) {
 60        poly_out.0.copy_from_slice(&poly_in.0);
 61        match LAZY {
 62            true => self.dft.forward_inplace_lazy(&mut poly_out.0),
 63            false => self.dft.forward_inplace(&mut poly_out.0),
 64        }
 65    }
 66
 67    pub fn intt<const LAZY: bool>(&self, poly_in: &Poly<u64>, poly_out: &mut Poly<u64>) {
 68        poly_out.0.copy_from_slice(&poly_in.0);
 69        match LAZY {
 70            true => self.dft.backward_inplace_lazy(&mut poly_out.0),
 71            false => self.dft.backward_inplace(&mut poly_out.0),
 72        }
 73    }
 74}
 75
 76impl Ring<u64> {
 77    #[inline(always)]
 78    pub fn add_inplace<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &mut Poly<u64>) {
 79        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
 80        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
 81        self.modulus
 82            .va_add_vb_into_vb::<CHUNK, REDUCE>(&a.0, &mut b.0);
 83    }
 84
 85    #[inline(always)]
 86    pub fn add<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &Poly<u64>, c: &mut Poly<u64>) {
 87        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
 88        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
 89        debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
 90        self.modulus
 91            .va_add_vb_into_vc::<CHUNK, REDUCE>(&a.0, &b.0, &mut c.0);
 92    }
 93
 94    #[inline(always)]
 95    pub fn add_scalar_inplace<const REDUCE: REDUCEMOD>(&self, b: &u64, a: &mut Poly<u64>) {
 96        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
 97        self.modulus.va_add_sb_into_va::<CHUNK, REDUCE>(b, &mut a.0);
 98    }
 99
100    #[inline(always)]
101    pub fn add_scalar<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &u64, c: &mut Poly<u64>) {
102        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
103        debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
104        self.modulus
105            .va_add_sb_into_vc::<CHUNK, REDUCE>(&a.0, b, &mut c.0);
106    }
107
108    #[inline(always)]
109    pub fn add_scalar_then_mul_scalar_barrett_inplace<const REDUCE: REDUCEMOD>(&self, b: &u64, c: &Barrett<u64>, a: &mut Poly<u64>) {
110        debug_assert!(a.n() == self.n(), "b.n()={} != n={}", a.n(), self.n());
111        self.modulus.va_add_sb_mul_sc_into_va::<CHUNK, REDUCE>(b, c, &mut a.0);
112    }
113
114    #[inline(always)]
115    pub fn add_scalar_then_mul_scalar_barrett<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &u64, c: &Barrett<u64>, d: &mut Poly<u64>) {
116        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
117        debug_assert!(d.n() == self.n(), "c.n()={} != n={}", d.n(), self.n());
118        self.modulus
119            .va_add_sb_mul_sc_into_vd::<CHUNK, REDUCE>(&a.0, b, c, &mut d.0);
120    }
121
122    #[inline(always)]
123    pub fn sub_inplace<const BRANGE:u8, const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &mut Poly<u64>) {
124        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
125        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
126        self.modulus
127            .va_sub_vb_into_vb::<CHUNK, BRANGE, REDUCE>(&a.0, &mut b.0);
128    }
129
130    #[inline(always)]
131    pub fn sub<const BRANGE:u8, const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &Poly<u64>, c: &mut Poly<u64>) {
132        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
133        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
134        debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
135        self.modulus
136            .va_sub_vb_into_vc::<CHUNK, BRANGE, REDUCE>(&a.0, &b.0, &mut c.0);
137    }
138
139    #[inline(always)]
140    pub fn neg<const ARANGE:u8, const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &mut Poly<u64>) {
141        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
142        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
143        self.modulus.va_neg_into_vb::<CHUNK, ARANGE, REDUCE>(&a.0, &mut b.0);
144    }
145
146    #[inline(always)]
147    pub fn neg_inplace<const ARANGE:u8,const REDUCE: REDUCEMOD>(&self, a: &mut Poly<u64>) {
148        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
149        self.modulus.va_neg_into_va::<CHUNK, ARANGE, REDUCE>(&mut a.0);
150    }
151
152    #[inline(always)]
153    pub fn mul_montgomery_external<const REDUCE: REDUCEMOD>(
154        &self,
155        a: &Poly<Montgomery<u64>>,
156        b: &Poly<u64>,
157        c: &mut Poly<u64>,
158    ) {
159        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
160        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
161        debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
162        self.modulus
163            .va_mont_mul_vb_into_vc::<CHUNK, REDUCE>(&a.0, &b.0, &mut c.0);
164    }
165
166    #[inline(always)]
167    pub fn mul_montgomery_external_inplace<const REDUCE: REDUCEMOD>(
168        &self,
169        a: &Poly<Montgomery<u64>>,
170        b: &mut Poly<u64>,
171    ) {
172        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
173        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
174        self.modulus
175            .va_mont_mul_vb_into_vb::<CHUNK, REDUCE>(&a.0, &mut b.0);
176    }
177
178    #[inline(always)]
179    pub fn mul_scalar<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &u64, c: &mut Poly<u64>) {
180        debug_assert!(a.n() == self.n(), "b.n()={} != n={}", a.n(), self.n());
181        debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
182        self.modulus.sa_barrett_mul_vb_into_vc::<CHUNK, REDUCE>(
183            &self.modulus.barrett.prepare(*b),
184            &a.0,
185            &mut c.0,
186        );
187    }
188
189    #[inline(always)]
190    pub fn mul_scalar_inplace<const REDUCE: REDUCEMOD>(&self, a: &u64, b: &mut Poly<u64>) {
191        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
192        self.modulus.sa_barrett_mul_vb_into_vb::<CHUNK, REDUCE>(
193            &self
194                .modulus
195                .barrett
196                .prepare(self.modulus.barrett.reduce::<BARRETT>(a)),
197            &mut b.0,
198        );
199    }
200
201    #[inline(always)]
202    pub fn mul_scalar_barrett_inplace<const REDUCE: REDUCEMOD>(
203        &self,
204        a: &Barrett<u64>,
205        b: &mut Poly<u64>,
206    ) {
207        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
208        self.modulus
209            .sa_barrett_mul_vb_into_vb::<CHUNK, REDUCE>(a, &mut b.0);
210    }
211
212    #[inline(always)]
213    pub fn mul_scalar_barrett<const REDUCE: REDUCEMOD>(
214        &self,
215        a: &Barrett<u64>,
216        b: &Poly<u64>,
217        c: &mut Poly<u64>,
218    ) {
219        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
220        self.modulus
221            .sa_barrett_mul_vb_into_vc::<CHUNK, REDUCE>(a, &b.0, &mut c.0);
222    }
223
224    #[inline(always)]
225    pub fn a_sub_b_mul_c_scalar_barrett<const VBRANGE: u8, const REDUCE: REDUCEMOD>(
226        &self,
227        a: &Poly<u64>,
228        b: &Poly<u64>,
229        c: &Barrett<u64>,
230        d: &mut Poly<u64>,
231    ) {
232        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
233        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
234        debug_assert!(d.n() == self.n(), "d.n()={} != n={}", d.n(), self.n());
235        self.modulus
236            .va_sub_vb_mul_sc_into_vd::<CHUNK, VBRANGE, REDUCE>(&a.0, &b.0, c, &mut d.0);
237    }
238
239    #[inline(always)]
240    pub fn a_sub_b_mul_c_scalar_barrett_inplace<const BRANGE: u8, const REDUCE: REDUCEMOD>(
241        &self,
242        a: &Poly<u64>,
243        c: &Barrett<u64>,
244        b: &mut Poly<u64>,
245    ) {
246        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
247        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
248        self.modulus
249            .va_sub_vb_mul_sc_into_vb::<CHUNK, BRANGE, REDUCE>(&a.0, c, &mut b.0);
250    }
251}