arnaucube/poulpy
1
0
Fork 0
mirror of https://github.com/arnaucube/poulpy.git synced 2026-02-10 05:06:44 +01:00
Code Issues Packages Projects Releases Wiki Activity

wip

Browse Source
This commit is contained in:
Jean-Philippe Bossuat
2025-01-02 18:05:16 +01:00
parent 8d84727fae
commit e23ee338c8
29 changed files with 331 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
math/Cargo.toml Normal file
View File

@@ -0,0 +1,20 @@
[package]
name = "math"
version = "0.1.0"
edition = "2021"
[dependencies]
primality-test = "0.3.0"
num-bigint = "0.4.6"
num-traits = "0.2.19"
prime_factorization = "1.0.5"
itertools = "0.13.0"
criterion = "0.5.1"
[[bench]]
name = "ntt"
harness = false
[[bench]]
name = "operations"
harness = false

122
math/benches/ntt.rs Normal file
View File

@@ -0,0 +1,122 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use math::{modulus::prime::Prime,dft::ntt::Table};
use math::dft::DFT;
fn forward_inplace(c: &mut Criterion) {
fn runner(prime_instance: Prime<u64>, nth_root: u64) -> Box<dyn FnMut()> {
let ntt_table: Table<u64> = Table::<u64>::new(prime_instance, nth_root);
let mut a: Vec<u64> = vec![0; (nth_root >> 1) as usize];
for i in 0..a.len(){
a[i] = i as u64;
}
Box::new(move || {
ntt_table.forward_inplace::<false>(&mut a)
})
}
let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> = c.benchmark_group("forward_inplace");
for log_nth_root in 11..18 {
let prime_instance: Prime<u64> = Prime::<u64>::new(0x1fffffffffe00001, 1);
let runners = [
("prime", {
runner(prime_instance, 1<<log_nth_root)
}),
];
for (name, mut runner) in runners {
let id = BenchmarkId::new(name, 1<<(log_nth_root-1));
b.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
}
}
}
fn forward_inplace_lazy(c: &mut Criterion) {
fn runner(prime_instance: Prime<u64>, nth_root: u64) -> Box<dyn FnMut()> {
let ntt_table: Table<u64> = Table::<u64>::new(prime_instance, nth_root);
let mut a: Vec<u64> = vec![0; (nth_root >> 1) as usize];
for i in 0..a.len(){
a[i] = i as u64;
}
Box::new(move || {
ntt_table.forward_inplace_lazy(&mut a)
})
}
let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> = c.benchmark_group("forward_inplace_lazy");
for log_nth_root in 11..17 {
let prime_instance: Prime<u64> = Prime::<u64>::new(0x1fffffffffe00001, 1);
let runners = [
("prime", {
runner(prime_instance, 1<<log_nth_root)
}),
];
for (name, mut runner) in runners {
let id = BenchmarkId::new(name, 1<<(log_nth_root-1));
b.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
}
}
}
fn backward_inplace(c: &mut Criterion) {
fn runner(prime_instance: Prime<u64>, nth_root: u64) -> Box<dyn FnMut()> {
let ntt_table: Table<u64> = Table::<u64>::new(prime_instance, nth_root);
let mut a: Vec<u64> = vec![0; (nth_root >> 1) as usize];
for i in 0..a.len(){
a[i] = i as u64;
}
Box::new(move || {
ntt_table.backward_inplace::<false>(&mut a)
})
}
let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> = c.benchmark_group("backward_inplace");
for log_nth_root in 11..18 {
let prime_instance: Prime<u64> = Prime::<u64>::new(0x1fffffffffe00001, 1);
let runners = [
("prime", {
runner(prime_instance, 1<<log_nth_root)
}),
];
for (name, mut runner) in runners {
let id = BenchmarkId::new(name, 1<<(log_nth_root-1));
b.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
}
}
}
fn backward_inplace_lazy(c: &mut Criterion) {
fn runner(prime_instance: Prime<u64>, nth_root: u64) -> Box<dyn FnMut()> {
let ntt_table: Table<u64> = Table::<u64>::new(prime_instance, nth_root);
let mut a: Vec<u64> = vec![0; (nth_root >> 1) as usize];
for i in 0..a.len(){
a[i] = i as u64;
}
Box::new(move || {
ntt_table.backward_inplace::<true>(&mut a)
})
}
let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> = c.benchmark_group("backward_inplace_lazy");
for log_nth_root in 11..17 {
let prime_instance: Prime<u64> = Prime::<u64>::new(0x1fffffffffe00001, 1);
let runners = [
("prime", {
runner(prime_instance, 1<<log_nth_root)
}),
];
for (name, mut runner) in runners {
let id = BenchmarkId::new(name, 1<<(log_nth_root-1));
b.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
}
}
}
criterion_group!(benches, forward_inplace, forward_inplace_lazy, backward_inplace, backward_inplace_lazy);
criterion_main!(benches);

109
math/benches/operations.rs Normal file
View File

@@ -0,0 +1,109 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use math::ring::Ring;
use math::modulus::VecOperations;
use math::modulus::montgomery::Montgomery;
use math::modulus::ONCE;
use math::CHUNK;
fn vec_add_unary(c: &mut Criterion) {
fn runner(r: Ring<u64>) -> Box<dyn FnMut()> {
let mut p0: math::poly::Poly<u64> = r.new_poly();
let mut p1: math::poly::Poly<u64> = r.new_poly();
for i in 0..p0.n(){
p0.0[i] = i as u64;
p1.0[i] = i as u64;
}
Box::new(move || {
r.modulus.vec_add_unary_assign::<CHUNK, ONCE>(&p0.0, &mut p1.0);
})
}
let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> = c.benchmark_group("add_vec_unary");
for log_n in 11..17 {
let n: usize = 1<<log_n as usize;
let q_base: u64 = 0x1fffffffffe00001u64;
let q_power: usize = 1usize;
let r: Ring<u64> = Ring::<u64>::new(n, q_base, q_power);
let runners = [
("prime", {
runner(r)
}),
];
for (name, mut runner) in runners {
let id = BenchmarkId::new(name, n);
b.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
}
}
}
fn vec_mul_montgomery_external_unary_assign(c: &mut Criterion) {
fn runner(r: Ring<u64>) -> Box<dyn FnMut()> {
let mut p0: math::poly::Poly<Montgomery<u64>> = r.new_poly();
let mut p1: math::poly::Poly<u64> = r.new_poly();
for i in 0..p0.n(){
p0.0[i] = r.modulus.montgomery.prepare::<ONCE>(i as u64);
p1.0[i] = i as u64;
}
Box::new(move || {
r.modulus.vec_mul_montgomery_external_unary_assign::<CHUNK, ONCE>(&p0.0, &mut p1.0);
})
}
let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> = c.benchmark_group("mul_vec_montgomery_external_unary_assign");
for log_n in 11..17 {
let n: usize = 1<<log_n as usize;
let q_base: u64 = 0x1fffffffffe00001u64;
let q_power: usize = 1usize;
let r: Ring<u64> = Ring::<u64>::new(n, q_base, q_power);
let runners = [
("prime", {
runner(r)
}),
];
for (name, mut runner) in runners {
let id = BenchmarkId::new(name, n);
b.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
}
}
}
fn vec_mul_montgomery_external_binary_assign(c: &mut Criterion) {
fn runner(r: Ring<u64>) -> Box<dyn FnMut()> {
let mut p0: math::poly::Poly<Montgomery<u64>> = r.new_poly();
let mut p1: math::poly::Poly<u64> = r.new_poly();
let mut p2: math::poly::Poly<u64> = r.new_poly();
for i in 0..p0.n(){
p0.0[i] = r.modulus.montgomery.prepare::<ONCE>(i as u64);
p1.0[i] = i as u64;
}
Box::new(move || {
r.modulus.vec_mul_montgomery_external_binary_assign::<CHUNK,ONCE>(&p0.0, & p1.0, &mut p2.0);
})
}
let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> = c.benchmark_group("mul_vec_montgomery_external_binary_assign");
for log_n in 11..17 {
let n: usize = 1<<log_n as usize;
let q_base: u64 = 0x1fffffffffe00001u64;
let q_power: usize = 1usize;
let r: Ring<u64> = Ring::<u64>::new(n, q_base, q_power);
let runners = [
("prime", {
runner(r)
}),
];
for (name, mut runner) in runners {
let id = BenchmarkId::new(name, n);
b.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
}
}
}
criterion_group!(benches, vec_add_unary, vec_mul_montgomery_external_unary_assign, vec_mul_montgomery_external_binary_assign);
criterion_main!(benches);

51
math/examples/main.rs Normal file
View File

@@ -0,0 +1,51 @@
use math::ring::Ring;
use math::modulus::prime::Prime;
use math::dft::ntt::Table;
fn main() {
// Example usage of `Prime<u64>`
let q_base: u64 = 65537; // Example prime base
let q_power: usize = 1; // Example power
let prime_instance: Prime<u64> = Prime::<u64>::new(q_base, q_power);
// Display the fields of `Prime` to verify
println!("Prime instance created:");
println!("q: {}", prime_instance.q());
println!("q_base: {}", prime_instance.q_base());
println!("q_power: {}", prime_instance.q_power());
let n: u64 = 32;
let nth_root: u64 = n<<1;
let ntt_table: Table<u64> = Table::<u64>::new(prime_instance, nth_root);
let mut a: Vec<u64> = vec![0; (nth_root >> 1) as usize];
for i in 0..a.len(){
a[i] = i as u64;
}
println!("{:?}", a);
ntt_table.forward_inplace::<false>(&mut a);
println!("{:?}", a);
ntt_table.backward_inplace::<false>(&mut a);
println!("{:?}", a);
let r : Ring<u64> = Ring::<u64>::new(n as usize, q_base, q_power);
let mut p0: math::poly::Poly<u64> = r.new_poly();
let mut p1: math::poly::Poly<u64> = r.new_poly();
for i in 0..p0.n(){
p0.0[i] = i as u64
}
r.automorphism(p0, (2*r.n-1) as u64, &mut p1);
println!("{:?}", p1);
}

