use std::fmt;
use std::ops::{Add, AddAssign, Mul, MulAssign};
use std::collections::BTreeMap;
use xops::binop;
use crate::core::num::*;
type I = u8;
type D = i8;
#[derive(Clone, PartialEq, Eq, Debug, Hash)]
pub struct MDeg(pub BTreeMap<I, D>);
#[derive(Clone, PartialEq, Eq, Debug, Hash)]
pub struct Term<F: Field> {
coef: F,
mdeg: MDeg,
}
#[derive(Clone, Debug, Hash)]
pub struct Polynomial<F: Field> {
terms: Vec<Term<F>>,
}
impl MDeg {
pub fn zero() -> Self {
MDeg(BTreeMap::new())
}
pub fn ones(n: I) -> Self {
Self::repeated(1, n)
}
fn repeated(deg: D, n: I) -> Self {
MDeg((0..n).map(|i| (i, deg)).collect())
}
pub fn from_tuples(tuples: &[(I, D)]) -> Self {
MDeg(
tuples
.iter()
.copied()
.filter(|&(_, deg)| deg != 0i8)
.collect()
)
}
pub fn single(idx: I, deg: D) -> Self {
MDeg::from_tuples(&[(idx, deg)])
}
pub fn degs(&self) -> impl Iterator<Item = (&I, &D)> {
self.0.iter()
}
pub fn degs_mut(&mut self) -> impl Iterator<Item = (&I, &mut D)> {
self.0.iter_mut()
}
pub fn total_deg(&self) -> D {
self.0.values().sum()
}
pub fn max_idx(&self) -> I {
*self.0.keys().max().unwrap_or(&0)
}
}
impl<F: Field> Term<F> {
#[inline]
pub fn from_coef_mdeg(coef: F, mdeg: MDeg) -> Self {
Term { coef, mdeg }
}
pub fn constant(coef: F) -> Self {
Term::from_coef_mdeg(coef, MDeg::zero())
}
pub fn zero() -> Self {
Term::constant(F::ZERO)
}
pub fn one() -> Self {
Term::constant(F::ONE)
}
pub fn monic(mdeg: MDeg) -> Self {
Term::from_coef_mdeg(F::ONE, mdeg)
}
pub fn var(idx: I, deg: D) -> Self {
Term::monic(MDeg::single(idx, deg))
}
pub fn eval(&self, x: &[F]) -> F {
if x.len() < self.mdeg.max_idx().into() {
panic!("incorrectly sized argument {:?} passed to term {:?}", x, self);
}
self.mdeg
.degs()
.map(|(&idx, °)| x[idx as usize].powi32(deg.into()))
.fold(self.coef, |acc, x| acc * x)
}
}
impl<F: Field> Polynomial<F> {
pub fn from_vec(v: Vec<Term<F>>) -> Self {
Polynomial { terms: v }
}
pub fn monomial(t: Term<F>) -> Self {
Self::from_vec(vec![t])
}
pub fn zero() -> Self {
Self::from_vec(Vec::new())
}
pub fn one() -> Self {
Self::monomial(Term::one())
}
pub fn term_count(&self) -> usize {
self.terms.len()
}
pub fn iter_terms(&self) -> impl Iterator<Item = &Term<F>> {
self.terms.iter()
}
pub fn eval(&self, x: &[F]) -> F {
self.terms
.iter()
.map(|t| t.eval(x))
.fold(F::ZERO, |acc, y| acc + y) }
}
impl AddAssign<&Self> for MDeg {
fn add_assign(&mut self, other: &Self) {
for (&idx, &other_deg) in other.degs() {
if let Some(self_deg_ref) = self.0.get_mut(&idx) {
*self_deg_ref += other_deg;
} else {
self.0.insert(idx, other_deg);
}
}
}
}
impl AddAssign for MDeg {
fn add_assign(&mut self, other: Self) {
*self += &other;
}
}
#[binop(refs_clone)]
impl Add for MDeg {
type Output = MDeg;
fn add(mut self, other: Self) -> Self::Output {
self += other;
self
}
}
impl<F: Field> MulAssign<&Self> for Term<F> {
fn mul_assign(&mut self, other: &Self) {
self.coef *= other.coef;
self.mdeg += &other.mdeg;
}
}
impl<F: Field> MulAssign for Term<F> {
fn mul_assign(&mut self, other: Self) {
*self *= &other;
}
}
#[binop(refs_clone)]
impl<F: Field> Mul for Term<F> {
type Output = Term<F>;
fn mul(mut self, other: Self) -> Self::Output {
self *= other;
self
}
}
impl<F: Field> AddAssign<&Term<F>> for Polynomial<F> {
fn add_assign(&mut self, rhs: &Term<F>) {
if let Some(term_ref) = self.terms.iter_mut().find(|t| t.mdeg == rhs.mdeg) {
term_ref.coef += rhs.coef; } else {
self.terms.push(rhs.clone());
}
}
}
impl<F: Field> AddAssign<Term<F>> for Polynomial<F> {
fn add_assign(&mut self, rhs: Term<F>) {
*self += &rhs;
}
}
#[binop(commute, refs_clone)]
impl<F: Field> Add<Term<F>> for Polynomial<F> {
type Output = Polynomial<F>;
fn add(mut self, rhs: Term<F>) -> Self::Output {
self += rhs;
self
}
}
impl<F: Field> AddAssign for Polynomial<F> {
fn add_assign(&mut self, other: Self) {
for t in other.terms {
*self += t
}
}
}
impl<F: Field> AddAssign<&Self> for Polynomial<F> {
fn add_assign(&mut self, other: &Self) {
for t in other.iter_terms() {
*self += t;
}
}
}
#[binop(refs_clone)]
impl<F: Field> Add for Polynomial<F> {
type Output = Polynomial<F>;
fn add(mut self, other: Self) -> Self::Output {
self += other;
self
}
}
impl<F: Field> MulAssign<&Term<F>> for Polynomial<F> {
fn mul_assign(&mut self, rhs: &Term<F>) {
for term_ref in self.terms.iter_mut() {
*term_ref *= rhs;
}
}
}
impl<F: Field> MulAssign<Term<F>> for Polynomial<F> {
fn mul_assign(&mut self, rhs: Term<F>) {
*self *= &rhs;
}
}
#[binop(refs_clone)]
impl<F: Field> Mul<Term<F>> for Polynomial<F> {
type Output = Polynomial<F>;
fn mul(mut self, rhs: Term<F>) -> Self::Output {
self *= rhs;
self
}
}
#[binop(derefs)]
impl<F: Field> Mul for &Polynomial<F> {
type Output = Polynomial<F>;
#[allow(clippy::suspicious_arithmetic_impl)]
fn mul(self, other: &Polynomial<F>) -> Self::Output {
other
.terms
.iter()
.fold(Polynomial::zero(), |acc, t| acc + self * t)
}
}
pub fn x<F: Field>(deg: D) -> Term<F> {
Term::var(0, deg)
}
pub fn y<F: Field>(deg: D) -> Term<F> {
Term::var(1, deg)
}
pub fn z<F: Field>(deg: D) -> Term<F> {
Term::var(0, deg)
}
impl fmt::Display for MDeg {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "MDeg{:?}", self.0)
}
}
impl<F: Field> fmt::Display for Term<F> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}·X{:?}", self.coef, self.mdeg.0)
}
}
impl<F: Field> fmt::Display for Polynomial<F> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let mut term_iter = self.terms.iter();
if let Some(term) = term_iter.next() {
write!(f, "{}", term)?;
} else {
return write!(f, "0_F[x]");
}
for term in term_iter {
write!(f, " + {}", term)?;
}
fmt::Result::Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_mdeg() {
let zero = MDeg::zero();
let ones = MDeg::ones(5);
println!("zero = {}", zero);
println!("ones = {}", ones);
let a = MDeg::from_tuples(&[(0, 3), (3, 5), (4, 2)]);
let b = MDeg::from_tuples(&[(0, 1), (1, 2), (4, 1), (5, 2)]);
let c = MDeg::from_tuples(&[(0, 4), (1, 2), (3, 5), (4, 3), (5, 2)]);
println!("\na = {}", a);
println!("b = {}", b);
let d = &a + &b;
println!("\nc = {}", d);
assert_eq!(a + b, c);
}
#[test]
fn test_term() {
let x = &Term::<f64>::var(0, 1);
let y = &Term::<f64>::var(1, 1);
let z = &Term::<f64>::var(2, 1);
println!("x = {}", x);
println!("y = {}", y);
println!("z = {}", z);
let t = &(x * y) * z;
println!("\nt = x * y * z = {}", t);
let c = Term::constant(37.0);
println!("\nc = {}", c);
let d = (&c * &c) * (&t * &t);
println!("\nc^2 * t^2 = {}", d);
println!("c^2 * t^2 = {:?}", d);
}
#[test]
fn test_poly() {
let c = |coef| Term::<f64>::constant(coef);
let x = |deg| Term::<f64>::var(0, deg);
let y = |deg| Term::<f64>::var(1, deg);
let z = |deg| Term::<f64>::var(2, deg);
let trm = |coef, [i, j, k]: [D; 3]| c(coef) * x(i) * y(j) * z(k);
println!("x^4 = {}", x(4));
println!("y^2 = {}", y(2));
println!("z^7 = {}", z(7));
let p = trm(5.0, [1, 2, 3]);
let q = trm(7.0, [4, 0, 2]);
let r = trm(13.0, [0, 5, 6]);
println!("\np = {}", p);
println!("q = {}", q);
println!("r = {}", r);
let f = Polynomial::from_vec(vec![p, q, r]);
println!("\nf = {}", f);
let p = trm(5.0, [1, 2, 3]);
let q = trm(7.0, [4, 0, 2]);
let r = trm(13.0, [0, 5, 6]);
let g = Polynomial::zero() + p + q + r;
println!("\nf = {}", g);
}
}