M math/src/ring/impl_u64/packing.rs → 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<u64>{
@@ -39,22 +41,55 @@
         x_pow
     }
 
-    pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(&self, polys: HashMap<Poly<u64>, usize>, log_gap: usize) -> Poly<u64>{
+    pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(&self, polys: &mut HashMap<usize, Poly<u64>>, log_gap: usize) -> Poly<u64>{
         
         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<usize> = polys.keys().copied().collect();
+        keys.sort();
+
+        let mut gap = 0usize;
+
+        if keys.len() > 1{
+            gap = max_pow2_gap(&keys);
+        }else{
+            gap = 1<<log_n;
+        }
+
+        let log_gap: usize = gap.log2();
+
         if !ZEROGARBAGE{
             if gap > 0 {
                 log_end -= log_gap;
             }
         }
-        */
         
-        let n_inv = self.modulus.inv(1<<(log_end - log_start));
+        let n_inv: Barrett<u64> = self.modulus.barrett.prepare(self.modulus.inv(1<<(log_end - log_start)));
+
+        for (_, poly) in polys.iter_mut() {
+            if !NTT{
+                self.ntt_inplace::<true>(poly);
+            }
+
+            self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv,  poly);
+        }
 
         Poly::<u64>::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()
 }