use num_bigint::{BigInt, BigUint, Sign};
use num_traits::Zero;
use crate::basis::Basis;
use crate::primes::mod_inverse;
use crate::RAYON_CHANNEL_THRESHOLD;
#[derive(Clone, Debug)]
pub struct RnsInt {
pub residues: Vec<u32>,
pub basis: Basis,
}
impl RnsInt {
pub fn from_bigint(n: &BigInt, basis: Basis) -> Self {
let residues = basis
.moduli()
.iter()
.map(|&m| {
let mm = BigInt::from(m);
let r = ((n % &mm) + &mm) % &mm;
r.to_biguint().unwrap().try_into().unwrap()
})
.collect();
RnsInt { residues, basis }
}
pub fn from_i64(n: i64, basis: Basis) -> Self {
Self::from_bigint(&BigInt::from(n), basis)
}
pub fn zero(basis: Basis) -> Self {
RnsInt { residues: vec![0; basis.len()], basis }
}
pub fn from_residues(residues: Vec<u32>, basis: Basis) -> Self {
RnsInt { residues, basis }
}
pub fn to_bigint(&self) -> BigInt {
let u = garner_crt(&self.residues, self.basis.moduli());
let m = self.basis.modulus_product();
if &u * 2u8 > m {
BigInt::from_biguint(Sign::Plus, u) - BigInt::from_biguint(Sign::Plus, m)
} else {
BigInt::from_biguint(Sign::Plus, u)
}
}
pub fn is_zero(&self) -> bool {
self.residues.iter().all(|&r| r == 0)
}
pub fn add(&self, other: &Self) -> Self {
let out = channel_map(&self.residues, &other.residues, self.basis.moduli(), add_channel);
Self::from_residues(out, self.basis.clone())
}
pub fn sub(&self, other: &Self) -> Self {
let out = channel_map(&self.residues, &other.residues, self.basis.moduli(), sub_channel);
Self::from_residues(out, self.basis.clone())
}
pub fn mul(&self, other: &Self) -> Self {
let out = channel_map(&self.residues, &other.residues, self.basis.moduli(), mul_channel);
Self::from_residues(out, self.basis.clone())
}
pub fn neg(&self) -> Self {
let out: Vec<u32> = self
.residues
.iter()
.zip(self.basis.moduli())
.map(|(&r, &m)| (m - r) % m)
.collect();
Self::from_residues(out, self.basis.clone())
}
}
fn channel_map(
a: &[u32],
b: &[u32],
moduli: &[u32],
f: impl Fn(u32, u32, u32) -> u32 + Sync + Send,
) -> Vec<u32> {
use rayon::prelude::*;
if a.len() >= RAYON_CHANNEL_THRESHOLD {
a.par_iter()
.zip(b.par_iter())
.zip(moduli.par_iter())
.map(|((&av, &bv), &m)| f(av, bv, m))
.collect()
} else {
a.iter()
.zip(b.iter())
.zip(moduli.iter())
.map(|((&av, &bv), &m)| f(av, bv, m))
.collect()
}
}
#[inline]
pub fn add_channel(a: u32, b: u32, m: u32) -> u32 {
(a + b) % m
}
#[inline]
pub fn sub_channel(a: u32, b: u32, m: u32) -> u32 {
(a + m - b) % m
}
#[inline]
pub fn mul_channel(a: u32, b: u32, m: u32) -> u32 {
(a * b) % m
}
pub fn garner_crt(residues: &[u32], moduli: &[u32]) -> BigUint {
let k = residues.len();
assert_eq!(k, moduli.len(), "residue/moduli length mismatch");
if k == 0 {
return BigUint::zero();
}
let mut c: Vec<u64> = residues.iter().map(|&r| r as u64).collect();
for i in 0..k {
let mi = moduli[i] as u64;
for j in 0..i {
let inv = mod_inverse(moduli[j] as u64 % mi, mi)
.expect("basis primes must be pairwise coprime for CRT");
let diff = (c[i] + mi - (c[j] % mi)) % mi;
c[i] = (diff * inv) % mi;
}
}
let mut result = BigUint::from(c[k - 1]);
for i in (0..k - 1).rev() {
result = result * BigUint::from(moduli[i]) + BigUint::from(c[i]);
}
result
}
pub fn crt_balanced(residues: &[u32], moduli: &[u32]) -> BigInt {
let u = garner_crt(residues, moduli);
let m: BigUint = moduli.iter().map(|&p| BigUint::from(p)).product();
if &u * 2u8 > m {
BigInt::from_biguint(Sign::Plus, u) - BigInt::from_biguint(Sign::Plus, m)
} else {
BigInt::from_biguint(Sign::Plus, u)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn b() -> Basis {
Basis::standard()
}
#[test]
fn roundtrip_positive() {
let a = RnsInt::from_i64(123_456_789, b());
assert_eq!(a.to_bigint(), BigInt::from(123_456_789));
}
#[test]
fn roundtrip_negative() {
let a = RnsInt::from_i64(-42, b());
assert_eq!(a.to_bigint(), BigInt::from(-42));
}
#[test]
fn add_sub_mul() {
let a = RnsInt::from_i64(1000, b());
let bb = RnsInt::from_i64(337, b());
assert_eq!(a.add(&bb).to_bigint(), BigInt::from(1337));
assert_eq!(a.sub(&bb).to_bigint(), BigInt::from(663));
assert_eq!(bb.sub(&a).to_bigint(), BigInt::from(-663));
assert_eq!(a.mul(&bb).to_bigint(), BigInt::from(337_000));
}
#[test]
fn garner_classic() {
assert_eq!(garner_crt(&[2, 3, 2], &[3, 5, 7]), BigUint::from(23u8));
assert_eq!(garner_crt(&[0, 1, 0], &[2, 3, 5]), BigUint::from(10u8));
}
#[test]
fn is_zero_works() {
assert!(RnsInt::zero(b()).is_zero());
assert!(!RnsInt::from_i64(1, b()).is_zero());
}
}