8
math/src/dft.rs Normal file
View File

@@ -0,0 +1,8 @@
pub mod ntt;
pub trait DFT<O> {
fn forward_inplace_lazy(&self, x: &mut [O]);
fn forward_inplace(&self, x: &mut [O]);
fn backward_inplace_lazy(&self, x: &mut [O]);
fn backward_inplace(&self, x: &mut [O]);
}

275
math/src/dft/ntt.rs Normal file
View File

@@ -0,0 +1,275 @@
use crate::modulus::montgomery::Montgomery;
use crate::modulus::barrett::Barrett;
use crate::modulus::prime::Prime;
use crate::modulus::ReduceOnce;
use crate::modulus::WordOps;
use crate::modulus::ONCE;
use crate::dft::DFT;
use itertools::izip;
#[allow(dead_code)]
pub struct Table<O>{
prime:Prime<O>,
psi: O,
psi_forward_rev:Vec<Barrett<u64>>,
psi_backward_rev: Vec<Barrett<u64>>,
q:O,
two_q:O,
four_q:O,
}
impl Table< u64> {
pub fn new(prime: Prime<u64>, nth_root: u64)->Self{
assert!(nth_root&(nth_root-1) == 0, "invalid argument: nth_root = {} is not a power of two", nth_root);
let psi: u64 = prime.primitive_nth_root(nth_root);
let psi_mont: Montgomery<u64> = prime.montgomery.prepare::<ONCE>(psi);
let psi_inv_mont: Montgomery<u64> = prime.montgomery.pow(psi_mont, prime.phi-1);
let mut psi_forward_rev: Vec<Barrett<u64>> = vec![Barrett(0, 0); (nth_root >> 1) as usize];
let mut psi_backward_rev: Vec<Barrett<u64>> = vec![Barrett(0, 0); (nth_root >> 1) as usize];
psi_forward_rev[0] = prime.barrett.prepare(1);
psi_backward_rev[0] = prime.barrett.prepare(1);
let log_nth_root_half: u32 = (nth_root>>1).log2() as _;
let mut powers_forward: u64 = 1u64;
let mut powers_backward: u64 = 1u64;
for i in 1..(nth_root>>1) as usize{
let i_rev: usize = i.reverse_bits_msb(log_nth_root_half);
prime.montgomery.mul_external_assign::<ONCE>(psi_mont, &mut powers_forward);
prime.montgomery.mul_external_assign::<ONCE>(psi_inv_mont, &mut powers_backward);
psi_forward_rev[i_rev] = prime.barrett.prepare(powers_forward);
psi_backward_rev[i_rev] = prime.barrett.prepare(powers_backward);
}
let q: u64 = prime.q();
Self{
prime: prime,
psi:psi,
psi_forward_rev: psi_forward_rev,
psi_backward_rev: psi_backward_rev,
q:q,
two_q:q<<1,
four_q:q<<2,
}
}
}
impl DFT<u64> for Table<u64>{
fn forward_inplace(&self, a: &mut [u64]){
self.forward_inplace::<false>(a)
}
fn forward_inplace_lazy(&self, a: &mut [u64]){
self.forward_inplace::<true>(a)
}
fn backward_inplace(&self, a: &mut [u64]){
self.backward_inplace::<false>(a)
}
fn backward_inplace_lazy(&self, a: &mut [u64]){
self.backward_inplace::<true>(a)
}
}
impl Table<u64>{
pub fn forward_inplace<const LAZY:bool>(&self, a: &mut [u64]){
self.forward_inplace_core::<LAZY, 0, 0>(a);
}
pub fn forward_inplace_core<const LAZY: bool, const SKIPSTART: u8, const SKIPEND: u8>(&self, a: &mut [u64]) {
let n: usize = a.len();
assert!(n & n-1 == 0, "invalid x.len()= {} must be a power of two", n);
let log_n: u32 = usize::BITS - ((n as usize)-1).leading_zeros();
let start: u32 = SKIPSTART as u32;
let end: u32 = log_n - (SKIPEND as u32);
for layer in start..end {
let (m, size) = (1 << layer, 1 << (log_n - layer - 1));
let t: usize = 2*size;
if layer == log_n - 1 {
if LAZY{
izip!(a.chunks_exact_mut(t), &self.psi_forward_rev[m..]).for_each(|(a, psi)| {
let (a, b) = a.split_at_mut(size);
self.dit_inplace::<false>(&mut a[0], &mut b[0], *psi);
debug_assert!(a[0] < self.two_q, "forward_inplace_core::<LAZY=true> output {} > {} (2q-1)", a[0], self.two_q-1);
debug_assert!(b[0] < self.two_q, "forward_inplace_core::<LAZY=true> output {} > {} (2q-1)", b[0], self.two_q-1);
});
}else{
izip!(a.chunks_exact_mut(t), &self.psi_forward_rev[m..]).for_each(|(a, psi)| {
let (a, b) = a.split_at_mut(size);
self.dit_inplace::<true>(&mut a[0], &mut b[0], *psi);
self.prime.barrett.reduce_assign(&mut a[0]);
self.prime.barrett.reduce_assign(&mut b[0]);
debug_assert!(a[0] < self.q, "forward_inplace_core::<LAZY=false> output {} > {} (q-1)", a[0], self.q-1);
debug_assert!(b[0] < self.q, "forward_inplace_core::<LAZY=false> output {} > {} (q-1)", b[0], self.q-1);
});
}
} else if t >= 16{
izip!(a.chunks_exact_mut(t), &self.psi_forward_rev[m..]).for_each(|(a, psi)| {
let (a, b) = a.split_at_mut(size);
izip!(a.chunks_exact_mut(8), b.chunks_exact_mut(8)).for_each(|(a, b)| {
self.dit_inplace::<true>(&mut a[0], &mut b[0], *psi);
self.dit_inplace::<true>(&mut a[1], &mut b[1], *psi);
self.dit_inplace::<true>(&mut a[2], &mut b[2], *psi);
self.dit_inplace::<true>(&mut a[3], &mut b[3], *psi);
self.dit_inplace::<true>(&mut a[4], &mut b[4], *psi);
self.dit_inplace::<true>(&mut a[5], &mut b[5], *psi);
self.dit_inplace::<true>(&mut a[6], &mut b[6], *psi);
self.dit_inplace::<true>(&mut a[7], &mut b[7], *psi);
});
});
}else{
izip!(a.chunks_exact_mut(t), &self.psi_forward_rev[m..]).for_each(|(a, psi)| {
let (a, b) = a.split_at_mut(size);
izip!(a, b).for_each(|(a, b)| self.dit_inplace::<true>(a, b, *psi));
});
}
}
}
#[inline(always)]
fn dit_inplace<const LAZY: bool>(&self, a: &mut u64, b: &mut u64, t: Barrett<u64>) {
debug_assert!(*a < self.four_q, "a:{} q:{}", a, self.four_q);
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_lazy(t, *b);
*b = a.wrapping_add(self.two_q-bt);
*a = a.wrapping_add(bt);
if !LAZY {
a.reduce_once_assign(self.two_q);
b.reduce_once_assign(self.two_q);
}
}
pub fn backward_inplace<const LAZY:bool>(&self, a: &mut [u64]){
self.backward_inplace_core::<LAZY, 0, 0>(a);
}
pub fn backward_inplace_core<const LAZY:bool, const SKIPSTART: u8, const SKIPEND: u8>(&self, a: &mut [u64]) {
let n: usize = a.len();
assert!(n & n-1 == 0, "invalid x.len()= {} must be a power of two", n);
let log_n = usize::BITS - ((n as usize)-1).leading_zeros();
let start: u32 = SKIPEND as u32;
let end: u32 = log_n - (SKIPSTART as u32);
for layer in (start..end).rev() {
let (m, size) = (1 << layer, 1 << (log_n - layer - 1));
let t: usize = 2*size;
if layer == 0 {
let n_inv: Barrett<u64> = self.prime.barrett.prepare(self.prime.inv(n as u64));
let psi: Barrett<u64> = self.prime.barrett.prepare(self.prime.barrett.mul_external(n_inv, self.psi_backward_rev[1].0));
izip!(a.chunks_exact_mut(2 * size)).for_each(
|a| {
let (a, b) = a.split_at_mut(size);
izip!(a.chunks_exact_mut(8), b.chunks_exact_mut(8)).for_each(|(a, b)| {
self.dif_last_inplace::<LAZY>(&mut a[0], &mut b[0], psi, n_inv);
self.dif_last_inplace::<LAZY>(&mut a[1], &mut b[1], psi, n_inv);
self.dif_last_inplace::<LAZY>(&mut a[2], &mut b[2], psi, n_inv);
self.dif_last_inplace::<LAZY>(&mut a[3], &mut b[3], psi, n_inv);
self.dif_last_inplace::<LAZY>(&mut a[4], &mut b[4], psi, n_inv);
self.dif_last_inplace::<LAZY>(&mut a[5], &mut b[5], psi, n_inv);
self.dif_last_inplace::<LAZY>(&mut a[6], &mut b[6], psi, n_inv);
self.dif_last_inplace::<LAZY>(&mut a[7], &mut b[7], psi, n_inv);
});
},
);
} else if t >= 16{
izip!(a.chunks_exact_mut(t), &self.psi_backward_rev[m..]).for_each(
|(a, psi)| {
let (a, b) = a.split_at_mut(size);
izip!(a.chunks_exact_mut(8), b.chunks_exact_mut(8)).for_each(|(a, b)| {
self.dif_inplace::<true>(&mut a[0], &mut b[0], *psi);
self.dif_inplace::<true>(&mut a[1], &mut b[1], *psi);
self.dif_inplace::<true>(&mut a[2], &mut b[2], *psi);
self.dif_inplace::<true>(&mut a[3], &mut b[3], *psi);
self.dif_inplace::<true>(&mut a[4], &mut b[4], *psi);
self.dif_inplace::<true>(&mut a[5], &mut b[5], *psi);
self.dif_inplace::<true>(&mut a[6], &mut b[6], *psi);
self.dif_inplace::<true>(&mut a[7], &mut b[7], *psi);
});
},
);
} else {
izip!(a.chunks_exact_mut(2 * size), &self.psi_backward_rev[m..]).for_each(
|(a, psi)| {
let (a, b) = a.split_at_mut(size);
izip!(a, b).for_each(|(a, b)| self.dif_inplace::<true>(a, b, *psi));
},
);
}
}
}
#[inline(always)]
fn dif_inplace<const LAZY: bool>(&self, a: &mut u64, b: &mut u64, t: Barrett<u64>) {
debug_assert!(*a < self.two_q, "a:{} q:{}", a, self.four_q);
debug_assert!(*b < self.two_q, "b:{} q:{}", b, self.four_q);
let d: u64 = self.prime.barrett.mul_external_lazy(t, *a + self.two_q - *b);
*a = a.wrapping_add(*b);
a.reduce_once_assign(self.two_q);
*b = d;
if !LAZY {
a.reduce_once_assign(self.q);
b.reduce_once_assign(self.q);
}
}
fn dif_last_inplace<const LAZY:bool>(&self, a: &mut u64, b: &mut u64, psi: Barrett<u64>, n_inv: Barrett<u64>){
debug_assert!(*a < self.two_q);
debug_assert!(*b < self.two_q);
if LAZY{
let d: u64 = self.prime.barrett.mul_external_lazy(psi, *a + self.two_q - *b);
*a = self.prime.barrett.mul_external_lazy(n_inv, *a + *b);
*b = d;
}else{
let d: u64 = self.prime.barrett.mul_external(psi, *a + self.two_q - *b);
*a = self.prime.barrett.mul_external(n_inv, *a + *b);
*b = d;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_ntt() {
let q_base: u64 = 0x800000000004001;
let q_power: usize = 1;
let prime_instance: Prime<u64> = Prime::<u64>::new(q_base, q_power);
let n: u64 = 32;
let two_nth_root: u64 = n<<1;
let ntt_table: Table<u64> = Table::<u64>::new(prime_instance, two_nth_root);
let mut a: Vec<u64> = vec![0; n as usize];
for i in 0..a.len(){
a[i] = i as u64;
}
let b: Vec<u64> = a.clone();
ntt_table.forward_inplace::<false>(&mut a);
ntt_table.backward_inplace::<false>(&mut a);
assert!(a == b);
}
}

120
math/src/lib.rs Normal file
View File

@@ -0,0 +1,120 @@
#![feature(bigint_helper_methods)]
#![feature(test)]
pub mod modulus;
pub mod dft;
pub mod ring;
pub mod poly;
pub const CHUNK: usize= 8;
pub mod macros{
#[macro_export]
macro_rules! apply_unary {
($self:expr, $f:expr, $a:expr, $CHUNK:expr) => {
match CHUNK{
8 => {
$a.chunks_exact_mut(8).for_each(|a| {
$f(&$self, &mut a[0]);
$f(&$self, &mut a[1]);
$f(&$self, &mut a[2]);
$f(&$self, &mut a[3]);
$f(&$self, &mut a[4]);
$f(&$self, &mut a[5]);
$f(&$self, &mut a[6]);
$f(&$self, &mut a[7]);
});
let n: usize = $a.len();
let m = n - (n&(CHUNK-1));
$a[m..].iter_mut().for_each(|a| {
$f(&$self, a);
});
},
_=>{
$a.iter_mut().for_each(|a| {
$f(&$self, a);
});
}
}
};
}
#[macro_export]
macro_rules! apply_binary {
($self:expr, $f:expr, $a:expr, $b:expr, $CHUNK:expr) => {
let n: usize = $a.len();
debug_assert!($b.len() == n, "invalid argument b: b.len() = {} != a.len() = {}", $b.len(), n);
debug_assert!(CHUNK&(CHUNK-1) == 0, "invalid CHUNK const: not a power of two");
match CHUNK{
8 => {
izip!($a.chunks_exact(8), $b.chunks_exact_mut(8)).for_each(|(a, b)| {
$f(&$self, &a[0], &mut b[0]);
$f(&$self, &a[1], &mut b[1]);
$f(&$self, &a[2], &mut b[2]);
$f(&$self, &a[3], &mut b[3]);
$f(&$self, &a[4], &mut b[4]);
$f(&$self, &a[5], &mut b[5]);
$f(&$self, &a[6], &mut b[6]);
$f(&$self, &a[7], &mut b[7]);
});
let m = n - (n&(CHUNK-1));
izip!($a[m..].iter(), $b[m..].iter_mut()).for_each(|(a, b)| {
$f(&$self, a, b);
});
},
_=>{
izip!($a.iter(), $b.iter_mut()).for_each(|(a, b)| {
$f(&$self, a, b);
});
}
}
};
}
#[macro_export]
macro_rules! apply_ternary {
($self:expr, $f:expr, $a:expr, $b:expr, $c:expr, $CHUNK:expr) => {
let n: usize = $a.len();
debug_assert!($b.len() == n, "invalid argument b: b.len() = {} != a.len() = {}", $b.len(), n);
debug_assert!(CHUNK&(CHUNK-1) == 0, "invalid CHUNK const: not a power of two");
match CHUNK{
8 => {
izip!($a.chunks_exact(8), $b.chunks_exact(8), $c.chunks_exact_mut(8)).for_each(|(a, b, c)| {
$f(&$self, &a[0], &b[0], &mut c[0]);
$f(&$self, &a[1], &b[1], &mut c[1]);
$f(&$self, &a[2], &b[2], &mut c[2]);
$f(&$self, &a[3], &b[3], &mut c[3]);
$f(&$self, &a[4], &b[4], &mut c[4]);
$f(&$self, &a[5], &b[5], &mut c[5]);
$f(&$self, &a[6], &b[6], &mut c[6]);
$f(&$self, &a[7], &b[7], &mut c[7]);
});
let m = n - (n&7);
izip!($a[m..].iter(), $b[m..].iter(), $c[m..].iter_mut()).for_each(|(a, b, c)| {
$f(&$self, a, b, c);
});
},
_=>{
izip!($a.iter(), $b.iter(), $c.iter_mut()).for_each(|(a, b, c)| {
$f(&$self, a, b, c);
});
}
}
};
}
}

125
math/src/modulus.rs Normal file
View File

@@ -0,0 +1,125 @@
pub mod prime;
pub mod barrett;
pub mod montgomery;
pub mod impl_u64;
pub type REDUCEMOD = u8;
pub const NONE: REDUCEMOD = 0;
pub const ONCE: REDUCEMOD = 1;
pub const TWICE: REDUCEMOD = 2;
pub const FOURTIMES: REDUCEMOD = 3;
pub const BARRETT: REDUCEMOD = 4;
pub const BARRETTLAZY: REDUCEMOD = 5;
pub trait WordOps<O>{
fn log2(self) -> O;
fn reverse_bits_msb(self, n:u32) -> O;
}
impl WordOps<u64> for u64{
#[inline(always)]
fn log2(self) -> u64{
(u64::BITS - (self-1).leading_zeros()) as _
}
#[inline(always)]
fn reverse_bits_msb(self, n: u32) -> u64{
self.reverse_bits() >> (usize::BITS - n)
}
}
impl WordOps<usize> for usize{
#[inline(always)]
fn log2(self) -> usize{
(usize::BITS - (self-1).leading_zeros()) as _
}
#[inline(always)]
fn reverse_bits_msb(self, n: u32) -> usize{
self.reverse_bits() >> (usize::BITS - n)
}
}
pub trait ReduceOnce<O>{
/// Assigns self-q to self if self >= q in constant time.
/// User must ensure that 2q fits in O.
fn reduce_once_constant_time_assign(&mut self, q: O);
/// Returns self-q if self >= q else self in constant time.
/// /// User must ensure that 2q fits in O.
fn reduce_once_constant_time(&self, q:O) -> O;
/// Assigns self-q to self if self >= q.
/// /// User must ensure that 2q fits in O.
fn reduce_once_assign(&mut self, q: O);
/// Returns self-q if self >= q else self.
/// /// User must ensure that 2q fits in O.
fn reduce_once(&self, q:O) -> O;
}
pub trait WordOperations<O>{
// Applies a parameterized modular reduction.
fn word_reduce_assign<const REDUCE:REDUCEMOD>(&self, x: &mut O);
// Assigns a + b to c.
fn word_add_binary_assign<const REDUCE:REDUCEMOD>(&self, a: &O, b:&O, c: &mut O);
// Assigns a + b to b.
fn word_add_unary_assign<const REDUCE:REDUCEMOD>(&self, a: &O, b: &mut O);
// Assigns a - b to c.
fn word_sub_binary_assign<const REDUCE:REDUCEMOD>(&self, a: &O, b:&O, c: &mut O);
// Assigns b - a to b.
fn word_sub_unary_assign<const REDUCE:REDUCEMOD>(&self, a: &O, b: &mut O);
// Assigns -a to a.
fn word_neg_unary_assign<const REDUCE:REDUCEMOD>(&self, a:&mut O);
// Assigns -a to b.
fn word_neg_binary_assign<const REDUCE:REDUCEMOD>(&self, a: &O, b:&mut O);
// Assigns a * 2^64 to b.
fn word_prepare_montgomery_assign<const REDUCE:REDUCEMOD>(&self, a: &O, b: &mut montgomery::Montgomery<O>);
// Assigns a * b to c.
fn word_mul_montgomery_external_binary_assign<const REDUCE:REDUCEMOD>(&self, a:&montgomery::Montgomery<O>, b:&O, c: &mut O);
// Assigns a * b to b.
fn word_mul_montgomery_external_unary_assign<const REDUCE:REDUCEMOD>(&self, a:&montgomery::Montgomery<O>, b:&mut O);
}
pub trait VecOperations<O>{
// Applies a parameterized modular reduction.
fn vec_reduce_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, x: &mut [O]);
// Assigns a[i] + b[i] to c[i]
fn vec_add_binary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[O], b:&[O], c: &mut [O]);
// Assigns a[i] + b[i] to b[i]
fn vec_add_unary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[O], b: &mut [O]);
// Assigns a[i] - b[i] to c[i]
fn vec_sub_binary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[O], b:&[O], c: &mut [O]);
// Assigns a[i] - b[i] to b[i]
fn vec_sub_unary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[O], b: &mut [O]);
// Assigns -a[i] to a[i].
fn vec_neg_unary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &mut [O]);
// Assigns -a[i] to a[i].
fn vec_neg_binary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[O], b: &mut [O]);
// Assigns a * 2^64 to b.
fn vec_prepare_montgomery_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[O], b: &mut [montgomery::Montgomery<O>]);
// Assigns a[i] * b[i] to c[i].
fn vec_mul_montgomery_external_binary_assign<const CHUNK:usize,const REDUCE:REDUCEMOD>(&self, a:&[montgomery::Montgomery<O>], b:&[O], c: &mut [O]);
// Assigns a[i] * b[i] to b[i].
fn vec_mul_montgomery_external_unary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a:&[montgomery::Montgomery<O>], b:&mut [O]);
}

