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 rescaling_constant(&self) -> Vec<Barrett<u64>> {
 23        let level = self.level();
 24        let q_scale: u64 = self.0[level].modulus.q;
 25        (0..level).map(|i| {self.0[i].modulus.barrett.prepare(self.0[i].modulus.q - self.0[i].modulus.inv(q_scale))}).collect()
 26    }
 27
 28    pub fn set_poly_from_bigint(&self, coeffs: &[BigInt], step:usize, a: &mut PolyRNS<u64>){
 29        let level = self.level();
 30        assert!(level <= a.level(), "invalid level: level={} > a.level()={}", level, a.level());
 31        (0..level).for_each(|i|{self.0[i].from_bigint(coeffs, step, a.at_mut(i))});
 32    }
 33}
 34
 35impl RingRNS<'_, u64>{
 36    
 37    #[inline(always)]
 38    pub fn add<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &PolyRNS<u64>, c: &mut PolyRNS<u64>){
 39        let level: usize = self.level();
 40        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
 41        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
 42        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
 43        debug_assert!(c.level() >= level, "c.level()={} < level={}", c.level(), level);
 44        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.add::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
 45    }
 46
 47    #[inline(always)]
 48    pub fn add_inplace<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
 49        let level: usize = self.level();
 50        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
 51        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
 52        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
 53        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.add_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
 54    }
 55
 56    #[inline(always)]
 57    pub fn sub<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &PolyRNS<u64>, c: &mut PolyRNS<u64>){
 58        let level: usize = self.level();
 59        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
 60        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
 61        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
 62        debug_assert!(c.level() >= level, "c.level()={} < level={}", c.level(), level);
 63        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.sub::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
 64    }
 65
 66    #[inline(always)]
 67    pub fn sub_inplace<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
 68        let level: usize = self.level();
 69        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
 70        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
 71        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
 72        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.sub_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
 73    }
 74
 75    #[inline(always)]
 76    pub fn neg<const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
 77        let level: usize = self.level();
 78        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
 79        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
 80        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
 81        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.neg::<REDUCE>(&a.0[i], &mut b.0[i]));
 82    }
 83
 84    #[inline(always)]
 85    pub fn neg_inplace<const REDUCE: REDUCEMOD>(&self, a: &mut PolyRNS<u64>){
 86        let level: usize = self.level();
 87        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
 88        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
 89        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.neg_inplace::<REDUCE>(&mut a.0[i]));
 90    }
 91
 92    #[inline(always)]
 93    pub fn  mul_montgomery_external<const REDUCE:REDUCEMOD>(&self, a:&PolyRNS<Montgomery<u64>>, b:&PolyRNS<u64>, c: &mut PolyRNS<u64>){
 94        let level: usize = self.level();
 95        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
 96        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
 97        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
 98        debug_assert!(c.level() >= level, "c.level()={} < level={}", c.level(), level);
 99        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]));
100    }
101
102    #[inline(always)]
103    pub fn mul_montgomery_external_inplace<const REDUCE:REDUCEMOD>(&self, a:&PolyRNS<Montgomery<u64>>, b:&mut PolyRNS<u64>){
104        let level: usize = self.level();
105        debug_assert!(self.max_level() <= level, "max_level={} < level={}", self.max_level(), level);
106        debug_assert!(a.level() >= level, "a.level()={} < level={}", a.level(), level);
107        debug_assert!(b.level() >= level, "b.level()={} < level={}", b.level(), level);
108        self.0.iter().take(level + 1).enumerate().for_each(|(i, ring)| ring.mul_montgomery_external_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
109    }
110}