[pack]: replaced HashMap by [Option<Poly<u64>]

2026-02-10 05:06:44 +01:00 · 2025-01-10 11:48:51 +01:00
parent 0886e439e0
commit 332e07dbbf
8 changed files with 232 additions and 168 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -11,6 +11,17 @@ dependencies = [
 "memchr",
 ]

+[[package]]
+name = "alga"
+version = "0.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4f823d037a7ec6ea2197046bafd4ae150e6bc36f9ca347404f46a46823fa84f2"
+dependencies = [
+ "approx",
+ "num-complex 0.2.4",
+ "num-traits",
+]
+
 [[package]]
 name = "anes"
 version = "0.1.6"
@@ -23,6 +34,15 @@ version = "1.0.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "55cc3b69f167a1ef2e161439aa98aed94e6028e5f9a59be9a6ffb47aef1651f9"

+[[package]]
+name = "approx"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f0e60b75072ecd4168020818c0107f2857bb6c4e64252d8d3983f6263b40a5c3"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "autocfg"
 version = "1.4.0"
@@ -178,12 +198,6 @@ version = "1.13.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0"

-[[package]]
-name = "equivalent"
-version = "1.0.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5"
-
 [[package]]
 name = "getrandom"
 version = "0.2.15"
@@ -206,10 +220,10 @@ dependencies = [
 ]

 [[package]]
-name = "hashbrown"
-version = "0.15.2"
+name = "hermit-abi"
+version = "0.3.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289"
+checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024"

 [[package]]
 name = "hermit-abi"
@@ -217,23 +231,13 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc"

-[[package]]
-name = "indexmap"
-version = "2.7.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "62f822373a4fe84d4bb149bf54e584a7f4abec90e072ed49cda0edea5b95471f"
-dependencies = [
- "equivalent",
- "hashbrown",
-]
-
 [[package]]
 name = "is-terminal"
 version = "0.4.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "261f68e344040fbd0edea105bef17c66edf46f984ddb1115b775ce31be948f4b"
 dependencies = [
- "hermit-abi",
+ "hermit-abi 0.4.0",
 "libc",
 "windows-sys 0.52.0",
 ]