37
math/src/modulus/barrett.rs Normal file
View File

@@ -0,0 +1,37 @@
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Barrett<O>(pub O, pub O);
impl<O> Barrett<O> {
#[inline(always)]
pub fn value(&self) -> &O {
&self.0
}
#[inline(always)]
pub fn quotient(&self) -> &O {
&self.1
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BarrettPrecomp<O>{
pub q: O,
pub lo:O,
pub hi:O,
pub one: Barrett<O>,
}
impl<O> BarrettPrecomp<O>{
#[inline(always)]
pub fn value_hi(&self) -> &O{
&self.hi
}
#[inline(always)]
pub fn value_lo(&self) -> &O{
&self.lo
}
}

84
math/src/modulus/impl_u64/barrett.rs Normal file
View File

@@ -0,0 +1,84 @@
use crate::modulus::barrett::{Barrett, BarrettPrecomp};
use crate::modulus::ReduceOnce;
use num_bigint::BigUint;
use num_traits::cast::ToPrimitive;
impl BarrettPrecomp<u64>{
pub fn new(q: u64) -> BarrettPrecomp<u64> {
let big_r: BigUint = (BigUint::from(1 as usize)<<((u64::BITS<<1) as usize)) / BigUint::from(q);
let lo: u64 = (&big_r & BigUint::from(u64::MAX)).to_u64().unwrap();
let hi: u64 = (big_r >> u64::BITS).to_u64().unwrap();
let mut precomp: BarrettPrecomp<u64> = Self{q, lo, hi, one:Barrett(0,0)};
precomp.one = precomp.prepare(1);
precomp
}
#[inline(always)]
pub fn one(&self) -> Barrett<u64> {
self.one
}
#[inline(always)]
pub fn prepare(&self, v: u64) -> Barrett<u64> {
debug_assert!(v < self.q);
let quotient: u64 = (((v as u128) << 64) / self.q as u128) as _;
Barrett(v, quotient)
}
/// Returns lhs mod q.
#[inline(always)]
pub fn reduce(&self, lhs: u64) -> u64{
let mut r: u64 = self.reduce_lazy(lhs);
r.reduce_once_assign(self.q);
r
}
/// Returns lhs mod q in range [0, 2q-1].
#[inline(always)]
pub fn reduce_lazy(&self, lhs: u64) -> u64{
let (_, mhi) = lhs.widening_mul(self.hi);
lhs - mhi.wrapping_mul(self.q)
}
/// Assigns lhs mod q to lhs.
#[inline(always)]
pub fn reduce_assign(&self, lhs: &mut u64){
self.reduce_lazy_assign(lhs);
lhs.reduce_once_assign(self.q);
}
/// Assigns lhs mod q in range [0, 2q-1] to lhs.
#[inline(always)]
pub fn reduce_lazy_assign(&self, lhs: &mut u64){
let (_, mhi) = lhs.widening_mul(self.hi);
*lhs = *lhs - mhi.wrapping_mul(self.q)
}
#[inline(always)]
pub fn mul_external(&self, lhs: Barrett<u64>, rhs: u64) -> u64 {
let mut r: u64 = self.mul_external_lazy(lhs, rhs);
r.reduce_once_assign(self.q);
r
}
#[inline(always)]
pub fn mul_external_assign(&self, lhs: Barrett<u64>, rhs: &mut u64){
self.mul_external_lazy_assign(lhs, rhs);
rhs.reduce_once_assign(self.q);
}
#[inline(always)]
pub fn mul_external_lazy(&self, lhs: Barrett<u64>, rhs: u64) -> u64 {
let mut r: u64 = rhs;
self.mul_external_lazy_assign(lhs, &mut r);
r
}
#[inline(always)]
pub fn mul_external_lazy_assign(&self, lhs: Barrett<u64>, rhs: &mut u64){
let t: u64 = ((*lhs.quotient() as u128 * *rhs as u128) >> 64) as _;
*rhs = (rhs.wrapping_mul(*lhs.value())).wrapping_sub(self.q.wrapping_mul(t));
}
}

143
math/src/modulus/impl_u64/generation.rs Normal file
View File

@@ -0,0 +1,143 @@
use crate::modulus::prime;
use prime::Prime;
use primality_test::is_prime;
impl NTTFriendlyPrimesGenerator<u64> {
pub fn new(bit_size: u64, nth_root: u64) -> Self{
let mut check_next_prime: bool = true;
let mut check_prev_prime: bool = true;
let next_prime = (1<<bit_size) + 1;
let mut prev_prime = next_prime;
if next_prime > nth_root.wrapping_neg(){
check_next_prime = false;
}
if prev_prime < nth_root{
check_prev_prime = false
}
prev_prime -= nth_root;
Self{
size: bit_size as f64,
check_next_prime,
check_prev_prime,
nth_root,
next_prime,
prev_prime,
}
}
pub fn next_upstream_primes(&mut self, k: usize) -> Vec<Prime>{
let mut primes: Vec<Prime> = Vec::with_capacity(k);
for i in 0..k{
primes.push(self.next_upstream_prime())
}
primes
}
pub fn next_downstream_primes(&mut self, k: usize) -> Vec<Prime>{
let mut primes: Vec<Prime> = Vec::with_capacity(k);
for i in 0..k{
primes.push(self.next_downstream_prime())
}
primes
}
pub fn next_alternating_primes(&mut self, k: usize) -> Vec<Prime>{
let mut primes: Vec<Prime> = Vec::with_capacity(k);
for i in 0..k{
primes.push(self.next_alternating_prime())
}
primes
}
pub fn next_upstream_prime(&mut self) -> Prime{
loop {
if self.check_next_prime{
if (self.next_prime as f64).log2() - self.size >= 0.5 || self.next_prime > 0xffff_ffff_ffff_ffff-self.nth_root{
self.check_next_prime = false;
panic!("prime list for upstream primes is exhausted (overlap with next bit-size or prime > 2^64)");
}
}else{
if is_prime(self.next_prime) {
let prime = Prime::new_unchecked(self.next_prime);
self.next_prime += self.nth_root;
return prime
}
self.next_prime += self.nth_root;
}
}
}
pub fn next_downstream_prime(&mut self) -> Prime{
loop {
if self.size - (self.prev_prime as f64).log2() >= 0.5 || self.prev_prime < self.nth_root{
self.check_next_prime = false;
panic!("prime list for downstream primes is exhausted (overlap with previous bit-size or prime < nth_root)")
}else{
if is_prime(self.prev_prime){
let prime = Prime::new_unchecked(self.next_prime);
self.prev_prime -= self.nth_root;
return prime
}
self.prev_prime -= self.nth_root;
}
}
}
pub fn next_alternating_prime(&mut self) -> Prime{
loop {
if !(self.check_next_prime || self.check_prev_prime){
panic!("prime list for upstream and downstream prime is exhausted for the (overlap with previous/next bit-size or NthRoot > prime > 2^64)")
}
if self.check_next_prime{
if (self.next_prime as f64).log2() - self.size >= 0.5 || self.next_prime > 0xffff_ffff_ffff_ffff-self.nth_root{
self.check_next_prime = false;
}else{
if is_prime(self.next_prime){
let prime = Prime::new_unchecked(self.next_prime);
self.next_prime += self.nth_root;
return prime
}
self.next_prime += self.nth_root;
}
}
if self.check_prev_prime {
if self.size - (self.prev_prime as f64).log2() >= 0.5 || self.prev_prime < self.nth_root{
self.check_prev_prime = false;
}else{
if is_prime(self.prev_prime){
let prime = Prime::new_unchecked(self.prev_prime);
self.prev_prime -= self.nth_root;
return prime
}
self.prev_prime -= self.nth_root;
}
}
}
}
}
#[cfg(test)]
mod test {
use crate::modulus::prime_generator;
#[test]
fn prime_generation() {
let nth_root: u64 = 1<<16 ;
let mut g: prime_generator::NTTFriendlyPrimesGenerator = prime_generator::NTTFriendlyPrimesGenerator::new(30, nth_root);
let primes = g.next_alternating_primes(10);
println!("{:?}", primes);
for prime in primes.iter(){
assert!(prime.q() % nth_root == 1);
}
}
}

32
math/src/modulus/impl_u64/mod.rs Normal file
View File

@@ -0,0 +1,32 @@
pub mod prime;
pub mod barrett;
pub mod montgomery;
pub mod operations;
use crate::modulus::ReduceOnce;
impl ReduceOnce<u64> for u64{
#[inline(always)]
fn reduce_once_constant_time_assign(&mut self, q: u64){
debug_assert!(q < 0x8000000000000000, "2q >= 2^64");
*self -= (q.wrapping_sub(*self)>>63)*q;
}
#[inline(always)]
fn reduce_once_constant_time(&self, q:u64) -> u64{
debug_assert!(q < 0x8000000000000000, "2q >= 2^64");
self - (q.wrapping_sub(*self)>>63)*q
}
#[inline(always)]
fn reduce_once_assign(&mut self, q: u64){
debug_assert!(q < 0x8000000000000000, "2q >= 2^64");
*self = *self.min(&mut self.wrapping_sub(q))
}
#[inline(always)]
fn reduce_once(&self, q:u64) -> u64{
debug_assert!(q < 0x8000000000000000, "2q >= 2^64");
*self.min(&mut self.wrapping_sub(q))
}
}

206
math/src/modulus/impl_u64/montgomery.rs Normal file
View File

@@ -0,0 +1,206 @@
use crate::modulus::ReduceOnce;
use crate::modulus::montgomery::{MontgomeryPrecomp, Montgomery};
use crate::modulus::barrett::BarrettPrecomp;
use crate::modulus::{REDUCEMOD, NONE, ONCE, TWICE, FOURTIMES, BARRETT, BARRETTLAZY};
extern crate test;
/// MontgomeryPrecomp is a set of methods implemented for MontgomeryPrecomp<u64>
/// enabling Montgomery arithmetic over u64 values.
impl MontgomeryPrecomp<u64>{
/// Returns an new instance of MontgomeryPrecomp<u64>.
/// This method will fail if gcd(q, 2^64) != 1.
#[inline(always)]
pub fn new(q: u64) -> MontgomeryPrecomp<u64>{
assert!(q & 1 != 0, "Invalid argument: gcd(q={}, radix=2^64) != 1", q);
let mut q_inv: u64 = 1;
let mut q_pow = q;
for _i in 0..63{
q_inv = q_inv.wrapping_mul(q_pow);
q_pow = q_pow.wrapping_mul(q_pow);
}
let mut precomp = Self{
q: q,
two_q: q<<1,
four_q: q<<2,
barrett: BarrettPrecomp::new(q),
q_inv: q_inv,
one: 0,
minus_one:0,
};
precomp.one = precomp.prepare::<ONCE>(1);
precomp.minus_one = q-precomp.one;
precomp
}
/// Returns 2^64 mod q as a Montgomery<u64>.
#[inline(always)]
pub fn one(&self) -> Montgomery<u64>{
self.one
}
/// Returns (q-1) * 2^64 mod q as a Montgomery<u64>.
#[inline(always)]
pub fn minus_one(&self) -> Montgomery<u64>{
self.minus_one
}
/// Applies a modular reduction on x based on REDUCE:
/// - LAZY: no modular reduction.
/// - ONCE: subtracts q if x >= q.
/// - FULL: maps x to x mod q using Barrett reduction.
#[inline(always)]
pub fn reduce<const REDUCE:REDUCEMOD>(&self, x: u64) -> u64{
let mut r: u64 = x;
self.reduce_assign::<REDUCE>(&mut r);
r
}
/// Applies a modular reduction on x based on REDUCE:
/// - LAZY: no modular reduction.
/// - ONCE: subtracts q if x >= q.
/// - FULL: maps x to x mod q using Barrett reduction.
#[inline(always)]
pub fn reduce_assign<const REDUCE:REDUCEMOD>(&self, x: &mut u64){
match REDUCE {
NONE =>{},
ONCE =>{x.reduce_once_assign(self.q)},
TWICE=>{x.reduce_once_assign(self.two_q)},
FOURTIMES =>{x.reduce_once_assign(self.four_q)},
BARRETT =>{self.barrett.reduce_assign(x)},
BARRETTLAZY =>{self.barrett.reduce_lazy_assign(x)},
_ => unreachable!("invalid REDUCE argument")
}
}
/// Returns lhs * 2^64 mod q as a Montgomery<u64>.
#[inline(always)]
pub fn prepare<const REDUCE:REDUCEMOD>(&self, lhs: u64) -> Montgomery<u64>{
let mut rhs: u64 = 0;
self.prepare_assign::<REDUCE>(lhs, &mut rhs);
rhs
}
/// Assigns lhs * 2^64 mod q to rhs.
#[inline(always)]
pub fn prepare_assign<const REDUCE:REDUCEMOD>(&self, lhs: u64, rhs: &mut Montgomery<u64>){
let (_, mhi) = lhs.widening_mul(*self.barrett.value_lo());
*rhs = (lhs.wrapping_mul(*self.barrett.value_hi()).wrapping_add(mhi)).wrapping_mul(self.q).wrapping_neg();
self.reduce_assign::<REDUCE>(rhs);
}
/// Returns lhs * (2^64)^-1 mod q as a u64.
#[inline(always)]
pub fn unprepare<const REDUCE:REDUCEMOD>(&self, lhs: Montgomery<u64>) -> u64{
let mut rhs = 0u64;
self.unprepare_assign::<REDUCE>(lhs, &mut rhs);
rhs
}
/// Assigns lhs * (2^64)^-1 mod q to rhs.
#[inline(always)]
pub fn unprepare_assign<const REDUCE:REDUCEMOD>(&self, lhs: Montgomery<u64>, rhs: &mut u64){
let (_, r) = self.q.widening_mul(lhs.wrapping_mul(self.q_inv));
*rhs = self.reduce::<REDUCE>(self.q.wrapping_sub(r));
}
/// Returns lhs * rhs * (2^{64})^-1 mod q.
#[inline(always)]
pub fn mul_external<const REDUCE:REDUCEMOD>(&self, lhs: Montgomery<u64>, rhs: u64) -> u64{
let mut r: u64 = rhs;
self.mul_external_assign::<REDUCE>(lhs, &mut r);
r
}
/// Assigns lhs * rhs * (2^{64})^-1 mod q to rhs.
#[inline(always)]
pub fn mul_external_assign<const REDUCE:REDUCEMOD>(&self, lhs: Montgomery<u64>, rhs: &mut u64){
let (mlo, mhi) = lhs.widening_mul(*rhs);
let (_, hhi) = self.q.widening_mul(mlo.wrapping_mul(self.q_inv));
*rhs = self.reduce::<REDUCE>(mhi.wrapping_sub(hhi).wrapping_add(self.q));
}
/// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1].
#[inline(always)]
pub fn mul_internal<const REDUCE:REDUCEMOD>(&self, lhs: Montgomery<u64>, rhs: Montgomery<u64>) -> Montgomery<u64>{
self.mul_external::<REDUCE>(lhs, rhs)
}
/// Assigns lhs * rhs * (2^{64})^-1 mod q to rhs.
#[inline(always)]
pub fn mul_internal_assign<const REDUCE:REDUCEMOD>(&self, lhs: Montgomery<u64>, rhs: &mut Montgomery<u64>){
self.mul_external_assign::<REDUCE>(lhs, rhs);
}
#[inline(always)]
pub fn add_internal(&self, lhs: Montgomery<u64>, rhs: Montgomery<u64>) -> Montgomery<u64>{
self.barrett.reduce(rhs + lhs)
}
/// Assigns lhs + rhs to rhs.
#[inline(always)]
pub fn add_internal_lazy_assign(&self, lhs: Montgomery<u64>, rhs: &mut Montgomery<u64>){
*rhs += lhs
}
/// Assigns lhs + rhs - q if (lhs + rhs) >= q to rhs.
#[inline(always)]
pub fn add_internal_reduce_once_assign<const LAZY:bool>(&self, lhs: Montgomery<u64>, rhs: &mut Montgomery<u64>){
self.add_internal_lazy_assign(lhs, rhs);
rhs.reduce_once_assign(self.q);
}
/// Returns lhs mod q in range [0, 2q-1].
#[inline(always)]
pub fn reduce_lazy_assign(&self, lhs: &mut u64){
self.barrett.reduce_lazy_assign(lhs)
}
/// Returns (x^exponent) * 2^64 mod q.
#[inline(always)]
pub fn pow(&self, x: Montgomery<u64>, exponent:u64) -> Montgomery<u64>{
let mut y: Montgomery<u64> = self.one();
let mut x_mut: Montgomery<u64> = x;
let mut i: u64 = exponent;
while i > 0{
if i & 1 == 1{
self.mul_internal_assign::<ONCE>(x_mut, &mut y);
}
self.mul_internal_assign::<ONCE>(x_mut, &mut x_mut);
i >>= 1;
}
y.reduce_once_assign(self.q);
y
}
}
#[cfg(test)]
mod tests {
use crate::modulus::montgomery;
use super::*;
use test::Bencher;
#[test]
fn test_mul_external() {
let q: u64 = 0x1fffffffffe00001;
let m_precomp = montgomery::MontgomeryPrecomp::new(q);
let x: u64 = 0x5f876e514845cc8b;
let y: u64 = 0xad726f98f24a761a;
let y_mont = m_precomp.prepare::<ONCE>(y);
assert!(m_precomp.mul_external::<ONCE>(y_mont, x) == (x as u128 * y as u128 % q as u128) as u64);
}
#[bench]
fn bench_mul_external(b: &mut Bencher){
let q: u64 = 0x1fffffffffe00001;
let m_precomp = montgomery::MontgomeryPrecomp::new(q);
let mut x: u64 = 0x5f876e514845cc8b;
let y: u64 = 0xad726f98f24a761a;
let y_mont = m_precomp.prepare::<ONCE>(y);
b.iter(|| m_precomp.mul_external_assign::<ONCE>(y_mont, &mut x));
}
}

132
math/src/modulus/impl_u64/operations.rs Normal file
View File

@@ -0,0 +1,132 @@
use crate::modulus::{WordOperations, VecOperations};
use crate::modulus::prime::Prime;
use crate::modulus::ReduceOnce;
use crate::modulus::montgomery::Montgomery;
use crate::modulus::REDUCEMOD;
use crate::{apply_unary, apply_binary, apply_ternary};
use itertools::izip;
impl WordOperations<u64> for Prime<u64>{
/// Applies a modular reduction on x based on REDUCE:
/// - LAZY: no modular reduction.
/// - ONCE: subtracts q if x >= q.
/// - TWO: subtracts 2q if x >= 2q.
/// - FOUR: subtracts 4q if x >= 4q.
/// - BARRETT: maps x to x mod q using Barrett reduction.
/// - BARRETTLAZY: maps x to x mod q using Barrett reduction with values in [0, 2q-1].
#[inline(always)]
fn word_reduce_assign<const REDUCE:REDUCEMOD>(&self, x: &mut u64){
self.montgomery.reduce_assign::<REDUCE>(x);
}
#[inline(always)]
fn word_add_binary_assign<const REDUCE:REDUCEMOD>(&self, a: &u64, b: &u64, c: &mut u64){
*c = a.wrapping_add(*b);
self.word_reduce_assign::<REDUCE>(c);
}
#[inline(always)]
fn word_add_unary_assign<const REDUCE:REDUCEMOD>(&self, a: &u64, b: &mut u64){
*b = a.wrapping_add(*b);
self.word_reduce_assign::<REDUCE>(b);
}
#[inline(always)]
fn word_sub_binary_assign<const REDUCE:REDUCEMOD>(&self, a: &u64, b: &u64, c: &mut u64){
*c = a.wrapping_add(self.q.wrapping_sub(*b)).reduce_once(self.q);
}
#[inline(always)]
fn word_sub_unary_assign<const REDUCE:REDUCEMOD>(&self, a: &u64, b: &mut u64){
*b = a.wrapping_add(self.q.wrapping_sub(*b)).reduce_once(self.q);
}
#[inline(always)]
fn word_neg_unary_assign<const REDUCE:REDUCEMOD>(&self, a: &mut u64){
*a = self.q.wrapping_sub(*a);
self.word_reduce_assign::<REDUCE>(a)
}
#[inline(always)]
fn word_neg_binary_assign<const REDUCE:REDUCEMOD>(&self, a: &u64, b: &mut u64){
*b = self.q.wrapping_sub(*a);
self.word_reduce_assign::<REDUCE>(b)
}
#[inline(always)]
fn word_prepare_montgomery_assign<const REDUCE:REDUCEMOD>(&self, a: &u64, b: &mut Montgomery<u64>){
self.montgomery.prepare_assign::<REDUCE>(*a, b);
}
#[inline(always)]
fn word_mul_montgomery_external_binary_assign<const REDUCE:REDUCEMOD>(&self, a: &Montgomery<u64>, b:&u64, c: &mut u64){
*c = self.montgomery.mul_external::<REDUCE>(*a, *b);
}
#[inline(always)]
fn word_mul_montgomery_external_unary_assign<const REDUCE:REDUCEMOD>(&self, lhs:&Montgomery<u64>, rhs:&mut u64){
self.montgomery.mul_external_assign::<REDUCE>(*lhs, rhs);
}
}
impl VecOperations<u64> for Prime<u64>{
/// Applies a modular reduction on x based on REDUCE:
/// - LAZY: no modular reduction.
/// - ONCE: subtracts q if x >= q.
/// - TWO: subtracts 2q if x >= 2q.
/// - FOUR: subtracts 4q if x >= 4q.
/// - BARRETT: maps x to x mod q using Barrett reduction.
/// - BARRETTLAZY: maps x to x mod q using Barrett reduction with values in [0, 2q-1].
#[inline(always)]
fn vec_reduce_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, x: &mut [u64]){
apply_unary!(self, Self::word_reduce_assign::<REDUCE>, x, CHUNK);
}
#[inline(always)]
fn vec_add_binary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[u64], b:&[u64], c:&mut [u64]){
apply_ternary!(self, Self::word_add_binary_assign::<REDUCE>, a, b, c, CHUNK);
}
#[inline(always)]
fn vec_add_unary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[u64], b:&mut [u64]){
apply_binary!(self, Self::word_add_unary_assign::<REDUCE>, a, b, CHUNK);
}
#[inline(always)]
fn vec_sub_binary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[u64], b:&[u64], c:&mut [u64]){
apply_ternary!(self, Self::word_sub_binary_assign::<REDUCE>, a, b, c, CHUNK);
}
#[inline(always)]
fn vec_sub_unary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[u64], b:&mut [u64]){
apply_binary!(self, Self::word_sub_unary_assign::<REDUCE>, a, b, CHUNK);
}
#[inline(always)]
fn vec_neg_unary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &mut [u64]){
apply_unary!(self, Self::word_neg_unary_assign::<REDUCE>, a, CHUNK);
}
#[inline(always)]
fn vec_neg_binary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[u64], b: &mut [u64]){
apply_binary!(self, Self::word_neg_binary_assign::<REDUCE>, a, b, CHUNK);
}
#[inline(always)]
fn vec_prepare_montgomery_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a: &[u64], b: &mut [Montgomery<u64>]){
apply_binary!(self, Self::word_prepare_montgomery_assign::<REDUCE>, a, b, CHUNK);
}
#[inline(always)]
fn vec_mul_montgomery_external_binary_assign<const CHUNK:usize,const REDUCE:REDUCEMOD>(&self, a:& [Montgomery<u64>], b:&[u64], c: &mut [u64]){
apply_ternary!(self, Self::word_mul_montgomery_external_binary_assign::<REDUCE>, a, b, c, CHUNK);
}
#[inline(always)]
fn vec_mul_montgomery_external_unary_assign<const CHUNK:usize, const REDUCE:REDUCEMOD>(&self, a:& [Montgomery<u64>], b:&mut [u64]){
apply_binary!(self, Self::word_mul_montgomery_external_unary_assign::<REDUCE>, a, b, CHUNK);
}
}

