added sampling

math/Cargo.toml

@@ -4,12 +4,14 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+num = "0.4.3"
 primality-test = "0.3.0"
 num-bigint = "0.4.6"
 num-traits = "0.2.19"
 prime_factorization = "1.0.5"
-itertools = "0.13.0"
+itertools = "0.14.0"
 criterion = "0.5.1"
+sampling = { path = "../sampling" }
 
 [[bench]]
 name = "ntt"
@@ -21,4 +23,8 @@ harness = false
 
 [[bench]]
 name = "ring_rns"
+harness = false
+
+[[bench]]
+name = "sampling"
 harness = false

math/benches/sampling.rs

@@ -0,0 +1,40 @@
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use math::ring::{Ring, RingRNS};
+use math::ring::impl_u64::ring_rns::new_rings;
+use math::poly::PolyRNS;
+use sampling::source::Source;
+
+fn fill_uniform(c: &mut Criterion) {
+    fn runner(r: RingRNS<u64>) -> Box<dyn FnMut() + '_> {
+        
+        let mut a: PolyRNS<u64> = r.new_polyrns();
+        let seed: [u8; 32] = [0;32];
+        let mut source: Source = Source::new(seed);
+
+        Box::new(move || {
+            r.fill_uniform(&mut source, &mut a);
+        })
+    }
+
+    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> = c.benchmark_group("fill_uniform");
+    for log_n in 11..18 {
+
+        let n = 1<<log_n;
+        let moduli: Vec<u64> = vec![0x1fffffffffe00001u64, 0x1fffffffffc80001u64, 0x1fffffffffb40001, 0x1fffffffff500001];
+        let rings: Vec<Ring<u64>> = new_rings(n, moduli);
+        let ring_rns: RingRNS<'_, u64> = RingRNS::new(&rings);
+
+        let runners = [
+            (format!("prime/n={}/level={}", n, ring_rns.level()), {
+                runner(ring_rns)
+            }),
+        ];
+
+        for (name, mut runner) in runners {
+            b.bench_with_input(name, &(), |b, _| b.iter(&mut runner));
+        }
+    }
+}
+
+criterion_group!(benches, fill_uniform);
+criterion_main!(benches);

math/src/dft/ntt.rs

