1use crate::ring::Ring;
  2use crate::dft::ntt::Table;
  3use crate::modulus::prime::Prime;
  4use crate::modulus::montgomery::Montgomery;
  5use crate::modulus::barrett::Barrett;
  6use crate::poly::Poly;
  7use crate::modulus::REDUCEMOD;
  8use crate::modulus::VecOperations;
  9use num_bigint::BigInt;
 10use num_traits::ToPrimitive;
 11use crate::CHUNK;
 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!(step <= a.n(), "invalid step: step={} > a.n()={}", step, a.n());
 25        assert!(coeffs.len() <= a.n() / step, "invalid coeffs: coeffs.len()={} > a.n()/step={}", coeffs.len(), a.n()/step);
 26        let q_big: BigInt = BigInt::from(self.modulus.q);
 27        a.0.iter_mut().step_by(step).enumerate().for_each(|(i, v)| *v = (&coeffs[i] % &q_big).to_u64().unwrap());
 28    }
 29}
 30
 31impl Ring<u64>{
 32    pub fn ntt_inplace<const LAZY:bool>(&self, poly: &mut Poly<u64>){
 33        match LAZY{
 34            true => self.dft.forward_inplace_lazy(&mut poly.0),
 35            false => self.dft.forward_inplace(&mut poly.0)
 36        }
 37    }
 38
 39    pub fn intt_inplace<const LAZY:bool>(&self, poly: &mut Poly<u64>){
 40        match LAZY{
 41            true => self.dft.forward_inplace_lazy(&mut poly.0),
 42            false => self.dft.forward_inplace(&mut poly.0)
 43        }
 44    }
 45
 46    pub fn ntt<const LAZY:bool>(&self, poly_in: &Poly<u64>, poly_out: &mut Poly<u64>){
 47        poly_out.0.copy_from_slice(&poly_in.0);
 48        match LAZY{
 49            true => self.dft.backward_inplace_lazy(&mut poly_out.0),
 50            false => self.dft.backward_inplace(&mut poly_out.0)
 51        }
 52    }
 53
 54    pub fn intt<const LAZY:bool>(&self, poly_in: &Poly<u64>, poly_out: &mut Poly<u64>){
 55        poly_out.0.copy_from_slice(&poly_in.0);
 56        match LAZY{
 57            true => self.dft.backward_inplace_lazy(&mut poly_out.0),
 58            false => self.dft.backward_inplace(&mut poly_out.0)
 59        }
 60    }
 61}
 62
 63impl Ring<u64>{
 64
 65    #[inline(always)]
 66    pub fn add_inplace<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &mut Poly<u64>){
 67        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
 68        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
 69        self.modulus.vec_add_unary_assign::<CHUNK, REDUCE>(&a.0, &mut b.0);
 70    }
 71
 72    #[inline(always)]
 73    pub fn add<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &Poly<u64>, c: &mut Poly<u64>){
 74        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
 75        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
 76        debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
 77        self.modulus.vec_add_binary_assign::<CHUNK, REDUCE>(&a.0, &b.0, &mut c.0);
 78    }
 79
 80    #[inline(always)]
 81    pub fn sub_inplace<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &mut Poly<u64>){
 82        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
 83        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
 84        self.modulus.vec_sub_unary_assign::<CHUNK, REDUCE>(&a.0, &mut b.0);
 85    }
 86
 87    #[inline(always)]
 88    pub fn sub<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &Poly<u64>, c: &mut Poly<u64>){
 89        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
 90        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
 91        debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
 92        self.modulus.vec_sub_binary_assign::<CHUNK, REDUCE>(&a.0, &b.0, &mut c.0);
 93    }
 94
 95    #[inline(always)]
 96    pub fn neg<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &mut Poly<u64>){
 97        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
 98        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
 99        self.modulus.vec_neg_binary_assign::<CHUNK, REDUCE>(&a.0, &mut b.0);
100    }
101
102    #[inline(always)]
103    pub fn neg_inplace<const REDUCE: REDUCEMOD>(&self, a: &mut Poly<u64>){
104        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
105        self.modulus.vec_neg_unary_assign::<CHUNK, REDUCE>(&mut a.0);
106    }
107
108    #[inline(always)]
109    pub fn  mul_montgomery_external<const REDUCE:REDUCEMOD>(&self, a:&Poly<Montgomery<u64>>, b:&Poly<u64>, c: &mut Poly<u64>){
110        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
111        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
112        debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
113        self.modulus.vec_mul_montgomery_external_binary_assign::<CHUNK, REDUCE>(&a.0, &b.0, &mut c.0);
114    }
115
116    #[inline(always)]
117    pub fn mul_montgomery_external_inplace<const REDUCE:REDUCEMOD>(&self, a:&Poly<Montgomery<u64>>, b:&mut Poly<u64>){
118        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
119        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
120        self.modulus.vec_mul_montgomery_external_unary_assign::<CHUNK, REDUCE>(&a.0, &mut b.0);
121    }
122
123    #[inline(always)]
124    pub fn mul_scalar_barrett_inplace<const REDUCE:REDUCEMOD>(&self, a:&Barrett<u64>, b:&mut Poly<u64>){
125        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
126        self.modulus.vec_mul_scalar_barrett_external_unary_assign::<CHUNK, REDUCE>(a, &mut b.0);
127    }
128
129    #[inline(always)]
130    pub fn mul_scalar_barrett<const REDUCE:REDUCEMOD>(&self, a:&Barrett<u64>, b: &Poly<u64>, c:&mut Poly<u64>){
131        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
132        self.modulus.vec_mul_scalar_barrett_external_binary_assign::<CHUNK, REDUCE>(a, &b.0, &mut c.0);
133    }
134
135    #[inline(always)]
136    pub fn sum_aqqmb_prod_c_scalar_barrett<const REDUCE:REDUCEMOD>(&self, a: &Poly<u64>, b: &Poly<u64>, c: &Barrett<u64>, d: &mut Poly<u64>){
137        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
138        debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
139        debug_assert!(d.n() == self.n(), "d.n()={} != n={}", d.n(), self.n());
140        self.modulus.vec_sum_aqqmb_prod_c_scalar_barrett_assign_d::<CHUNK, REDUCE>(&a.0, &b.0, c, &mut d.0);
141    }
142}