@@ -295,7 +299,6 @@ name = "math"
 version = "0.1.0"
 dependencies = [
 "criterion",
- "indexmap",
 "itertools 0.14.0",
 "num",
 "num-bigint",
@@ -305,6 +308,17 @@ dependencies = [
 "prime_factorization",
 "rand_distr",
 "sampling",
+ "sprs",
+]
+
+[[package]]
+name = "matrixmultiply"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9380b911e3e96d10c1f415da0876389aaf1b56759054eeb0de7df940c456ba1a"
+dependencies = [
+ "autocfg",
+ "rawpointer",
 ]

 [[package]]
@@ -313,6 +327,21 @@ version = "2.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"

+[[package]]
+name = "ndarray"
+version = "0.16.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841"
+dependencies = [
+ "matrixmultiply",
+ "num-complex 0.4.6",
+ "num-integer",
+ "num-traits",
+ "portable-atomic",
+ "portable-atomic-util",
+ "rawpointer",
+]
+
 [[package]]
 name = "num"
 version = "0.4.3"
@@ -320,7 +349,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23"
 dependencies = [
 "num-bigint",
- "num-complex",
+ "num-complex 0.4.6",
 "num-integer",
 "num-iter",
 "num-rational",
@@ -337,6 +366,16 @@ dependencies = [
 "num-traits",
 ]

+[[package]]
+name = "num-complex"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6b19411a9719e753aff12e5187b74d60d3dc449ec3f4dc21e3989c3f554bc95"
+dependencies = [
+ "autocfg",
+ "num-traits",
+]
+
 [[package]]
 name = "num-complex"
 version = "0.4.6"
@@ -387,6 +426,16 @@ dependencies = [
 "libm",
 ]

+[[package]]
+name = "num_cpus"
+version = "1.16.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43"
+dependencies = [
+ "hermit-abi 0.3.9",
+ "libc",
+]
+
 [[package]]
 name = "once_cell"
 version = "1.20.2"
@@ -427,6 +476,21 @@ dependencies = [
 "plotters-backend",
 ]

+[[package]]
+name = "portable-atomic"
+version = "1.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "280dc24453071f1b63954171985a0b0d30058d287960968b9b2aca264c8d4ee6"
+
+[[package]]
+name = "portable-atomic-util"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d8a2f0d8d040d7848a709caf78912debcc3f33ee4b3cac47d73d1e1069e83507"
+dependencies = [
+ "portable-atomic",
+]
+
 [[package]]
 name = "ppv-lite86"
 version = "0.2.20"
@@ -511,6 +575,12 @@ dependencies = [
 "rand",
 ]

+[[package]]
+name = "rawpointer"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
+
 [[package]]
 name = "rayon"
 version = "1.10.0"
@@ -615,6 +685,27 @@ dependencies = [
 "serde",
 ]

+[[package]]
+name = "smallvec"
+version = "1.13.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67"
+
+[[package]]
+name = "sprs"
+version = "0.11.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "704ef26d974e8a452313ed629828cd9d4e4fa34667ca1ad9d6b1fffa43c6e166"
+dependencies = [
+ "alga",
+ "ndarray",
+ "num-complex 0.4.6",
+ "num-traits",
+ "num_cpus",
+ "rayon",
+ "smallvec",
+]
+
 [[package]]
 name = "syn"
 version = "2.0.90"
--- a/math/Cargo.toml
+++ b/math/Cargo.toml
@@ -13,7 +13,8 @@ prime_factorization = "1.0.5"
 itertools = "0.14.0"
 criterion = "0.5.1"
 rand_distr = "0.4.3"
-indexmap = "2.7.0"
+sprs = "0.11.2"
+
 sampling = { path = "../sampling" }

 [[bench]]
--- a/math/src/modulus/impl_u64/operations.rs
+++ b/math/src/modulus/impl_u64/operations.rs
@@ -100,7 +100,11 @@ impl ScalarOperations<u64> for Prime<u64> {
    }

    #[inline(always)]
-    fn sa_prepare_montgomery_into_sb<const REDUCE: REDUCEMOD>(&self, a: &u64, b: &mut Montgomery<u64>) {
+    fn sa_prepare_montgomery_into_sb<const REDUCE: REDUCEMOD>(
+        &self,
+        a: &u64,
+        b: &mut Montgomery<u64>,
+    ) {
        self.montgomery.prepare_assign::<REDUCE>(*a, b);
    }

@@ -330,7 +334,13 @@ impl VectorOperations<u64> for Prime<u64> {
        a: &[u64],
        b: &mut [Montgomery<u64>],
    ) {
-        apply_vv!(self, Self::sa_prepare_montgomery_into_sb::<REDUCE>, a, b, CHUNK);
+        apply_vv!(
+            self,
+            Self::sa_prepare_montgomery_into_sb::<REDUCE>,
+            a,
+            b,
+            CHUNK
+        );
    }

    #[inline(always)]
@@ -338,7 +348,12 @@ impl VectorOperations<u64> for Prime<u64> {
        &self,
        a: &mut [Montgomery<u64>],
    ) {
-        apply_v!(self, Self::sa_prepare_montgomery_into_sa::<REDUCE>, a, CHUNK);
+        apply_v!(
+            self,
+            Self::sa_prepare_montgomery_into_sa::<REDUCE>,
+            a,
+            CHUNK
+        );
    }

    #[inline(always)]
--- a/math/src/ring.rs
+++ b/math/src/ring.rs
@@ -2,8 +2,8 @@ pub mod impl_u64;

 use crate::dft::DFT;
 use crate::modulus::prime::Prime;
-use crate::poly::{Poly, PolyRNS};
 use crate::modulus::WordOps;
+use crate::poly::{Poly, PolyRNS};
 use num::traits::Unsigned;
 use std::rc::Rc;

@@ -14,8 +14,7 @@ pub struct Ring<O: Unsigned> {
 }

 impl<O: Unsigned> Ring<O> {
-
-    pub fn log_n(&self) -> usize{
+    pub fn log_n(&self) -> usize {
        return self.n().log2();
    }

@@ -31,8 +30,7 @@ impl<O: Unsigned> Ring<O> {
 pub struct RingRNS<O: Unsigned>(pub Vec<Rc<Ring<O>>>);

 impl<O: Unsigned> RingRNS<O> {
-
-    pub fn log_n(&self) -> usize{
+    pub fn log_n(&self) -> usize {
        return self.n().log2();
    }

--- a/math/src/ring/impl_u64/mod.rs
+++ b/math/src/ring/impl_u64/mod.rs
@@ -1,6 +1,6 @@
 pub mod automorphism;
+pub mod packing;
 pub mod rescaling_rns;
 pub mod ring;
 pub mod ring_rns;
 pub mod sampling;
-pub mod packing;
--- a/math/src/ring/impl_u64/packing.rs
+++ b/math/src/ring/impl_u64/packing.rs
@@ -1,188 +1,143 @@
-use std::collections::HashMap;
+use crate::modulus::barrett::Barrett;
+use crate::modulus::{WordOps, ONCE};
 use crate::poly::Poly;
 use crate::ring::Ring;
-use crate::modulus::{ONCE, WordOps};
-use crate::modulus::barrett::Barrett;
 use std::cmp::min;
+use std::collections::HashSet;
 use std::mem::transmute;

-
-impl Ring<u64>{
-
+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>>{
+    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;
+        (0..log_n).for_each(|i| {
+            let mut idx: usize = 1 << i;

-            if INV{
+            if INV {
                idx = self.n() - idx;
            }

            x_pow.push(self.new_poly());

-            if i == 0{
+            if i == 0 {
                x_pow[i].0[idx] = self.modulus.montgomery.one();
                self.ntt_inplace::<false>(&mut x_pow[i]);
-            }else{
+            } 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]);
+                self.a_mul_b_montgomery_into_c::<ONCE>(&left[i - 1], &left[i - 1], &mut right[0]);
            }
        });

-        if INV{
+        if INV {
            self.a_neg_into_a::<1, ONCE>(&mut x_pow[0]);
        }

-        if !NTT{
+        if !NTT {
            x_pow.iter_mut().for_each(|x| self.intt_inplace::<false>(x));
        }

        x_pow
    }

-    pub fn pack<'a, const ZEROGARBAGE: bool, const NTT: bool>(&self, polys: &'a mut HashMap<usize, &'a mut Poly<u64>>, log_gap: usize) -> &'a Poly<u64>{
-        
+    pub fn pack<const ZEROGARBAGE: bool, const NTT: bool>(
+        &self,
+        polys: &mut [Option<Poly<u64>>],
+        log_gap: usize,
+    ) {
        let log_n: usize = self.log_n();
-        let log_nth_root: usize = log_n+1;
-        let nth_root: usize = 1<<log_nth_root;
+        let log_nth_root: usize = log_n + 1;
+        let nth_root: usize = 1 << log_nth_root;
        let log_start: usize = log_n - log_gap;
        let mut log_end: usize = log_n;

-        let mut keys: Vec<usize> = polys.keys().copied().collect();
-        keys.sort();
+        let mut indices: Vec<usize> = Vec::<usize>::new();

-        let mut gap = 0usize;
-
-        if keys.len() > 1{
-            gap = max_pow2_gap(&keys);
-        }else{
-            gap = 1<<log_n;
+        // Retrives non-empty indexes
+        polys.iter().enumerate().for_each(|(i, poly)| {
+            if Some(poly) != None {
+                indices.push(i);
            }
+        });

-        let log_gap: usize = gap.log2();
+        let gap: usize = max_gap(&indices);

-        if !ZEROGARBAGE{
+        let set: HashSet<_> = indices.into_iter().collect();
+
+        let max_pow2_gap_divisor: usize = 1 << gap.trailing_zeros();
+
+        if !ZEROGARBAGE {
            if gap > 0 {
-                log_end -= log_gap;
+                log_end -= max_pow2_gap_divisor;
            }
        }

-        let n_inv: Barrett<u64> = self.modulus.barrett.prepare(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{
+        set.iter().for_each(|i| {
+            if let Some(poly) = polys[*i].as_mut() {
+                if !NTT {
                    self.ntt_inplace::<true>(poly);
                }
-
                self.a_mul_b_scalar_barrett_into_a::<ONCE>(&n_inv, poly);
            }
+        });

        let x_pow2: Vec<Poly<u64>> = self.gen_x_pow_2::<true, false>(log_n);
        let mut tmpa: Poly<u64> = self.new_poly();
        let mut tmpb: Poly<u64> = self.new_poly();

-        for i in log_start..log_end{
+        for i in log_start..log_end {
+            let t: usize = 1 << (log_n - 1 - i);

-            let t: usize = 1<<(log_n-1-i);
+            let (polys_lo, polys_hi) = polys.split_at_mut(t);

-            for j in 0..t{
+            for j in 0..t {
+                if let Some(poly_hi) = polys_hi[j].as_mut() {
+                    self.a_mul_b_montgomery_into_a::<ONCE>(&x_pow2[log_n - i - 1], poly_hi);

-                let option_lo: Option<&&mut Poly<u64>> = polys.get(&i);
-                let option_hi: Option<&&mut Poly<u64>> = polys.get(&(i+t));
-                let mut hi_exists: bool = false;
-
-                match option_hi{
-                    Some(hi) =>{
-
-                            // Unsafe code is necessary because two mutables references are
-                            // accessed from the map.
-                            unsafe{
-                                self.a_mul_b_montgomery_into_a::<ONCE>(&x_pow2[log_n-i-1],  transmute(*hi as *const Poly<u64> as *mut Poly<u64>));
-                            }
-
-                            hi_exists = true;
-
-                            match option_lo{
-                                Some(lo) =>{
-
-                                    self.a_sub_b_into_c::<1, ONCE>(lo, hi, &mut tmpa);
-                                    
-                                    // Ensures unsafe blocks are "safe".
-                                    let ptr_hi: *mut Poly<u64> = *hi as *const Poly<u64> as *mut Poly<u64>;
-                                    let ptr_lo: *mut Poly<u64> = *lo as *const Poly<u64> as *mut Poly<u64>;
-                                    assert!(ptr_hi != ptr_lo, "something went seriously wrong");
-
-                                    unsafe{
-                                        self.a_add_b_into_b::<ONCE>(hi, transmute(ptr_lo));
-                                    }
-                                }
-                                None =>{
-                                    unsafe{
-                                        polys.insert(j, transmute(*hi as *const Poly<u64> as *mut Poly<u64>));
-                                    }
-                                },
-                            }
-
-                            polys.remove(&(j+t));
-                        }
-
-                    None =>{},  
-                }
-
-                let option_lo: Option<&&mut Poly<u64>> = polys.get(&i);
-                let option_hi: Option<&&mut Poly<u64>> = polys.get(&(i+t));
-
-                match option_lo{
-                    Some(lo) =>{
-                        
-                        let gal_el: usize = self.galois_element(1<<(i-1), i == 0, log_nth_root);
-
-                        if hi_exists{
-                            self.automorphism::<true>(&tmpa, gal_el, 2<<self.log_n(), &mut tmpb);
-                        }else{
-                            self.automorphism::<true>(*lo, gal_el, nth_root, &mut tmpa);
-                        }
-                        unsafe{
-                            self.a_add_b_into_b::<ONCE>(&tmpa, transmute(*lo as *const Poly<u64> as *mut Poly<u64>));
+                    if let Some(poly_lo) = polys_lo[j].as_mut() {
+                        self.a_sub_b_into_c::<1, ONCE>(poly_lo, poly_hi, &mut tmpa);
+                        self.a_add_b_into_b::<ONCE>(poly_hi, poly_lo);
+                    } else {
+                        std::mem::swap(&mut polys_lo[j], &mut polys_hi[j]);
                    }
                }

-                    None =>{
-                        match option_hi{
-                            Some(hi) =>{
-                                let gal_el: usize = self.galois_element(1<<(i-1), i == 0, log_nth_root);
+                if let Some(poly_lo) = polys_lo[j].as_mut() {
+                    let gal_el: usize = self.galois_element(1 << (i - 1), i == 0, log_nth_root);

-                                self.automorphism::<true>(*hi, gal_el, nth_root, &mut tmpa);
-
-                                unsafe{
-                                    self.a_sub_b_into_a::<1, ONCE>(&tmpa, transmute(*hi as *const Poly<u64> as *mut Poly<u64>))
-                                }
+                    if !polys_hi[j].is_none() {
+                        self.automorphism::<true>(&tmpa, gal_el, 2 << self.log_n(), &mut tmpb);
+                    } else {
+                        self.automorphism::<true>(poly_lo, gal_el, nth_root, &mut tmpa);
                    }

-                            None =>{}
-                        }
-                    }
-                }
-            }
-        }
+                    self.a_add_b_into_b::<ONCE>(&tmpa, poly_lo);
+                } else if let Some(poly_hi) = polys_hi[j].as_mut() {
+                    let gal_el: usize = self.galois_element(1 << (i - 1), i == 0, log_nth_root);

-        *polys.get(&0).unwrap()
+                    self.automorphism::<true>(poly_hi, gal_el, nth_root, &mut tmpa);
+                    self.a_sub_b_into_a::<1, ONCE>(&tmpa, poly_hi)
+                }
+            }
+        }
    }
 }

-
-// Returns the largest
-fn max_pow2_gap(vec: &[usize]) -> usize{
+// Returns the largest gap.
+fn max_gap(vec: &[usize]) -> usize {
    let mut gap: usize = usize::MAX;
-    for i in 1..vec.len(){
-        let (l, r) = (vec[i-1], vec[i]);
+    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{
+        gap = min(gap, r - l);
+        if gap == 1 {
            break;
        }
-    };
-    1 << gap.trailing_zeros()
+    }
+    gap
 }
--- a/math/src/ring/impl_u64/ring.rs
+++ b/math/src/ring/impl_u64/ring.rs
@@ -42,24 +42,24 @@ impl Ring<u64> {
    }

    // Returns GALOISGENERATOR^gen_1 * (-1)^gen_2 mod 2^log_nth_root.
-    pub fn galois_element(&self, gen_1: usize, gen_2: bool, log_nth_root: usize) -> usize{
+    pub fn galois_element(&self, gen_1: usize, gen_2: bool, log_nth_root: usize) -> usize {
        let mut gal_el: usize = 1;
        let mut gen_1_pow: usize = GALOISGENERATOR;
        let mut e: usize = gen_1;
-        while e > 0{
-            if e & 1 == 1{
+        while e > 0 {
+            if e & 1 == 1 {
                gal_el = gal_el.wrapping_mul(gen_1_pow);
            }

            gen_1_pow *= gen_1_pow;
-            e>>=1;
+            e >>= 1;
        }

-        let nth_root = 1<<log_nth_root;
-        gal_el &= (nth_root-1);
+        let nth_root = 1 << log_nth_root;
+        gal_el &= (nth_root - 1);

-        if gen_2{
-            return nth_root - gal_el
+        if gen_2 {
+            return nth_root - gal_el;
        }
        gal_el
    }
@@ -223,9 +223,13 @@ impl Ring<u64> {
    }

    #[inline(always)]
-    pub fn a_prepare_montgomery_into_a<const REDUCE: REDUCEMOD>(&self, a: &mut Poly<Montgomery<u64>>){
+    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);
+        self.modulus
+            .va_prepare_montgomery_into_va::<CHUNK, REDUCE>(&mut a.0);
    }

    #[inline(always)]
--- a/sampling/src/source.rs
+++ b/sampling/src/source.rs
@@ -36,7 +36,7 @@ impl Source {
    }
 }

-impl RngCore for Source{
+impl RngCore for Source {
    #[inline(always)]
    fn next_u32(&mut self) -> u32 {
        self.source.next_u32()
@@ -53,7 +53,7 @@ impl RngCore for Source{
    }

    #[inline(always)]
-    fn try_fill_bytes(&mut self, bytes: &mut [u8]) -> Result<(), rand_core::Error>{
+    fn try_fill_bytes(&mut self, bytes: &mut [u8]) -> Result<(), rand_core::Error> {
        self.source.try_fill_bytes(bytes)
    }
 }