From 2b916b03c51c857d5863510796526b3984d578a0 Mon Sep 17 00:00:00 2001
From: Jean-Philippe Bossuat <jeanphilippe.bossuat@gmail.com>
Date: Wed, 8 Jan 2025 17:45:51 +0100
Subject: [PATCH] wip on plaintext ring packing

---
 math/src/modulus.rs                     | 14 +++++-
 math/src/modulus/impl_u64/operations.rs | 17 ++++++-
 math/src/ring.rs                        | 11 +++++
 math/src/ring/impl_u64/mod.rs           |  1 +
 math/src/ring/impl_u64/packing.rs       | 60 +++++++++++++++++++++++++
 math/src/ring/impl_u64/ring.rs          |  6 +++
 6 files changed, 106 insertions(+), 3 deletions(-)
 create mode 100644 math/src/ring/impl_u64/packing.rs

diff --git a/math/src/modulus.rs b/math/src/modulus.rs
index 4d9db52..7accd6e 100644
--- a/math/src/modulus.rs
+++ b/math/src/modulus.rs
@@ -93,8 +93,14 @@ pub trait ScalarOperations<O> {
     // Assigns -a to b.
     fn sa_neg_into_sb<const SARANGE: u8, const REDUCE: REDUCEMOD>(&self, a: &O, b: &mut O);
 
+    // Assigns a * 2^64 to a.
+    fn sa_prepare_montgomery_into_sa<const REDUCE: REDUCEMOD>(
+        &self,
+        a: &mut montgomery::Montgomery<O>,
+    );
+
     // Assigns a * 2^64 to b.
-    fn sa_prep_mont_into_sb<const REDUCE: REDUCEMOD>(
+    fn sa_prepare_montgomery_into_sb<const REDUCE: REDUCEMOD>(
         &self,
         a: &O,
         b: &mut montgomery::Montgomery<O>,
@@ -252,6 +258,12 @@ pub trait VectorOperations<O> {
         b: &mut [montgomery::Montgomery<O>],
     );
 
+    // vec(a) <- vec(a).
+    fn va_prepare_montgomery_into_va<const CHUNK: usize, const REDUCE: REDUCEMOD>(
+        &self,
+        a: &mut [montgomery::Montgomery<O>],
+    );
+
     // vec(c) <- vec(a) * vec(b).
     fn va_mont_mul_vb_into_vc<const CHUNK: usize, const REDUCE: REDUCEMOD>(
         &self,
diff --git a/math/src/modulus/impl_u64/operations.rs b/math/src/modulus/impl_u64/operations.rs
index 5be7ec8..2cdbe9f 100644
--- a/math/src/modulus/impl_u64/operations.rs
+++ b/math/src/modulus/impl_u64/operations.rs
@@ -95,7 +95,12 @@ impl ScalarOperations<u64> for Prime<u64> {
     }
 
     #[inline(always)]
-    fn sa_prep_mont_into_sb<const REDUCE: REDUCEMOD>(&self, a: &u64, b: &mut Montgomery<u64>) {
+    fn sa_prepare_montgomery_into_sa<const REDUCE: REDUCEMOD>(&self, a: &mut Montgomery<u64>) {
+        *a = self.montgomery.prepare::<REDUCE>(*a);
+    }
+
+    #[inline(always)]
+    fn sa_prepare_montgomery_into_sb<const REDUCE: REDUCEMOD>(&self, a: &u64, b: &mut Montgomery<u64>) {
         self.montgomery.prepare_assign::<REDUCE>(*a, b);
     }
 
@@ -325,9 +330,17 @@ impl VectorOperations<u64> for Prime<u64> {
         a: &[u64],
         b: &mut [Montgomery<u64>],
     ) {
-        apply_vv!(self, Self::sa_prep_mont_into_sb::<REDUCE>, a, b, CHUNK);
+        apply_vv!(self, Self::sa_prepare_montgomery_into_sb::<REDUCE>, a, b, CHUNK);
     }
 
+    #[inline(always)]
+    fn va_prepare_montgomery_into_va<const CHUNK: usize, const REDUCE: REDUCEMOD>(
+        &self,
+        a: &mut [Montgomery<u64>],
+    ) {
+        apply_v!(self, Self::sa_prepare_montgomery_into_sa::<REDUCE>, a, CHUNK);
+    }
+    
     #[inline(always)]
     fn va_mont_mul_vb_into_vc<const CHUNK: usize, const REDUCE: REDUCEMOD>(
         &self,
diff --git a/math/src/ring.rs b/math/src/ring.rs
index f8ef639..6d6edab 100644
--- a/math/src/ring.rs
+++ b/math/src/ring.rs
@@ -3,6 +3,7 @@ pub mod impl_u64;
 use crate::dft::DFT;
 use crate::modulus::prime::Prime;
 use crate::poly::{Poly, PolyRNS};
+use crate::modulus::WordOps;
 use num::traits::Unsigned;
 use std::sync::Arc;
 
@@ -13,6 +14,11 @@ pub struct Ring<O: Unsigned> {
 }
 
 impl<O: Unsigned> Ring<O> {
+
+    pub fn log_n(&self) -> usize{
+        return self.n().log2();
+    }
+
     pub fn n(&self) -> usize {
         return self.n;
     }
@@ -25,6 +31,11 @@ impl<O: Unsigned> Ring<O> {
 pub struct RingRNS<O: Unsigned>(pub Vec<Arc<Ring<O>>>);
 
 impl<O: Unsigned> RingRNS<O> {
+
+    pub fn log_n(&self) -> usize{
+        return self.n().log2();
+    }
+
     pub fn n(&self) -> usize {
         self.0[0].n()
     }
diff --git a/math/src/ring/impl_u64/mod.rs b/math/src/ring/impl_u64/mod.rs
index 8915bbd..c22475e 100644
--- a/math/src/ring/impl_u64/mod.rs
+++ b/math/src/ring/impl_u64/mod.rs
@@ -3,3 +3,4 @@ pub mod rescaling_rns;
 pub mod ring;
 pub mod ring_rns;
 pub mod sampling;
+pub mod packing;
diff --git a/math/src/ring/impl_u64/packing.rs b/math/src/ring/impl_u64/packing.rs
new file mode 100644
index 0000000..77ed53e
--- /dev/null
+++ b/math/src/ring/impl_u64/packing.rs
@@ -0,0 +1,60 @@
+use std::collections::HashMap;
+use crate::poly::Poly;
+use crate::ring::Ring;
+use crate::modulus::ONCE;
+
+
+impl Ring<u64>{
+
+    // Generates a vector storing {X^{2^0}, X^{2^1}, .., X^{2^log_n}}.
+    pub fn gen_x_pow_2<const NTT: bool, const INV: bool>(&self, log_n: usize) -> Vec<Poly<u64>>{
+        let mut x_pow: Vec<Poly<u64>> = Vec::<Poly<u64>>::with_capacity(log_n);
+
+        (0..log_n).for_each(|i|{
+            let mut idx: usize = 1<<i;
+
+            if INV{
+                idx = self.n() - idx;
+            }
+
+            x_pow.push(self.new_poly());
+
+            if i == 0{
+                x_pow[i].0[idx] = self.modulus.montgomery.one();
+                self.ntt_inplace::<false>(&mut x_pow[i]);
+            }else{
+                let (left, right) = x_pow.split_at_mut(i);
+                self.a_mul_b_montgomery_into_c::<ONCE>(&left[i-1], &left[i-1], &mut right[0]);
+            }
+        });
+
+        if INV{
+            self.a_neg_into_a::<1, ONCE>(&mut x_pow[0]);
+        }
+
+        if !NTT{
+            x_pow.iter_mut().for_each(|x| self.intt_inplace::<false>(x));
+        }
+
+        x_pow
+    }
+
+    pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(&self, polys: HashMap<Poly<u64>, usize>, log_gap: usize) -> Poly<u64>{
+        
+        let log_n = self.log_n();
+        let log_start = log_n - log_gap;
+        let log_end = log_n;
+
+        /* 
+        if !ZEROGARBAGE{
+            if gap > 0 {
+                log_end -= log_gap;
+            }
+        }
+        */
+        
+        let n_inv = self.modulus.inv(1<<(log_end - log_start));
+
+        Poly::<u64>::default()
+    }
+}
\ No newline at end of file
diff --git a/math/src/ring/impl_u64/ring.rs b/math/src/ring/impl_u64/ring.rs
index 5c22b7b..dcbf18f 100644
--- a/math/src/ring/impl_u64/ring.rs
+++ b/math/src/ring/impl_u64/ring.rs
@@ -198,6 +198,12 @@ impl Ring<u64> {
             .va_neg_into_va::<CHUNK, ARANGE, REDUCE>(&mut a.0);
     }
 
+    #[inline(always)]
+    pub fn a_prepare_montgomery_into_a<const REDUCE: REDUCEMOD>(&self, a: &mut Poly<Montgomery<u64>>){
+        debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
+        self.modulus.va_prepare_montgomery_into_va::<CHUNK, REDUCE>(&mut a.0);
+    }
+
     #[inline(always)]
     pub fn a_mul_b_montgomery_into_c<const REDUCE: REDUCEMOD>(
         &self,