arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/math/src/ring/impl_u64/rescaling_rns.rs
math/src/ring/impl_u64/rescaling_rns.rsRaw
  1use crate::ring::Ring;
  2use crate::ring::RingRNS;
  3use crate::poly::PolyRNS;
  4use crate::modulus::barrett::Barrett;
  5use crate::scalar::ScalarRNS;
  6use crate::modulus::ONCE;
  7extern crate test;
  8
  9impl RingRNS<'_, u64>{
 10
 11    /// Updates b to floor(a / q[b.level()]).
 12    pub fn div_floor_by_last_modulus<const NTT:bool>(&self, a: &PolyRNS<u64>, buf: &mut PolyRNS<u64>, b: &mut PolyRNS<u64>){
 13        debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
 14        debug_assert!(b.level() >= a.level()-1, "invalid input b: b.level()={} < a.level()-1={}", b.level(), a.level()-1);
 15
 16        let level = self.level();
 17        let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
 18        
 19        if NTT{
 20            let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
 21            self.0[level].intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
 22            for (i, r) in self.0[0..level].iter().enumerate(){
 23                r.ntt::<false>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
 24                r.sum_aqqmb_prod_c_scalar_barrett::<ONCE>(&buf_ntt_qi_scaling[0], a.at(i), &rescaling_constants.0[i], b.at_mut(i));
 25            }
 26        }else{
 27            for (i, r) in self.0[0..level].iter().enumerate(){
 28                r.sum_aqqmb_prod_c_scalar_barrett::<ONCE>(a.at(level), a.at(i), &rescaling_constants.0[i], b.at_mut(i));
 29            }
 30        }
 31    }
 32
 33    /// Updates a to floor(a / q[b.level()]).
 34    /// Expects a to be in the NTT domain.
 35    pub fn div_floor_by_last_modulus_inplace<const NTT:bool>(&self, buf: &mut PolyRNS<u64>, a: &mut PolyRNS<u64>){
 36        debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
 37
 38        let level = self.level();
 39        let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
 40
 41        if NTT{
 42            let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
 43            self.0[level].intt::<true>(a.at(level), &mut buf_ntt_q_scaling[0]);
 44            for (i, r) in self.0[0..level].iter().enumerate(){
 45                r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
 46                r.sum_aqqmb_prod_c_scalar_barrett_inplace::<ONCE>(&buf_ntt_qi_scaling[0], &rescaling_constants.0[i], a.at_mut(i));
 47            }
 48        }else{
 49            let (a_i, a_level) = buf.0.split_at_mut(level);
 50            for (i, r) in self.0[0..level].iter().enumerate(){
 51                r.sum_aqqmb_prod_c_scalar_barrett_inplace::<ONCE>(&a_level[0], &rescaling_constants.0[i], &mut a_i[i]);
 52            }
 53        }
 54    }
 55
 56    /// Updates b to floor(a / prod_{level - nb_moduli}^{level} q[i])
 57    pub fn div_floor_by_last_moduli<const NTT:bool>(&self, nb_moduli:usize, a: &PolyRNS<u64>, buf: &mut PolyRNS<u64>, c: &mut PolyRNS<u64>){
 58
 59        debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
 60        debug_assert!(c.level() >= a.level()-1, "invalid input b: b.level()={} < a.level()-1={}", c.level(), a.level()-1);
 61        debug_assert!(nb_moduli <= a.level(), "invalid input nb_moduli: nb_moduli={} > a.level()={}", nb_moduli, a.level());
 62
 63        if nb_moduli == 0{
 64            if a != c{
 65                c.copy(a);
 66            }
 67        }else{
 68            if NTT{
 69                self.intt::<false>(a, buf);
 70                (0..nb_moduli).for_each(|i|{self.at_level(self.level()-i).div_floor_by_last_modulus_inplace::<false>(&mut PolyRNS::<u64>::default(), buf)});
 71                self.at_level(self.level()-nb_moduli).ntt::<false>(buf, c);
 72            }else{
 73                self.div_floor_by_last_modulus::<false>(a, buf, c);
 74                (1..nb_moduli).for_each(|i|{self.at_level(self.level()-i).div_floor_by_last_modulus_inplace::<false>(buf, c)});
 75            }
 76        }
 77    }
 78
 79    /// Updates a to floor(a / prod_{level - nb_moduli}^{level} q[i])
 80    pub fn div_floor_by_last_moduli_inplace<const NTT:bool>(&self, nb_moduli:usize, buf: &mut PolyRNS<u64>, a: &mut PolyRNS<u64>){
 81        debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
 82        debug_assert!(nb_moduli <= a.level(), "invalid input nb_moduli: nb_moduli={} > a.level()={}", nb_moduli, a.level());
 83        if NTT{
 84            self.intt::<false>(a, buf);
 85            (0..nb_moduli).for_each(|i|{self.at_level(self.level()-i).div_floor_by_last_modulus_inplace::<false>(&mut PolyRNS::<u64>::default(), buf)});
 86            self.at_level(self.level()-nb_moduli).ntt::<false>(buf, a);
 87        }else{
 88            (0..nb_moduli).for_each(|i|{self.at_level(self.level()-i).div_floor_by_last_modulus_inplace::<false>(buf, a)});
 89        } 
 90    }
 91
 92    /// Updates b to round(a / q[b.level()]).
 93    /// Expects b to be in the NTT domain.
 94    pub fn div_round_by_last_modulus<const NTT:bool>(&self, a: &PolyRNS<u64>, buf: &mut PolyRNS<u64>, b: &mut PolyRNS<u64>){
 95        debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
 96        debug_assert!(b.level() >= a.level()-1, "invalid input b: b.level()={} < a.level()-1={}", b.level(), a.level()-1);
 97
 98        let level: usize = self.level();
 99        let r_last: &Ring<u64> = &self.0[level];
100        let q_level_half: u64 = r_last.modulus.q>>1;
101        let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
102        let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
103
104        if NTT{
105            r_last.intt::<false>(a.at(level), &mut buf_ntt_q_scaling[0]);
106            r_last.add_scalar_inplace::<ONCE>(&q_level_half, &mut buf_ntt_q_scaling[0]);
107            for (i, r) in self.0[0..level].iter().enumerate(){
108                r_last.add_scalar::<ONCE>(&buf_ntt_q_scaling[0], &q_level_half, &mut buf_ntt_qi_scaling[0]);
109                r.ntt_inplace::<false>(&mut buf_ntt_qi_scaling[0]);
110                r.sum_aqqmb_prod_c_scalar_barrett::<ONCE>(&buf_ntt_qi_scaling[0], a.at(i), &rescaling_constants.0[i], b.at_mut(i));
111            }
112        }else{
113
114        }
115        
116    }
117
118    /// Updates a to round(a / q[b.level()]).
119    /// Expects a to be in the NTT domain.
120    pub fn div_round_by_last_modulus_inplace<const NTT:bool>(&self, buf: &mut PolyRNS<u64>, a: &mut PolyRNS<u64>){
121        debug_assert!(self.level() <= a.level(), "invalid input a: self.level()={} > a.level()={}", self.level(), a.level());
122
123        let level = self.level();
124        let r_last: &Ring<u64> = &self.0[level];
125        let q_level_half: u64 = r_last.modulus.q>>1;
126        let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
127        let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
128
129        if NTT{
130            r_last.intt::<true>(a.at(level), &mut buf_ntt_q_scaling[0]);
131            r_last.add_scalar_inplace::<ONCE>(&q_level_half, &mut buf_ntt_q_scaling[0]);
132            for (i, r) in self.0[0..level].iter().enumerate(){
133                r_last.add_scalar::<ONCE>(&buf_ntt_q_scaling[0], &q_level_half, &mut buf_ntt_qi_scaling[0]);
134                r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
135                r.sum_aqqmb_prod_c_scalar_barrett_inplace::<ONCE>(&buf_ntt_qi_scaling[0], &rescaling_constants.0[i], a.at_mut(i));
136            }
137        }
138        
139    }
140}
141
142