From 2e446c6b76698ebb99cfaf89b842dee08b42d83b Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Tue, 14 Jan 2025 13:49:24 +0100
Subject: [PATCH] wip

---
 math/src/ring/impl_u64/packing.rs |  6 ++--
 math/tests/automorphism.rs        | 58 ++++++++-----------------------
 2 files changed, 18 insertions(+), 46 deletions(-)

diff --git a/math/src/ring/impl_u64/packing.rs b/math/src/ring/impl_u64/packing.rs
index 2bccd80..fb2f32f 100644
--- a/math/src/ring/impl_u64/packing.rs
+++ b/math/src/ring/impl_u64/packing.rs
@@ -4,12 +4,11 @@ use crate::modulus::ONCE;
 use crate::poly::Poly;
 use crate::ring::Ring;
 use std::cmp::min;
-use std::rc::Rc;
 
 impl Ring<u64> {
     pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(
         &self,
-        polys: &mut Vec<Option<Poly<u64>>>,
+        polys: &mut Vec<Option<&mut Poly<u64>>>,
         log_gap: usize,
     ) {
         let log_n: usize = self.log_n();
@@ -46,7 +45,7 @@ impl Ring<u64> {
             .prepare(self.modulus.inv(1 << (log_end - log_start)));
 
         indices.iter().for_each(|i| {
-            if let Some(poly) = polys[*i].as_mut() {
+            if let Some(poly) = polys[*i].as_deref_mut() {
                 if !NTT {
                     self.ntt_inplace::<true>(poly);
                 }
@@ -126,6 +125,7 @@ fn max_gap(vec: &[usize]) -> usize {
     gap
 }
 
+
 pub struct StreamRepacker {
     accumulators: Vec<Accumulator>,
     tmp_a: Poly<u64>,
diff --git a/math/tests/automorphism.rs b/math/tests/automorphism.rs
index b39abbe..7ea4144 100644
--- a/math/tests/automorphism.rs
+++ b/math/tests/automorphism.rs
@@ -62,74 +62,46 @@ fn packing_u64() {
     let ring: Ring<u64> = Ring::new(n, q_base, q_power);
 
     sub_test("test_packing_u64::<NTT:false>", || {
-        test_packing_full_u64::<false>(&ring)
+        test_packing_sparse_u64::<false>(&ring, 1)
     });
     sub_test("test_packing_u64::<NTT:true>", || {
-        test_packing_full_u64::<true>(&ring)
+        test_packing_sparse_u64::<true>(&ring, 1)
     });
     sub_test("test_packing_sparse_u64::<NTT:false>", || {
-        test_packing_sparse_u64::<false>(&ring)
+        test_packing_sparse_u64::<false>(&ring, 3)
     });
     sub_test("test_packing_sparse_u64::<NTT:true>", || {
-        test_packing_sparse_u64::<true>(&ring)
+        test_packing_sparse_u64::<true>(&ring, 3)
     });
 }
 
-fn test_packing_full_u64<const NTT: bool>(ring: &Ring<u64>) {
+fn test_packing_sparse_u64<const NTT: bool>(ring: &Ring<u64>, gap: usize) {
     let n: usize = ring.n();
 
-    let mut result: Vec<Option<Poly<u64>>> = vec![None; n];
+    let mut result: Vec<Option<&mut Poly<u64>>> = Vec::with_capacity(n);
+    result.resize_with(n, || None);
 
-    for i in 0..n {
-        let mut poly: Poly<u64> = ring.new_poly();
-        poly.fill(&(1 + i as u64));
+    let mut polys: Vec<Poly<u64>> = vec![ring.new_poly(); (n+gap-1)/gap];
+
+    polys.iter_mut().enumerate().for_each(|(i , poly)|{
+        poly.fill(&((1 + i*gap) as u64));
         if NTT {
-            ring.ntt_inplace::<false>(&mut poly);
+            ring.ntt_inplace::<false>(poly);
         }
-
-        result[i] = Some(poly);
-    }
+        result[i*gap] = Some(poly);
+    });
 
     ring.pack::<true, NTT>(&mut result, ring.log_n());
 
     if let Some(poly) = result[0].as_mut() {
-        if NTT {
-            ring.intt_inplace::<false>(poly);
-        }
 
-        poly.0
-            .iter()
-            .enumerate()
-            .for_each(|(i, x)| assert_eq!(*x, 1 + i as u64));
-    }
-}
-
-fn test_packing_sparse_u64<const NTT: bool>(ring: &Ring<u64>) {
-    let n: usize = ring.n();
-
-    let mut result: Vec<Option<Poly<u64>>> = vec![None; n];
-
-    let gap: usize = 3;
-
-    for i in (0..n).step_by(gap) {
-        let mut poly: Poly<u64> = ring.new_poly();
-        poly.fill(&(1 + i as u64));
-        if NTT {
-            ring.ntt_inplace::<false>(&mut poly);
-        }
-        result[i] = Some(poly);
-    }
-
-    ring.pack::<true, NTT>(&mut result, ring.log_n());
-
-    if let Some(poly) = result[0].as_mut() {
         if NTT {
             ring.intt_inplace::<false>(poly);
         }
 
         poly.0.iter().enumerate().for_each(|(i, x)| {
             if i % gap == 0 {
-                assert_eq!(*x, 1 + i as u64)
+                assert_eq!(*x, (1+i) as u64)
             } else {
                 assert_eq!(*x, 0u64)
             }