218
math/src/modulus/impl_u64/prime.rs Normal file
View File

@@ -0,0 +1,218 @@
use crate::modulus::prime::Prime;
use crate::modulus::montgomery::{Montgomery, MontgomeryPrecomp};
use crate::modulus::barrett::BarrettPrecomp;
use crate::modulus::ONCE;
use primality_test::is_prime;
use prime_factorization::Factorization;
impl Prime<u64>{
/// Returns a new instance of Prime<u64>.
/// Panics if q_base is not a prime > 2 and
/// if q_base^q_power would overflow u64.
pub fn new(q_base: u64, q_power: usize) -> Self{
assert!(is_prime(q_base) && q_base > 2);
Self::new_unchecked(q_base, q_power)
}
/// Returns a new instance of Prime<u64>.
/// Does not check if q_base is a prime > 2.
/// Panics if q_base^q_power would overflow u64.
pub fn new_unchecked(q_base: u64, q_power: usize) -> Self {
let mut q = q_base;
for _i in 1..q_power{
q *= q_base
}
assert!(q.next_power_of_two().ilog2() <= 61);
let mut phi = q_base-1;
for _i in 1..q_power{
phi *= q_base
}
let mut prime: Prime<u64> = Self {
q:q,
two_q:q<<1,
four_q:q<<2,
q_base:q_base,
q_power:q_power,
factors: Vec::new(),
montgomery:MontgomeryPrecomp::new(q),
barrett:BarrettPrecomp::new(q),
phi:phi,
};
prime.check_factors();
prime
}
pub fn q(&self) -> u64{
self.q
}
pub fn q_base(&self) -> u64{
self.q_base
}
pub fn q_power(&self) -> usize{
self.q_power
}
/// Returns x^exponen mod q.
#[inline(always)]
pub fn pow(&self, x: u64, exponent: u64) -> u64{
let mut y_mont: Montgomery<u64> = self.montgomery.one();
let mut x_mont: Montgomery<u64> = self.montgomery.prepare::<ONCE>(x);
let mut i: u64 = exponent;
while i > 0{
if i & 1 == 1{
self.montgomery.mul_internal_assign::<ONCE>(x_mont, &mut y_mont);
}
self.montgomery.mul_internal_assign::<ONCE>(x_mont, &mut x_mont);
i >>= 1;
}
self.montgomery.unprepare::<ONCE>(y_mont)
}
/// Returns x^-1 mod q.
/// User must ensure that x is not divisible by q_base.
#[inline(always)]
pub fn inv(&self, x: u64) -> u64{
self.pow(x, self.phi-1)
}
}
impl Prime<u64>{
/// Returns the smallest nth primitive root of q_base.
pub fn primitive_root(&self) -> u64{
let mut candidate: u64 = 1u64;
let mut not_found: bool = true;
while not_found{
candidate += 1;
for &factor in &self.factors{
if pow(candidate, (self.q_base-1)/factor, self.q_base) == 1{
not_found = true;
break
}
not_found = false;
}
}
if not_found{
panic!("failed to find a primitive root for q_base={}", self.q_base)
}
candidate
}
/// Returns an nth primitive root of q = q_base^q_power in Montgomery.
pub fn primitive_nth_root(&self, nth_root:u64) -> u64{
assert!(self.q & (nth_root-1) == 1, "invalid prime: q = {} % nth_root = {} = {} != 1", self.q, nth_root, self.q & (nth_root-1));
let psi: u64 = self.primitive_root();
// nth primitive root mod q_base: psi_nth^(prime.q_base-1)/nth_root mod q_base
let psi_nth_q_base: u64 = pow(psi, (self.q_base-1)/nth_root, self.q_base);
// lifts nth primitive root mod q_base to q = q_base^q_power
let psi_nth_q: u64 = self.hensel_lift(psi_nth_q_base, nth_root);
assert!(self.pow(psi_nth_q, nth_root) == 1, "invalid nth primitive root: psi^nth_root != 1 mod q");
assert!(self.pow(psi_nth_q, nth_root>>1) == self.q-1, "invalid nth primitive root: psi^(nth_root/2) != -1 mod q");
psi_nth_q
}
/// Checks if the field self.factor is populated.
/// If not, factorize q_base-1 and populates self.factor.
/// If yes, checks that it contains the unique factors of q_base-1.
pub fn check_factors(&mut self){
if self.factors.len() == 0{
let factors = Factorization::run(self.q_base-1).prime_factor_repr();
let mut distincts_factors: Vec<u64> = Vec::with_capacity(factors.len());
for factor in factors.iter(){
distincts_factors.push(factor.0)
}
self.factors = distincts_factors
}else{
let mut q_base: u64 = self.q_base;
for &factor in &self.factors{
if !is_prime(factor){
panic!("invalid factor list: factor {} is not prime", factor)
}
while q_base%factor != 0{
q_base /= factor
}
}
if q_base != 1{
panic!("invalid factor list: does not fully divide q_base: q_base % (all factors) = {}", q_base)
}
}
}
/// Returns (psi + a * q_base)^{nth_root} = 1 mod q = q_base^q_power given psi^{nth_root} = 1 mod q_base.
/// Panics if psi^{nth_root} != 1 mod q_base.
fn hensel_lift(&self, psi: u64, nth_root: u64) -> u64{
assert!(pow(psi, nth_root, self.q_base)==1, "invalid argument psi: psi^nth_root = {} != 1", pow(psi, nth_root, self.q_base));
let mut psi_mont: Montgomery<u64> = self.montgomery.prepare::<ONCE>(psi);
let nth_root_mont: Montgomery<u64> = self.montgomery.prepare::<ONCE>(nth_root);
for _i in 1..self.q_power{
let psi_pow: Montgomery<u64> = self.montgomery.pow(psi_mont, nth_root-1);
let num: Montgomery<u64> = self.montgomery.one() + self.q - self.montgomery.mul_internal::<ONCE>(psi_pow, psi_mont);
let mut den: Montgomery<u64> = self.montgomery.mul_internal::<ONCE>(nth_root_mont, psi_pow);
den = self.montgomery.pow(den, self.phi-1);
psi_mont = self.montgomery.add_internal(psi_mont, self.montgomery.mul_internal::<ONCE>(num, den));
}
self.montgomery.unprepare::<ONCE>(psi_mont)
}
}
/// Returns x^exponent mod q.
/// This function internally instantiate a new MontgomeryPrecomp<u64>
/// To be used when called only a few times and if there
/// is no Prime instantiated with q.
pub fn pow(x:u64, exponent:u64, q:u64) -> u64{
let montgomery: MontgomeryPrecomp<u64> = MontgomeryPrecomp::<u64>::new(q);
let mut y_mont: Montgomery<u64> = montgomery.one();
let mut x_mont: Montgomery<u64> = montgomery.prepare::<ONCE>(x);
let mut i: u64 = exponent;
while i > 0{
if i & 1 == 1{
montgomery.mul_internal_assign::<ONCE>(x_mont, &mut y_mont);
}
montgomery.mul_internal_assign::<ONCE>(x_mont, &mut x_mont);
i >>= 1;
}
montgomery.unprepare::<ONCE>(y_mont)
}

