use flint_sys::{self, deps, fmpz::*, fmpz_mod::*, fmpz_mod_poly::*};
use rug::Integer;
use rug_fft;
use serde::de::Deserializer;
use serde::ser::Serializer;
use serde::{Deserialize, Serialize};
use std::cmp::*;
use std::fmt::{self, Debug, Display, Formatter};
use std::mem::MaybeUninit;
use std::ops::*;
mod flint_rug_bridge;
pub struct ModPoly {
raw: fmpz_mod_poly_struct,
ctx: fmpz_mod_ctx,
modulus: Integer,
}
impl ModPoly {
pub fn new(modulus: Integer) -> Self {
unsafe {
let mut raw = MaybeUninit::uninit();
let mut ctx = MaybeUninit::uninit();
let mut flint_modulus = flint_rug_bridge::int_to_fmpz(&modulus);
fmpz_mod_ctx_init(ctx.as_mut_ptr(), &flint_modulus);
fmpz_clear(&mut flint_modulus);
let ctx = ctx.assume_init();
fmpz_mod_poly_init(raw.as_mut_ptr(), &ctx as *const _ as *mut _);
ModPoly {
raw: raw.assume_init(),
ctx,
modulus,
}
}
}
pub fn from_int(modulus: Integer, mut constant: Integer) -> Self {
constant %= &modulus;
let mut this = ModPoly::new(modulus);
this.set_coefficient(0, &constant);
this
}
pub fn with_capacity(modulus: Integer, n: usize) -> Self {
unsafe {
let mut raw = MaybeUninit::uninit();
let mut flint_modulus = flint_rug_bridge::int_to_fmpz(&modulus);
let mut ctx = MaybeUninit::uninit();
fmpz_mod_ctx_init(ctx.as_mut_ptr(), &flint_modulus);
fmpz_clear(&mut flint_modulus);
let ctx = ctx.assume_init();
fmpz_mod_poly_init2(
raw.as_mut_ptr(),
n as deps::slong,
&ctx as *const _ as *mut _,
);
ModPoly {
raw: raw.assume_init(),
modulus,
ctx,
}
}
}
pub fn interpolate_from_mul_subgroup(mut ys: Vec<Integer>, m: Integer, w: &Integer) -> Self {
rug_fft::cooley_tukey_radix_2_intt(&mut ys, &m, w);
let mut p = ModPoly::with_capacity(m, ys.len());
for (i, c) in ys.iter().enumerate() {
p.set_coefficient(i, c);
}
p
}
pub fn evaluate_over_mul_subgroup(&self, w: &Integer, n: usize) -> Vec<Integer> {
let mut cs: Vec<Integer> = (0..n)
.into_iter()
.map(|i| self.get_coefficient(i))
.collect();
rug_fft::cooley_tukey_radix_2_ntt(&mut cs, &self.modulus, w);
cs
}
pub fn with_roots(xs: impl IntoIterator<Item = Integer>, m: &Integer) -> Self {
let mut ps = xs
.into_iter()
.map(|x| {
let mut p = ModPoly::new(m.clone());
p.set_coefficient_ui(1, 1);
p.set_coefficient(0, &-x);
p
})
.collect::<Vec<_>>();
while ps.len() > 1 {
for i in 0..(ps.len() / 2) {
let back = ps.pop().unwrap();
ps[i] *= &back;
}
}
ps.pop().unwrap_or_else(|| {
let mut p = ModPoly::new(m.clone());
p.set_coefficient_ui(0, 1);
p
})
}
pub fn reserve(&mut self, n: usize) {
unsafe {
fmpz_mod_poly_realloc(&mut self.raw, n as deps::slong, &mut self.ctx);
}
}
pub fn evaluate(&self, i: &Integer) -> Integer {
unsafe {
let mut in_ = flint_rug_bridge::int_to_fmpz(i);
let mut out = fmpz::default();
fmpz_init(&mut out);
fmpz_mod_poly_evaluate_fmpz(
&mut out,
&self.raw as *const _ as *mut _,
&mut in_,
&self.ctx as *const _ as *mut _,
);
let out_rug = flint_rug_bridge::fmpz_to_int(&out);
fmpz_clear(&mut in_);
fmpz_clear(&mut out);
out_rug
}
}
pub fn modulus(&self) -> &Integer {
&self.modulus
}
pub fn get_coefficient(&self, i: usize) -> Integer {
unsafe {
let mut c = fmpz::default();
fmpz_init(&mut c);
fmpz_mod_poly_get_coeff_fmpz(
&mut c,
&self.raw as *const _ as *mut _,
i as deps::slong,
&self.ctx as *const _ as *mut _,
);
let c_gmp = flint_rug_bridge::fmpz_to_int(&c);
fmpz_clear(&mut c);
c_gmp % &self.modulus
}
}
pub fn set_coefficient(&mut self, i: usize, c: &Integer) {
unsafe {
let mut c_flint = flint_rug_bridge::int_to_fmpz(c);
fmpz_mod_poly_set_coeff_fmpz(
&mut self.raw,
i as deps::slong,
&mut c_flint,
&mut self.ctx,
);
fmpz_clear(&mut c_flint);
}
}
pub fn set_coefficient_ui(&mut self, i: usize, c: usize) {
unsafe {
fmpz_mod_poly_set_coeff_ui(
&mut self.raw,
i as deps::slong,
c as deps::ulong,
&mut self.ctx,
);
}
}
pub fn len(&self) -> usize {
unsafe {
fmpz_mod_poly_length(
&self.raw as *const _ as *mut _,
&self.ctx as *const _ as *mut _,
) as usize
}
}
pub fn neg(&mut self) {
unsafe {
fmpz_mod_poly_neg(&mut self.raw, &mut self.raw, &mut self.ctx);
}
}
pub fn add(&mut self, other: &Self) {
assert_eq!(self.modulus, other.modulus);
unsafe {
fmpz_mod_poly_add(
&mut self.raw,
&mut self.raw,
&other.raw as *const _ as *mut _,
&mut self.ctx,
);
}
}
pub fn sub(&mut self, other: &Self) {
assert_eq!(self.modulus, other.modulus);
unsafe {
fmpz_mod_poly_sub(
&mut self.raw,
&mut self.raw,
&other.raw as *const _ as *mut _,
&mut self.ctx,
);
}
}
pub fn sub_from(&mut self, other: &Self) {
assert_eq!(self.modulus, other.modulus);
unsafe {
fmpz_mod_poly_sub(
&mut self.raw,
&other.raw as *const _ as *mut _,
&mut self.raw,
&mut self.ctx,
);
}
}
pub fn mul(&mut self, other: &Self) {
assert_eq!(self.modulus, other.modulus);
unsafe {
fmpz_mod_poly_mul(
&mut self.raw,
&mut self.raw,
&other.raw as *const _ as *mut _,
&mut self.ctx,
);
}
}
pub fn divrem(&self, other: &Self) -> (ModPoly, ModPoly) {
assert_eq!(self.modulus, other.modulus);
let mut q = ModPoly::new(self.modulus.clone());
let mut r = ModPoly::new(self.modulus.clone());
unsafe {
fmpz_mod_poly_divrem(
&mut q.raw,
&mut r.raw,
&self.raw as *const _ as *mut _,
&other.raw as *const _ as *mut _,
&self.ctx as *const _ as *mut _,
);
}
(q, r)
}
pub fn div(&mut self, other: &Self) {
assert_eq!(self.modulus, other.modulus);
let mut r = ModPoly::new(self.modulus.clone());
unsafe {
fmpz_mod_poly_divrem(
&mut self.raw,
&mut r.raw,
&mut self.raw,
&other.raw as *const _ as *mut _,
&mut self.ctx,
);
}
}
pub fn div_from(&mut self, other: &Self) {
assert_eq!(self.modulus, other.modulus);
let mut r = ModPoly::new(self.modulus.clone());
unsafe {
fmpz_mod_poly_divrem(
&mut self.raw,
&mut r.raw,
&other.raw as *const _ as *mut _,
&mut self.raw,
&mut self.ctx,
);
}
}
pub fn rem(&mut self, other: &Self) {
assert_eq!(self.modulus, other.modulus);
let mut q = ModPoly::new(self.modulus.clone());
unsafe {
fmpz_mod_poly_divrem(
&mut q.raw,
&mut self.raw,
&mut self.raw,
&other.raw as *const _ as *mut _,
&mut self.ctx,
);
}
}
pub fn rem_from(&mut self, other: &Self) {
assert_eq!(self.modulus, other.modulus);
let mut q = ModPoly::new(self.modulus.clone());
unsafe {
fmpz_mod_poly_divrem(
&mut q.raw,
&mut self.raw,
&other.raw as *const _ as *mut _,
&mut self.raw,
&mut self.ctx,
);
}
}
pub fn sqr(&mut self) {
unsafe {
fmpz_mod_poly_sqr(&mut self.raw, &mut self.raw, &mut self.ctx);
}
}
}
impl Clone for ModPoly {
fn clone(&self) -> Self {
let mut this = ModPoly::new(self.modulus.clone());
unsafe {
fmpz_mod_poly_set(
&mut this.raw,
&self.raw as *const _ as *mut _,
&self.ctx as *const _ as *mut _,
);
}
this
}
}
impl Drop for ModPoly {
fn drop(&mut self) {
unsafe { fmpz_mod_poly_clear(&mut self.raw, &mut self.ctx) }
}
}
impl PartialEq<ModPoly> for ModPoly {
fn eq(&self, other: &ModPoly) -> bool {
unsafe {
fmpz_mod_poly_equal(
&self.raw as *const _ as *mut _,
&other.raw as *const _ as *mut _,
&self.ctx as *const _ as *mut _,
) != 0
}
}
}
impl Eq for ModPoly {}
impl Debug for ModPoly {
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
f.debug_struct("ModPoly")
.field("modulus", &self.modulus)
.field(
"coefficients",
&(0..self.len())
.map(|i| self.get_coefficient(i))
.collect::<Vec<_>>(),
)
.finish()
}
}
impl Display for ModPoly {
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
let n = self.len();
let mut first = true;
for i in 0..n {
let j = n - i - 1;
let c = self.get_coefficient(j);
if c != 0 {
if !first {
write!(f, " + ")?;
}
write!(f, "{}", c)?;
if j != 0 {
write!(f, "x^{}", j)?;
}
first = false;
}
}
if first {
write!(f, "0")?;
}
Ok(())
}
}
macro_rules! impl_self_binary {
($Big:ty,
$func:ident,
$from_func:ident,
$Trait:ident { $method:ident },
$TraitAssign:ident { $method_assign:ident }
) => {
impl $Trait<$Big> for $Big {
type Output = $Big;
#[inline]
fn $method(mut self, rhs: $Big) -> $Big {
self.$method_assign(rhs);
self
}
}
impl $Trait<&$Big> for $Big {
type Output = $Big;
#[inline]
fn $method(mut self, rhs: &$Big) -> $Big {
self.$method_assign(rhs);
self
}
}
impl $Trait<$Big> for &$Big {
type Output = $Big;
#[inline]
fn $method(self, mut rhs: $Big) -> $Big {
<$Big>::$from_func(&mut rhs, self);
rhs
}
}
impl $TraitAssign<$Big> for $Big {
#[inline]
fn $method_assign(&mut self, rhs: $Big) {
<$Big>::$func(self, &rhs)
}
}
impl $TraitAssign<&$Big> for $Big {
#[inline]
fn $method_assign(&mut self, rhs: &$Big) {
<$Big>::$func(self, rhs)
}
}
};
}
macro_rules! impl_int_binary {
($Big:ty,
$Base:ty,
$func:ident,
$from_func:ident,
$lift_func:ident,
$Trait:ident { $method:ident },
$TraitAssign:ident { $method_assign:ident }
) => {
impl $Trait<$Base> for $Big {
type Output = $Big;
#[inline]
fn $method(mut self, rhs: $Base) -> $Big {
let rhs = <$Big>::$lift_func(self.modulus.clone(), rhs);
<$Big>::$func(&mut self, &rhs);
self
}
}
impl $Trait<$Big> for $Base {
type Output = $Big;
#[inline]
fn $method(self, mut rhs: $Big) -> $Big {
let lhs = <$Big>::$lift_func(rhs.modulus.clone(), self);
<$Big>::$from_func(&mut rhs, &lhs);
rhs
}
}
impl $TraitAssign<$Base> for $Big {
#[inline]
fn $method_assign(&mut self, rhs: $Base) {
let rhs = <$Big>::$lift_func(self.modulus.clone(), rhs);
<$Big>::$func(self, &rhs)
}
}
};
}
impl_self_binary!(ModPoly, add, add, Add { add }, AddAssign { add_assign });
impl_int_binary!(
ModPoly,
Integer,
add,
add,
from_int,
Add { add },
AddAssign { add_assign }
);
impl_self_binary!(
ModPoly,
sub,
sub_from,
Sub { sub },
SubAssign { sub_assign }
);
impl_int_binary!(
ModPoly,
Integer,
sub,
sub_from,
from_int,
Sub { sub },
SubAssign { sub_assign }
);
impl_self_binary!(ModPoly, mul, mul, Mul { mul }, MulAssign { mul_assign });
impl_int_binary!(
ModPoly,
Integer,
mul,
mul,
from_int,
Mul { mul },
MulAssign { mul_assign }
);
impl_self_binary!(
ModPoly,
div,
div_from,
Div { div },
DivAssign { div_assign }
);
impl_int_binary!(
ModPoly,
Integer,
div,
div_from,
from_int,
Div { div },
DivAssign { div_assign }
);
impl_self_binary!(
ModPoly,
rem,
rem_from,
Rem { rem },
RemAssign { rem_assign }
);
impl_int_binary!(
ModPoly,
Integer,
rem,
rem_from,
from_int,
Rem { rem },
RemAssign { rem_assign }
);
use std::convert::From;
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
pub struct ModPolySer {
pub modulus: Integer,
pub coefficients: Vec<Integer>,
}
impl From<ModPolySer> for ModPoly {
fn from(other: ModPolySer) -> ModPoly {
let mut inner = ModPoly::new(other.modulus.clone());
for (i, c) in other.coefficients.into_iter().enumerate() {
inner.set_coefficient(i, &c);
}
inner
}
}
impl From<&ModPoly> for ModPolySer {
fn from(other: &ModPoly) -> ModPolySer {
let modulus = other.modulus().clone();
let coefficients = (0..(other.len()))
.into_iter()
.map(|i| other.get_coefficient(i).clone())
.collect();
ModPolySer {
modulus,
coefficients,
}
}
}
impl Serialize for ModPoly {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
ModPolySer::from(self).serialize(serializer)
}
}
impl<'de> Deserialize<'de> for ModPoly {
fn deserialize<D>(deserializer: D) -> Result<ModPoly, D::Error>
where
D: Deserializer<'de>,
{
ModPolySer::deserialize(deserializer).map(ModPoly::from)
}
}
#[cfg(test)]
mod test {
use super::*;
use quickcheck;
use quickcheck_macros;
#[test]
fn init() {
let p = Integer::from(17);
let f = ModPoly::new(p);
assert_eq!(f.len(), 0);
}
#[test]
fn from_const() {
let p = Integer::from(17);
let f = ModPoly::from_int(p, Integer::from(0));
assert_eq!(f.len(), 0);
assert_eq!(f.evaluate(&Integer::from(0)), Integer::from(0));
}
#[test]
fn just_set() {
let p = Integer::from(17);
let mut f = ModPoly::new(p);
f.set_coefficient_ui(0, 1);
assert_eq!(f.len(), 1);
f.set_coefficient_ui(5, 1);
assert_eq!(f.len(), 6);
f.set_coefficient_ui(5, 0);
assert_eq!(f.len(), 1);
}
#[test]
fn set_get() {
let p = Integer::from(17);
let mut f = ModPoly::new(p);
f.set_coefficient_ui(0, 1);
assert_eq!(f.get_coefficient(0), Integer::from(1));
f.set_coefficient(5, &Integer::from(5));
for i in 1..5 {
assert_eq!(f.get_coefficient(i), Integer::from(0));
}
assert_eq!(f.get_coefficient(5), Integer::from(5));
}
#[test]
fn add() {
let p = Integer::from(17);
let mut f = ModPoly::new(p.clone());
f.set_coefficient_ui(0, 1);
let mut g = ModPoly::new(p);
g.set_coefficient_ui(3, 1);
let h = f.clone() + g.clone();
assert_eq!(h.get_coefficient(0), Integer::from(1));
assert_eq!(h.get_coefficient(1), Integer::from(0));
assert_eq!(h.get_coefficient(2), Integer::from(0));
assert_eq!(h.get_coefficient(3), Integer::from(1));
assert_eq!(h.len(), 4);
assert_eq!(h, f.clone() + &g);
assert_eq!(h, &f + g.clone());
assert_eq!(h, g.clone() + Integer::from(1));
assert_eq!(h, Integer::from(1) + g.clone());
}
#[test]
fn sub() {
let p = Integer::from(17);
let mut f = ModPoly::new(p.clone());
f.set_coefficient_ui(0, 1);
let mut g = ModPoly::new(p);
g.set_coefficient_ui(3, 1);
let h = f.clone() - g.clone();
assert_eq!(h.get_coefficient(0), Integer::from(1));
assert_eq!(h.get_coefficient(1), Integer::from(0));
assert_eq!(h.get_coefficient(2), Integer::from(0));
assert_eq!(h.get_coefficient(3), Integer::from(16));
assert_eq!(h.len(), 4);
assert_eq!(h, f.clone() - &g);
assert_eq!(h, &f - g.clone());
assert_eq!(h, Integer::from(1) - g.clone());
}
#[test]
fn mul() {
let p = Integer::from(17);
let mut f = ModPoly::new(p.clone());
f.set_coefficient_ui(1, 2);
let mut g = ModPoly::new(p);
g.set_coefficient_ui(3, 1);
let h = f.clone() * g.clone();
assert_eq!(h.get_coefficient(0), Integer::from(0));
assert_eq!(h.get_coefficient(1), Integer::from(0));
assert_eq!(h.get_coefficient(2), Integer::from(0));
assert_eq!(h.get_coefficient(3), Integer::from(0));
assert_eq!(h.get_coefficient(4), Integer::from(2));
assert_eq!(h.len(), 5);
assert_eq!(h, f.clone() * &g);
assert_eq!(h, &f * g.clone());
assert_eq!(h, h.clone() * Integer::from(1));
assert_eq!(h, Integer::from(1) * h.clone());
}
#[test]
fn mul_wrap() {
let p = Integer::from(17);
let mut g = ModPoly::new(p);
g.set_coefficient_ui(3, 1);
g.set_coefficient_ui(0, 5);
let h = g.clone() * Integer::from(4);
assert_eq!(h.get_coefficient(0), Integer::from(3));
assert_eq!(h.get_coefficient(1), Integer::from(0));
assert_eq!(h.get_coefficient(2), Integer::from(0));
assert_eq!(h.get_coefficient(3), Integer::from(4));
assert_eq!(h.len(), 4);
}
#[test]
fn div() {
let p = Integer::from(17);
let mut f = ModPoly::new(p.clone());
f.set_coefficient_ui(1, 1);
let mut g = ModPoly::new(p);
g.set_coefficient_ui(3, 1);
let h = g.clone() / f.clone();
assert_eq!(h.get_coefficient(0), Integer::from(0));
assert_eq!(h.get_coefficient(1), Integer::from(0));
assert_eq!(h.get_coefficient(2), Integer::from(1));
assert_eq!(h.len(), 3);
assert_eq!(h, g.clone() / &f);
assert_eq!(h, &g / f.clone());
assert_eq!(h, h.clone() / Integer::from(1));
}
fn test_interpolate_from_mul_subgroup(
ys: Vec<isize>,
m: usize,
w: usize,
expected_cs: Vec<isize>,
) {
let n = ys.len();
let p = ModPoly::interpolate_from_mul_subgroup(
ys.into_iter().map(Integer::from).collect(),
Integer::from(m),
&Integer::from(w),
);
for i in 0..n {
assert_eq!(
p.get_coefficient(i),
expected_cs[i],
"Difference in coefficient {}: expected {} but got {}",
i,
expected_cs[i],
p.get_coefficient(i)
);
}
}
#[test]
fn interpolate_zero_mod_5() {
test_interpolate_from_mul_subgroup(vec![0, 0, 0, 0], 5, 2, vec![0, 0, 0, 0]);
}
#[test]
fn interpolate_const_mod_5() {
test_interpolate_from_mul_subgroup(vec![3, 3, 3, 3], 5, 2, vec![3, 0, 0, 0]);
}
#[test]
fn interpolate_line_mod_5() {
test_interpolate_from_mul_subgroup(vec![1, 0, 3, 4], 5, 2, vec![2, 4, 0, 0]);
}
#[test]
fn interpolate_poly_mod_5() {
test_interpolate_from_mul_subgroup(vec![4, 0, 0, 0], 5, 2, vec![1, 1, 1, 1]);
}
#[derive(Debug, Clone)]
struct Usize16([u32; 16]);
impl quickcheck::Arbitrary for Usize16 {
fn arbitrary<G: quickcheck::Gen>(g: &mut G) -> Self {
let mut a = [0u32; 16];
for i in &mut a {
*i = g.next_u32();
}
Usize16(a)
}
}
#[quickcheck_macros::quickcheck]
fn test_interpolate_rountrip(ys: Usize16) -> bool {
let m = Integer::from(17);
let w = Integer::from(3);
let Usize16(mut ys) = ys;
for i in &mut ys {
*i %= 17;
}
let ys: Vec<Integer> = ys.iter().cloned().map(Integer::from).collect();
let p = ModPoly::interpolate_from_mul_subgroup(ys.clone(), m.clone(), &w);
let ys2 = p.evaluate_over_mul_subgroup(&w, 16);
ys == ys2
}
}