1use crate::ring::{Ring, RingRNS};
  2use crate::poly::PolyRNS;
  3use crate::modulus::montgomery::Montgomery;
  4use crate::modulus::barrett::Barrett;
  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) -> Vec<Barrett<u64>> {
 29        let level = self.level();
 30        let q_scale: u64 = self.0[level].modulus.q;
 31        (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 set_poly_from_bigint(&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 set_bigint_from_poly(&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    
 69    #[inline(always)]
 70    pub fn add<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &PolyRNS<u64>, c: &mut PolyRNS<u64>){
 71        let level: usize = self.level();
 72        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
 73        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
 74        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
 75        debug_assert!(c.level() >= level, "c.level()={} < level={}", c.level(), level);
 76        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.add::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
 77    }
 78
 79    #[inline(always)]
 80    pub fn add_inplace<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
 81        let level: usize = self.level();
 82        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
 83        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
 84        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
 85        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.add_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
 86    }
 87
 88    #[inline(always)]
 89    pub fn sub<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &PolyRNS<u64>, c: &mut PolyRNS<u64>){
 90        let level: usize = self.level();
 91        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
 92        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
 93        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
 94        debug_assert!(c.level() >= level, "c.level()={} < level={}", c.level(), level);
 95        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.sub::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
 96    }
 97
 98    #[inline(always)]
 99    pub fn sub_inplace<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
100        let level: usize = self.level();
101        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
102        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
103        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
104        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.sub_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
105    }
106
107    #[inline(always)]
108    pub fn neg<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
109        let level: usize = self.level();
110        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
111        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
112        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
113        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.neg::<REDUCE>(&a.0[i], &mut b.0[i]));
114    }
115
116    #[inline(always)]
117    pub fn neg_inplace<const REDUCE: REDUCEMOD>(&self, a: &mut PolyRNS<u64>){
118        let level: usize = self.level();
119        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
120        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
121        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.neg_inplace::<REDUCE>(&mut a.0[i]));
122    }
123
124    #[inline(always)]
125    pub fn  mul_montgomery_external<const REDUCE:REDUCEMOD>(&self, a:&PolyRNS<Montgomery<u64>>, b:&PolyRNS<u64>, c: &mut PolyRNS<u64>){
126        let level: usize = self.level();
127        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
128        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
129        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
130        debug_assert!(c.level() >= level, "c.level()={} < level={}", c.level(), level);
131        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.mul_montgomery_external::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
132    }
133
134    #[inline(always)]
135    pub fn mul_montgomery_external_inplace<const REDUCE:REDUCEMOD>(&self, a:&PolyRNS<Montgomery<u64>>, b:&mut PolyRNS<u64>){
136        let level: usize = self.level();
137        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
138        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
139        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
140        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.mul_montgomery_external_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
141    }
142}