circuit-cas/src/poly/flat.rs

use itertools::Itertools;
use std::fmt::{self, Display};

use std::ops::{Add, Sub, BitXor, Mul};

use crate::fmt::num_to_superscript;
use super::var::{StaticVar, Var};
use std::collections::HashMap;

#[derive(Clone, Debug)]
pub struct Poly<V: Var> {
    mono: HashMap<Mono<V>, i32>,
}

impl<V: Var> Default for Poly<V> {
    fn default() -> Self {
        Poly {
            mono: Default::default(),
        }
    }
}

impl<V: Var> Display for Poly<V> {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        match self.mono.is_empty() {
            true => write!(fmt, "∅"),
            false => write!(
                fmt,
                "{}",
                self.mono
                    .iter()
                    .map(|(m, c)| format!("{}{}", c, m))
                    .join(" + ")
            ),
        }
    }
}

impl<V: Var, T: IntoIterator> From<T> for Poly<V>
where
    Poly<V>: FromIterator<<T as IntoIterator>::Item>,
{
    fn from(value: T) -> Self {
        value.into_iter().collect()
    }
}

impl<V: Var, U: Into<Mono<V>>> FromIterator<(i32, U)> for Poly<V> {
    fn from_iter<T: IntoIterator<Item = (i32, U)>>(iter: T) -> Self {
        Poly {
            mono: iter.into_iter().map(|(c, m)| (m.into(), c)).collect(),
        }
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
pub struct Mono<V: Var> {
    pub term: Vec<(V, u32)>,
}

impl<V:Var> Mono<V>{
    pub fn contains(&self, other:&Mono<V>)->bool{
        let mut self_it=self.term.iter().peekable();
        let mut other_it=other.term.iter().peekable();

        while let Some((o_term,o_exp))=other_it.peek(){
            if let Some((s_term,s_exp))=self_it.peek(){
                if s_term<o_term{
                    self_it.next();
                    continue;
                }else if s_term>o_term{
                    return false;
                }else if o_exp<=s_exp{
                    self_it.next();
                    other_it.next();
                }else{
                    return false;
                }
            }else{
                return false;
            }
        }

        true
    }
}

impl<V: Var, T: IntoIterator> From<T> for Mono<V>
where
    Mono<V>: FromIterator<<T as IntoIterator>::Item>,
{
    fn from(value: T) -> Self {
        value.into_iter().collect()
    }
}

impl<V: Var, U: Into<V>> FromIterator<(U, u32)> for Mono<V> {
    fn from_iter<T: IntoIterator<Item = (U, u32)>>(iter: T) -> Self {
        let mut term = iter
            .into_iter()
            .map(|(t, pow)| (t.into(), pow))
            .collect::<Vec<(V, u32)>>();
        term.sort();

        // Check duplicate variables
        assert!(
            (term[..])
                .windows(2)
                .all(|window| window[0].0 != window[1].0)
        );

        Mono { term }
    }
}

impl<V: Var> Display for Mono<V> {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        write!(
            fmt,
            "{}",
            self.term
                .iter()
                .map(|(t, p)| match p {
                    1 => format!("{t}"),
                    _ => format!("{t}{}",num_to_superscript(p.to_string())),
                })
                .join("")
        )
    }
}

impl BitXor<u32> for StaticVar {
    type Output = Mono<StaticVar>;

    fn bitxor(self, exp: u32) -> Self::Output {
        [(self, exp)].into()
    }
}

impl<'a> BitXor<u32> for &'a StaticVar {
    type Output = Mono<StaticVar>;

    fn bitxor(self, exp: u32) -> Self::Output {
        [(self.clone(), exp)].into()
    }
}

impl<V: Var> Mul for Mono<V> {
    type Output = Self;

    fn mul(self, other: Mono<V>) -> Self::Output {
        let mut a_term = self.term.into_iter().peekable();
        let mut b_term = other.term.into_iter().peekable();

        let mut result: Vec<(V, u32)> = Default::default();

        loop {
            match (a_term.peek(), b_term.peek()) {
                (Some((a_var, _)), Some((b_var, _))) => {
                    if a_var < b_var {
                        result.push(a_term.next().unwrap());
                    } else if a_var > b_var {
                        result.push(b_term.next().unwrap());
                    } else {
                        let (var, a_exp) = a_term.next().unwrap();
                        let (_, b_exp) = b_term.next().unwrap();
                        result.push((var, a_exp + b_exp));
                    }
                }
                (Some(a), None) => {
                    result.push(a.clone());
                    a_term.next();
                }
                (None, Some(b)) => {
                    result.push(b.clone());
                    b_term.next();
                }
                (None, None) => {
                    break;
                }
            }
        }

        Mono { term: result }
    }
}

impl<V: Var> Mul<Mono<V>> for i32 {
    type Output = Poly<V>;

    fn mul(self, mono: Mono<V>) -> Self::Output {
        let mut poly: HashMap<Mono<V>, i32> = Default::default();

        poly.insert(mono, self);

        Poly { mono: poly }
    }
}

impl<V: Var> Add for Poly<V> {
    type Output = Poly<V>;

    fn add(mut self, other: Poly<V>) -> Self::Output {
        for (mono, coeff) in other.mono {
            let entry = self.mono.entry(mono).or_insert(0);
            *entry += coeff;
        }
        self.mono.retain(|_, &mut coeff| coeff != 0);
        self
    }
}

impl<V: Var> Sub for Poly<V> {
    type Output = Poly<V>;

    fn sub(mut self, other: Poly<V>) -> Self::Output {
        for (mono, coeff) in other.mono {
            let entry = self.mono.entry(mono).or_insert(0);
            *entry -= coeff;
        }
        self.mono.retain(|_, &mut coeff| coeff != 0);
        self
    }
}