18
math/src/modulus/montgomery.rs Normal file
View File

@@ -0,0 +1,18 @@
use crate::modulus::barrett::BarrettPrecomp;
/// Montgomery is a generic struct storing
/// an element in the Montgomery domain.
pub type Montgomery<O> = O;
/// MontgomeryPrecomp is a generic struct storing
/// precomputations for Montgomery arithmetic.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct MontgomeryPrecomp<O>{
pub q: O,
pub two_q: O,
pub four_q: O,
pub barrett: BarrettPrecomp<O>,
pub q_inv: O,
pub one: Montgomery<O>,
pub minus_one: Montgomery<O>,
}

23
math/src/modulus/prime.rs Normal file
View File

@@ -0,0 +1,23 @@
use crate::modulus::montgomery::MontgomeryPrecomp;
use crate::modulus::barrett::BarrettPrecomp;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Prime<O> {
pub q: O, /// q_base^q_powers
pub two_q: O,
pub four_q: O,
pub q_base: O,
pub q_power: usize,
pub factors: Vec<O>, /// distinct factors of q-1
pub montgomery: MontgomeryPrecomp<O>,
pub barrett:BarrettPrecomp<O>,
pub phi: O,
}
pub struct NTTFriendlyPrimesGenerator<O>{
pub size: f64,
pub next_prime: O,
pub prev_prime: O,
pub check_next_prime: bool,
pub check_prev_prime: bool,
}