@@ -150,7 +150,6 @@ impl Table<u64>{
         debug_assert!(*b < self.four_q, "b:{} q:{}", b, self.four_q);
         a.reduce_once_assign(self.two_q);
         let bt: u64 = self.prime.barrett.mul_external::<NONE>(t, *b);
-        debug_assert!(bt < self.two_q, "bt:{} two_q:{}", bt, self.two_q);
         *b = *a + self.two_q-bt;
         *a += bt;
         if !LAZY {

math/src/lib.rs

@@ -5,6 +5,7 @@ pub mod modulus;
 pub mod dft;
 pub mod ring;
 pub mod poly;
+pub mod scalar;
 
 pub const CHUNK: usize= 8;
 

math/src/modulus.rs

@@ -15,6 +15,7 @@ pub const BARRETTLAZY: REDUCEMOD = 5;
 pub trait WordOps<O>{
     fn log2(self) -> O;
     fn reverse_bits_msb(self, n:u32) -> O;
+    fn mask(self) -> O;
 }
 
 impl WordOps<u64> for u64{
@@ -26,6 +27,10 @@ impl WordOps<u64> for u64{
     fn reverse_bits_msb(self, n: u32) -> u64{
         self.reverse_bits() >> (usize::BITS - n)
     }
+    #[inline(always)]
+    fn mask(self) -> u64{
+        (1<<self.log2())-1
+    }
 }
 
 impl WordOps<usize> for usize{
@@ -37,6 +42,10 @@ impl WordOps<usize> for usize{
     fn reverse_bits_msb(self, n: u32) -> usize{
         self.reverse_bits() >> (usize::BITS - n)
     }
+    #[inline(always)]
+    fn mask(self) -> usize{
+        (1<<self.log2())-1
+    }
 }
 
 pub trait ReduceOnce<O>{

math/src/modulus/barrett.rs

@@ -14,8 +14,6 @@ impl<O> Barrett<O> {
     }
 }
 
-pub struct BarrettRNS<O>(pub Vec<Barrett<O>>);
-
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct BarrettPrecomp<O>{
     pub q: O,

math/src/ring.rs

@@ -1,17 +1,18 @@
 pub mod impl_u64;
 
+use num::traits::Unsigned;
 use crate::modulus::prime::Prime;
 use crate::poly::{Poly, PolyRNS};
 use crate::dft::DFT;
 
 
-pub struct Ring<O>{
+pub struct Ring<O: Unsigned>{
     pub n:usize,
     pub modulus:Prime<O>,
     pub dft:Box<dyn DFT<O>>,
 }
 
-impl<O> Ring<O>{
+impl<O: Unsigned> Ring<O>{
     pub fn n(&self) -> usize{
         return self.n
     }
@@ -21,9 +22,9 @@ impl<O> Ring<O>{
     }
 }
 
-pub struct RingRNS<'a, O>(& 'a [Ring<O>]);
+pub struct RingRNS<'a, O: Unsigned>(& 'a [Ring<O>]);
 
-impl<O: Copy> RingRNS<'_, O> {
+impl<O: Unsigned> RingRNS<'_, O> {
 
     pub fn n(&self) -> usize{
         self.0[0].n()

math/src/ring/impl_u64/mod.rs

@@ -1,4 +1,5 @@
 pub mod automorphism;
 pub mod ring;
 pub mod ring_rns;
-pub mod rescaling_rns;
+pub mod rescaling_rns;
+pub mod sampling;

math/src/ring/impl_u64/rescaling_rns.rs

@@ -1,6 +1,7 @@
 use crate::ring::RingRNS;
-use crate::poly::{Poly, PolyRNS};
-use crate::modulus::barrett::BarrettRNS;
+use crate::poly::PolyRNS;
+use crate::modulus::barrett::Barrett;
+use crate::scalar::ScalarRNS;
 use crate::modulus::ONCE;
 extern crate test;
 
@@ -12,7 +13,7 @@ impl RingRNS<'_, u64>{
         assert!(b.level() >= a.level()-1, "invalid input b: b.level()={} < a.level()-1={}", b.level(), a.level()-1);
         let level = self.level();
         self.0[level].intt::<false>(a.at(level), buf.at_mut(0));
-        let rescaling_constants: BarrettRNS<u64> = self.rescaling_constant();
+        let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
         let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
         for (i, r) in self.0[0..level].iter().enumerate(){
             r.ntt::<false>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
@@ -25,7 +26,7 @@ impl RingRNS<'_, u64>{
     pub fn div_floor_by_last_modulus_ntt_inplace(&self, buf: &mut PolyRNS<u64>, b: &mut PolyRNS<u64>){
         let level = self.level();
         self.0[level].intt::<true>(b.at(level), buf.at_mut(0));
-        let rescaling_constants: BarrettRNS<u64> = self.rescaling_constant();
+        let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
         let (buf_ntt_q_scaling, buf_ntt_qi_scaling) = buf.0.split_at_mut(1);
         for (i, r) in self.0[0..level].iter().enumerate(){
             r.ntt::<true>(&buf_ntt_q_scaling[0], &mut buf_ntt_qi_scaling[0]);
@@ -37,7 +38,7 @@ impl RingRNS<'_, u64>{
     pub fn div_floor_by_last_modulus(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
         assert!(b.level() >= a.level()-1, "invalid input b: b.level()={} < a.level()-1={}", b.level(), a.level()-1);
         let level = self.level();
-        let rescaling_constants:crate::modulus::barrett::BarrettRNS<u64>  = self.rescaling_constant();
+        let rescaling_constants:ScalarRNS<Barrett<u64>>  = self.rescaling_constant();
         for (i, r) in self.0[0..level].iter().enumerate(){
             r.sum_aqqmb_prod_c_scalar_barrett::<ONCE>(a.at(level), a.at(i), &rescaling_constants.0[i], b.at_mut(i));
         }
@@ -46,7 +47,7 @@ impl RingRNS<'_, u64>{
     /// Updates a to floor(b / q[b.level()]).
     pub fn div_floor_by_last_modulus_inplace(&self, a: &mut PolyRNS<u64>){
         let level = self.level();
-        let rescaling_constants: BarrettRNS<u64> = self.rescaling_constant();
+        let rescaling_constants: ScalarRNS<Barrett<u64>> = self.rescaling_constant();
         let (a_i, a_level) = a.split_at_mut(level);
         for (i, r) in self.0[0..level].iter().enumerate(){
             r.sum_aqqmb_prod_c_scalar_barrett_inplace::<ONCE>(&a_level[0], &rescaling_constants.0[i], &mut a_i[i]);
@@ -77,8 +78,10 @@ impl RingRNS<'_, u64>{
 #[cfg(test)]
 mod tests {
     use num_bigint::BigInt;
+    use num_bigint::Sign;
     use crate::ring::Ring;
     use crate::ring::impl_u64::ring_rns::new_rings;
+    use sampling::source::Source;
     use super::*;
 
     #[test]
@@ -86,21 +89,20 @@ mod tests {
         let n = 1<<10;
         let moduli: Vec<u64> = vec![0x1fffffffffc80001u64, 0x1fffffffffe00001u64];
         let rings: Vec<Ring<u64>> = new_rings(n, moduli);
-        let ring_rns = RingRNS::new(&rings);
+        let ring_rns: RingRNS<'_, u64> = RingRNS::new(&rings);
+        let seed: [u8; 32] = [0;32];
+        let mut source: Source = Source::new(seed);
 
         let mut a: PolyRNS<u64> = ring_rns.new_polyrns();
         let mut b: PolyRNS<u64> = ring_rns.new_polyrns();
         let mut c: PolyRNS<u64> = ring_rns.at_level(ring_rns.level()-1).new_polyrns();
 
-        // Allocates an rns poly with values [0..n]
-        let mut coeffs_a: Vec<BigInt> = (0..n).map(|i|{BigInt::from(i)}).collect();
-        ring_rns.from_bigint_inplace(&coeffs_a, 1, &mut a);
+        // Allocates a random PolyRNS
+        ring_rns.fill_uniform(&mut source, &mut a);
 
-        // Scales by q_level both a and coeffs_a
-        let scalar: u64 = ring_rns.0[ring_rns.level()].modulus.q;
-        ring_rns.mul_scalar_inplace::<ONCE>(&scalar, &mut a);
-        let scalar_big = BigInt::from(scalar);
-        coeffs_a.iter_mut().for_each(|a|{*a *= &scalar_big});
+        // Maps PolyRNS to [BigInt]
+        let mut coeffs_a: Vec<BigInt> = (0..n).map(|i|{BigInt::from(i)}).collect();
+        ring_rns.at_level(a.level()).to_bigint_inplace(&a, 1, &mut coeffs_a);
 
         // Performs c = intt(ntt(a) / q_level)
         ring_rns.ntt_inplace::<false>(&mut a);
@@ -112,7 +114,14 @@ mod tests {
         ring_rns.at_level(c.level()).to_bigint_inplace(&c, 1, &mut coeffs_c);
 
         // Performs floor division on a
-        coeffs_a.iter_mut().for_each(|a|{*a /= &scalar_big});
+        let scalar_big = BigInt::from(ring_rns.0[ring_rns.level()].modulus.q);
+        coeffs_a.iter_mut().for_each(|a|{
+            // Emulates floor division in [0, q-1] and maps to [-(q-1)/2, (q-1)/2-1]
+            *a /= &scalar_big;
+            if a.sign() == Sign::Minus {
+                *a -= 1;
+            }
+        });
 
     	assert!(coeffs_a == coeffs_c);
     }

math/src/ring/impl_u64/ring_rns.rs

@@ -1,7 +1,8 @@
 use crate::ring::{Ring, RingRNS};
 use crate::poly::PolyRNS;
 use crate::modulus::montgomery::Montgomery;
-use crate::modulus::barrett::BarrettRNS;
+use crate::modulus::barrett::Barrett;
+use crate::scalar::ScalarRNS;
 use crate::modulus::REDUCEMOD;
 use num_bigint::BigInt;
 
@@ -25,10 +26,10 @@ impl<'a> RingRNS<'a, u64>{
         modulus
     }
 
-    pub fn rescaling_constant(&self) -> BarrettRNS<u64> {
+    pub fn rescaling_constant(&self) -> ScalarRNS<Barrett<u64>> {
         let level = self.level();
         let q_scale: u64 = self.0[level].modulus.q;
-        BarrettRNS((0..level).map(|i| {self.0[i].modulus.barrett.prepare(self.0[i].modulus.q - self.0[i].modulus.inv(q_scale))}).collect())
+        ScalarRNS((0..level).map(|i| {self.0[i].modulus.barrett.prepare(self.0[i].modulus.q - self.0[i].modulus.inv(q_scale))}).collect())
     }
 
     pub fn from_bigint_inplace(&self, coeffs: &[BigInt], step:usize, a: &mut PolyRNS<u64>){

math/src/ring/impl_u64/sampling.rs

@@ -0,0 +1,18 @@
+use sampling::source::Source;
+use crate::modulus::WordOps;
+use crate::ring::{Ring, RingRNS};
+use crate::poly::{Poly, PolyRNS};
+
+impl Ring<u64>{
+    pub fn fill_uniform(&self, source: &mut Source, a: &mut Poly<u64>){
+        let max:u64  = self.modulus.q;
+        let mask: u64 = max.mask();
+        a.0.iter_mut().for_each(|a|{*a = source.next_u64n(max, mask)});
+    }
+}
+
+impl RingRNS<'_, u64>{
+    pub fn fill_uniform(&self, source: &mut Source, a: &mut PolyRNS<u64>){
+        self.0.iter().enumerate().for_each(|(i, r)|{r.fill_uniform(source, a.at_mut(i))});
+    }
+}

math/src/scalar.rs

@@ -0,0 +1,2 @@
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ScalarRNS<O>(pub Vec<O>);