diff --git a/math/src/ring/impl_u64/packing.rs b/math/src/ring/impl_u64/packing.rs index 77ed53e..33a8147 100644 --- a/math/src/ring/impl_u64/packing.rs +++ b/math/src/ring/impl_u64/packing.rs @@ -1,7 +1,9 @@ use std::collections::HashMap; use crate::poly::Poly; use crate::ring::Ring; -use crate::modulus::ONCE; +use crate::modulus::{ONCE, WordOps}; +use crate::modulus::barrett::Barrett; +use std::cmp::min; impl Ring{ @@ -39,22 +41,55 @@ impl Ring{ x_pow } - pub fn pack(&self, polys: HashMap, usize>, log_gap: usize) -> Poly{ + pub fn pack(&self, polys: &mut HashMap>, log_gap: usize) -> Poly{ let log_n = self.log_n(); let log_start = log_n - log_gap; - let log_end = log_n; + let mut log_end = log_n; + + let mut keys: Vec = polys.keys().copied().collect(); + keys.sort(); + + let mut gap = 0usize; + + if keys.len() > 1{ + gap = max_pow2_gap(&keys); + }else{ + gap = 1< 0 { log_end -= log_gap; } } - */ - let n_inv = self.modulus.inv(1<<(log_end - log_start)); + let n_inv: Barrett = self.modulus.barrett.prepare(self.modulus.inv(1<<(log_end - log_start))); + + for (_, poly) in polys.iter_mut() { + if !NTT{ + self.ntt_inplace::(poly); + } + + self.a_mul_b_scalar_barrett_into_a::(&n_inv, poly); + } Poly::::default() } +} + +// Returns the largest +fn max_pow2_gap(vec: &[usize]) -> usize{ + let mut gap: usize = usize::MAX; + for i in 1..vec.len(){ + let (l, r) = (vec[i-1], vec[i]); + assert!(l > r, "invalid input vec: not sorted"); + gap = min(gap, r-l); + if gap == 1{ + break; + } + }; + 1 << gap.trailing_zeros() } \ No newline at end of file