58
math/src/poly.rs Normal file
View File

@@ -0,0 +1,58 @@
pub mod poly;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Poly<O>(pub Vec<O>);
impl<O> Poly<O>where
O: Default + Clone,
{
pub fn new(n: usize) -> Self{
Self(vec![O::default();n])
}
pub fn buffer_size(&self) -> usize{
return self.0.len()
}
pub fn from_buffer(&mut self, n: usize, buf: &mut [O]){
assert!(buf.len() >= n, "invalid buffer: buf.len()={} < n={}", buf.len(), n);
self.0 = Vec::from(&buf[..n]);
}
pub fn n(&self) -> usize{
return self.0.len()
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PolyRNS<O>(pub Vec<Poly<O>>);
impl<O> PolyRNS<O>where
O: Default + Clone,
{
pub fn new(n: usize) -> Self{
Self(vec![Poly::<O>::new(n);n])
}
pub fn n(&self) -> usize{
self.0[0].n()
}
pub fn level(&self) -> usize{
self.0.len()-1
}
pub fn buffer_size(&self) -> usize{
self.n() * (self.level()+1)
}
pub fn from_buffer(&mut self, n: usize, level: usize, buf: &mut [O]){
assert!(buf.len() >= n * (level+1), "invalid buffer: buf.len()={} < n * (level+1)={}", buf.len(), level+1);
self.0.clear();
for chunk in buf.chunks_mut(n).take(level+1) {
let mut poly: Poly<O> = Poly(Vec::new());
poly.from_buffer(n, chunk);
self.0.push(poly);
}
}
}

1
math/src/poly/poly.rs Normal file
View File

@@ -0,0 +1 @@

14
math/src/ring.rs Normal file
View File

@@ -0,0 +1,14 @@
pub mod impl_u64;
use crate::modulus::prime::Prime;
use crate::dft::DFT;
pub struct Ring<O>{
pub n:usize,
pub modulus:Prime<O>,
pub dft:Box<dyn DFT<O>>,
}
pub struct RingRNS<O>(pub Vec<Ring<O>>);

43
math/src/ring/impl_u64/automorphism.rs Normal file
View File

@@ -0,0 +1,43 @@
use crate::modulus::WordOps;
use crate::ring::Ring;
use crate::poly::Poly;
/// Returns a lookup table for the automorphism X^{i} -> X^{i * k mod nth_root}.
/// Method will panic if n or nth_root are not power-of-two.
/// Method will panic if gal_el is not coprime with nth_root.
pub fn automorphism_index_ntt(n: usize, nth_root:u64, gal_el: u64) -> Vec<u64>{
assert!(n&(n-1) != 0, "invalid n={}: not a power-of-two", n);
assert!(nth_root&(nth_root-1) != 0, "invalid nth_root={}: not a power-of-two", n);
assert!(gal_el & 1 == 1, "invalid gal_el={}: not coprime with nth_root={}", gal_el, nth_root);
let mask = nth_root-1;
let log_nth_root: u32 = nth_root.log2() as u32;
let mut index: Vec<u64> = Vec::with_capacity(n);
for i in 0..n{
let i_rev: usize = 2*i.reverse_bits_msb(log_nth_root)+1;
let gal_el_i: u64 = (gal_el * (i_rev as u64) & mask)>>1;
index.push(gal_el_i.reverse_bits_msb(log_nth_root));
}
index
}
impl Ring<u64>{
pub fn automorphism(&self, a:Poly<u64>, gal_el: u64, b:&mut Poly<u64>){
debug_assert!(a.n() == b.n(), "invalid inputs: a.n() = {} != b.n() = {}", a.n(), b.n());
debug_assert!(gal_el&1 == 1, "invalid gal_el = {}: not odd", gal_el);
let n: usize = a.n();
let mask: u64 = (n-1) as u64;
let log_n: usize = n.log2();
let q: u64 = self.modulus.q();
let b_vec: &mut _ = &mut b.0;
let a_vec: &_ = &a.0;
a_vec.iter().enumerate().for_each(|(i, ai)|{
let gal_el_i: u64 = i as u64 * gal_el;
let sign: u64 = (gal_el_i>>log_n) & 1;
let i_out: u64 = gal_el_i & mask;
b_vec[i_out as usize] = ai * (sign^1) | (q - ai) * sign
});
}
}

3
math/src/ring/impl_u64/mod.rs Normal file
View File

@@ -0,0 +1,3 @@
pub mod automorphism;
pub mod ring;
pub mod ring_rns;

102
math/src/ring/impl_u64/ring.rs Normal file
View File

@@ -0,0 +1,102 @@
use crate::ring::Ring;
use crate::dft::ntt::Table;
use crate::modulus::prime::Prime;
use crate::modulus::montgomery::Montgomery;
use crate::poly::Poly;
use crate::modulus::REDUCEMOD;
use crate::modulus::VecOperations;
use crate::CHUNK;
impl Ring<u64>{
pub fn new(n:usize, q_base:u64, q_power:usize) -> Self{
let prime: Prime<u64> = Prime::<u64>::new(q_base, q_power);
Self {
n: n,
modulus: prime.clone(),
dft: Box::new(Table::<u64>::new(prime, (2 * n) as u64)),
}
}
pub fn n(&self) -> usize{
return self.n
}
pub fn new_poly(&self) -> Poly<u64>{
Poly::<u64>::new(self.n())
}
}
impl Ring<u64>{
pub fn ntt_inplace<const LAZY:bool>(&self, poly: &mut Poly<u64>){
match LAZY{
true => self.dft.forward_inplace_lazy(&mut poly.0),
false => self.dft.forward_inplace(&mut poly.0)
}
}
pub fn intt_inplace<const LAZY:bool>(&self, poly: &mut Poly<u64>){
match LAZY{
true => self.dft.forward_inplace_lazy(&mut poly.0),
false => self.dft.forward_inplace(&mut poly.0)
}
}
pub fn ntt<const LAZY:bool>(&self, poly_in: &Poly<u64>, poly_out: &mut Poly<u64>){
poly_out.0.copy_from_slice(&poly_in.0);
match LAZY{
true => self.dft.backward_inplace_lazy(&mut poly_out.0),
false => self.dft.backward_inplace(&mut poly_out.0)
}
}
pub fn intt<const LAZY:bool>(&self, poly_in: &Poly<u64>, poly_out: &mut Poly<u64>){
poly_out.0.copy_from_slice(&poly_in.0);
match LAZY{
true => self.dft.backward_inplace_lazy(&mut poly_out.0),
false => self.dft.backward_inplace(&mut poly_out.0)
}
}
}
impl Ring<u64>{
#[inline(always)]
pub fn add_inplace<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &mut Poly<u64>){
self.modulus.vec_add_unary_assign::<CHUNK, REDUCE>(&a.0, &mut b.0);
}
#[inline(always)]
pub fn add<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &Poly<u64>, c: &mut Poly<u64>){
self.modulus.vec_add_binary_assign::<CHUNK, REDUCE>(&a.0, &b.0, &mut c.0);
}
#[inline(always)]
pub fn sub_inplace<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &mut Poly<u64>){
self.modulus.vec_sub_unary_assign::<CHUNK, REDUCE>(&a.0, &mut b.0);
}
#[inline(always)]
pub fn sub<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &Poly<u64>, c: &mut Poly<u64>){
self.modulus.vec_sub_binary_assign::<CHUNK, REDUCE>(&a.0, &b.0, &mut c.0);
}
#[inline(always)]
pub fn neg<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &mut Poly<u64>){
self.modulus.vec_neg_binary_assign::<CHUNK, REDUCE>(&a.0, &mut b.0);
}
#[inline(always)]
pub fn neg_inplace<const REDUCE: REDUCEMOD>(&self, a: &mut Poly<u64>){
self.modulus.vec_neg_unary_assign::<CHUNK, REDUCE>(&mut a.0);
}
#[inline(always)]
pub fn mul_montgomery_external<const REDUCE:REDUCEMOD>(&self, a:&Poly<Montgomery<u64>>, b:&Poly<u64>, c: &mut Poly<u64>){
self.modulus.vec_mul_montgomery_external_binary_assign::<CHUNK, REDUCE>(&a.0, &b.0, &mut c.0);
}
#[inline(always)]
pub fn mul_montgomery_external_inplace<const REDUCE:REDUCEMOD>(&self, a:&Poly<Montgomery<u64>>, b:&mut Poly<u64>){
self.modulus.vec_mul_montgomery_external_unary_assign::<CHUNK, REDUCE>(&a.0, &mut b.0);
}
}

114
math/src/ring/impl_u64/ring_rns.rs Normal file
View File

@@ -0,0 +1,114 @@
use crate::ring::{Ring, RingRNS};
use crate::poly::PolyRNS;
use crate::modulus::montgomery::Montgomery;
use crate::modulus::REDUCEMOD;
impl RingRNS<u64>{
pub fn new(n:usize, moduli: Vec<u64>) -> Self{
assert!(!moduli.is_empty(), "moduli cannot be empty");
let rings: Vec<Ring<u64>> = moduli
.into_iter()
.map(|prime: u64| Ring::new(n, prime, 1))
.collect();
RingRNS(rings)
}
pub fn n(&self) -> usize{
self.0[0].n()
}
pub fn max_level(&self) -> usize{
self.0.len()-1
}
}
impl RingRNS<u64>{
}
impl RingRNS<u64>{
#[inline(always)]
pub fn add<const LEVEL:usize, const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &PolyRNS<u64>, c: &mut PolyRNS<u64>){
debug_assert!(self.max_level() <= LEVEL, "max_level={} < LEVEL={}", self.max_level(), LEVEL);
debug_assert!(a.n() >= self.n(), "a.n()={} < n={}", a.n(), self.n());
debug_assert!(b.n() >= self.n(), "b.n()={} < n={}", b.n(), self.n());
debug_assert!(c.n() >= self.n(), "c.n()={} < n={}", c.n(), self.n());
debug_assert!(a.level() >= LEVEL, "a.level()={} < LEVEL={}", a.level(), LEVEL);
debug_assert!(b.level() >= LEVEL, "b.level()={} < LEVEL={}", b.level(), LEVEL);
debug_assert!(c.level() >= LEVEL, "c.level()={} < LEVEL={}", c.level(), LEVEL);
self.0.iter().take(LEVEL + 1).enumerate().for_each(|(i, ring)| ring.add::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
}
#[inline(always)]
pub fn add_inplace<const LEVEL:usize, const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
debug_assert!(self.max_level() <= LEVEL, "max_level={} < LEVEL={}", self.max_level(), LEVEL);
debug_assert!(a.n() >= self.n(), "a.n()={} < n={}", a.n(), self.n());
debug_assert!(b.n() >= self.n(), "b.n()={} < n={}", b.n(), self.n());
debug_assert!(a.level() >= LEVEL, "a.level()={} < LEVEL={}", a.level(), LEVEL);
debug_assert!(b.level() >= LEVEL, "b.level()={} < LEVEL={}", b.level(), LEVEL);
self.0.iter().take(LEVEL + 1).enumerate().for_each(|(i, ring)| ring.add_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
}
#[inline(always)]
pub fn sub<const LEVEL:usize, const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &PolyRNS<u64>, c: &mut PolyRNS<u64>){
debug_assert!(self.max_level() <= LEVEL, "max_level={} < LEVEL={}", self.max_level(), LEVEL);
debug_assert!(a.n() >= self.n(), "a.n()={} < n={}", a.n(), self.n());
debug_assert!(b.n() >= self.n(), "b.n()={} < n={}", b.n(), self.n());
debug_assert!(c.n() >= self.n(), "c.n()={} < n={}", c.n(), self.n());
debug_assert!(a.level() >= LEVEL, "a.level()={} < LEVEL={}", a.level(), LEVEL);
debug_assert!(b.level() >= LEVEL, "b.level()={} < LEVEL={}", b.level(), LEVEL);
debug_assert!(c.level() >= LEVEL, "c.level()={} < LEVEL={}", c.level(), LEVEL);
self.0.iter().take(LEVEL + 1).enumerate().for_each(|(i, ring)| ring.sub::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
}
#[inline(always)]
pub fn sub_inplace<const LEVEL:usize, const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
debug_assert!(self.max_level() <= LEVEL, "max_level={} < LEVEL={}", self.max_level(), LEVEL);
debug_assert!(a.n() >= self.n(), "a.n()={} < n={}", a.n(), self.n());
debug_assert!(b.n() >= self.n(), "b.n()={} < n={}", b.n(), self.n());
debug_assert!(a.level() >= LEVEL, "a.level()={} < LEVEL={}", a.level(), LEVEL);
debug_assert!(b.level() >= LEVEL, "b.level()={} < LEVEL={}", b.level(), LEVEL);
self.0.iter().take(LEVEL + 1).enumerate().for_each(|(i, ring)| ring.sub_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
}
#[inline(always)]
pub fn neg<const LEVEL:usize, const REDUCE: REDUCEMOD>(&self, a: &PolyRNS<u64>, b: &mut PolyRNS<u64>){
debug_assert!(self.max_level() <= LEVEL, "max_level={} < LEVEL={}", self.max_level(), LEVEL);
debug_assert!(a.n() >= self.n(), "a.n()={} < n={}", a.n(), self.n());
debug_assert!(b.n() >= self.n(), "b.n()={} < n={}", b.n(), self.n());
debug_assert!(a.level() >= LEVEL, "a.level()={} < LEVEL={}", a.level(), LEVEL);
debug_assert!(b.level() >= LEVEL, "b.level()={} < LEVEL={}", b.level(), LEVEL);
self.0.iter().take(LEVEL + 1).enumerate().for_each(|(i, ring)| ring.neg::<REDUCE>(&a.0[i], &mut b.0[i]));
}
#[inline(always)]
pub fn neg_inplace<const LEVEL:usize, const REDUCE: REDUCEMOD>(&self, a: &mut PolyRNS<u64>){
debug_assert!(self.max_level() <= LEVEL, "max_level={} < LEVEL={}", self.max_level(), LEVEL);
debug_assert!(a.n() >= self.n(), "a.n()={} < n={}", a.n(), self.n());
debug_assert!(a.level() >= LEVEL, "a.level()={} < LEVEL={}", a.level(), LEVEL);
self.0.iter().take(LEVEL + 1).enumerate().for_each(|(i, ring)| ring.neg_inplace::<REDUCE>(&mut a.0[i]));
}
#[inline(always)]
pub fn mul_montgomery_external<const LEVEL:usize, const REDUCE:REDUCEMOD>(&self, a:&PolyRNS<Montgomery<u64>>, b:&PolyRNS<u64>, c: &mut PolyRNS<u64>){
debug_assert!(self.max_level() <= LEVEL, "max_level={} < LEVEL={}", self.max_level(), LEVEL);
debug_assert!(a.n() >= self.n(), "a.n()={} < n={}", a.n(), self.n());
debug_assert!(b.n() >= self.n(), "b.n()={} < n={}", b.n(), self.n());
debug_assert!(c.n() >= self.n(), "c.n()={} < n={}", c.n(), self.n());
debug_assert!(a.level() >= LEVEL, "a.level()={} < LEVEL={}", a.level(), LEVEL);
debug_assert!(b.level() >= LEVEL, "b.level()={} < LEVEL={}", b.level(), LEVEL);
debug_assert!(c.level() >= LEVEL, "c.level()={} < LEVEL={}", c.level(), LEVEL);
self.0.iter().take(LEVEL + 1).enumerate().for_each(|(i, ring)| ring.mul_montgomery_external::<REDUCE>(&a.0[i], &b.0[i], &mut c.0[i]));
}
#[inline(always)]
pub fn mul_montgomery_external_inplace<const LEVEL:usize, const REDUCE:REDUCEMOD>(&self, a:&PolyRNS<Montgomery<u64>>, b:&mut PolyRNS<u64>){
debug_assert!(self.max_level() <= LEVEL, "max_level={} < LEVEL={}", self.max_level(), LEVEL);
debug_assert!(a.n() >= self.n(), "a.n()={} < n={}", a.n(), self.n());
debug_assert!(b.n() >= self.n(), "b.n()={} < n={}", b.n(), self.n());
debug_assert!(a.level() >= LEVEL, "a.level()={} < LEVEL={}", a.level(), LEVEL);
debug_assert!(b.level() >= LEVEL, "b.level()={} < LEVEL={}", b.level(), LEVEL);
self.0.iter().take(LEVEL + 1).enumerate().for_each(|(i, ring)| ring.mul_montgomery_external_inplace::<REDUCE>(&a.0[i], &mut b.0[i]));
}
}