simsym/
poly.rs

1//! Univariate polynomial normal form: combine like terms, order by descending degree.
2
3use crate::expr::{add, const_, mul, neg, pow, sub, Expr, ExprKind};
4use crate::rational::Rational;
5use crate::symbol::Symbol;
6use std::collections::BTreeMap;
7
8/// If `expr` is a polynomial in a single variable, return it in standard form (highest degree first).
9pub fn polynomial_normal_form(expr: Expr) -> Expr {
10    try_polynomial_normal_form(expr.clone()).unwrap_or(expr)
11}
12
13pub fn try_polynomial_normal_form(expr: Expr) -> Option<Expr> {
14    let mut summands = Vec::new();
15    flatten_summands(&expr, Rational::one(), &mut summands);
16    build_polynomial(&summands)
17}
18
19fn flatten_summands(expr: &Expr, sign: Rational, out: &mut Vec<Expr>) {
20    if sign.is_zero() {
21        return;
22    }
23    match expr.kind() {
24        ExprKind::Add(a, b) => {
25            flatten_summands(a, sign, out);
26            flatten_summands(b, sign, out);
27        }
28        ExprKind::Sub(a, b) => {
29            flatten_summands(a, sign, out);
30            flatten_summands(b, -sign, out);
31        }
32        ExprKind::Neg(e) => flatten_summands(e, -sign, out),
33        _ => out.push(scale_expr(sign, expr)),
34    }
35}
36
37fn scale_expr(sign: Rational, e: &Expr) -> Expr {
38    if sign.is_one() {
39        e.clone()
40    } else if sign == -Rational::one() {
41        neg(e.clone())
42    } else {
43        mul(const_(sign), e.clone())
44    }
45}
46
47fn build_polynomial(summands: &[Expr]) -> Option<Expr> {
48    let mut by_degree: BTreeMap<i64, Rational> = BTreeMap::new();
49    let mut var: Option<Symbol> = None;
50
51    for term in summands {
52        let (v_opt, degree, coeff) = parse_univariate_term(term)?;
53        if let Some(v) = v_opt {
54            if let Some(existing) = var {
55                if existing != v {
56                    return None;
57                }
58            } else {
59                var = Some(v);
60            }
61        } else if degree != 0 {
62            return None;
63        }
64        *by_degree.entry(degree).or_insert(Rational::zero()) += coeff;
65    }
66
67    let var = match var {
68        Some(v) => v,
69        None => {
70            let c = by_degree.get(&0).copied().unwrap_or(Rational::zero());
71            return Some(const_(c));
72        }
73    };
74    let mut degrees: Vec<i64> = by_degree
75        .iter()
76        .filter(|(_, c)| !c.is_zero())
77        .map(|(&d, _)| d)
78        .collect();
79    if degrees.is_empty() {
80        return Some(const_(Rational::zero()));
81    }
82    degrees.sort_by(|a, b| b.cmp(a));
83
84    let mut acc = term_for_degree(var, degrees[0], by_degree[&degrees[0]])?;
85    for &d in &degrees[1..] {
86        let c = by_degree[&d];
87        let t = term_for_degree(var, d, if c.is_positive() { c } else { -c })?;
88        acc = if c.is_positive() {
89            add(acc, t)
90        } else {
91            sub(acc, t)
92        };
93    }
94    Some(acc)
95}
96
97fn parse_univariate_term(e: &Expr) -> Option<(Option<Symbol>, i64, Rational)> {
98    match e.kind() {
99        ExprKind::Neg(inner) => {
100            let (v, d, c) = parse_univariate_term(inner)?;
101            Some((v, d, -c))
102        }
103        _ => {
104            let (coeff, mono) = coeff_monomial(e)?;
105            match mono.kind() {
106                ExprKind::Const(_) => Some((None, 0, coeff)),
107                ExprKind::Var(s) => Some((Some(*s), 1, coeff)),
108                ExprKind::Pow(base, exp) => {
109                    let s = match base.kind() {
110                        ExprKind::Var(s) => *s,
111                        _ => return None,
112                    };
113                    let n = as_const(exp)?.as_integer()?;
114                    if n < 0 {
115                        return None;
116                    }
117                    Some((Some(s), n, coeff))
118                }
119                _ => None,
120            }
121        }
122    }
123}
124
125fn term_for_degree(var: Symbol, degree: i64, coeff: Rational) -> Option<Expr> {
126    if coeff.is_zero() {
127        return Some(const_(Rational::zero()));
128    }
129    if degree == 0 {
130        return Some(const_(coeff));
131    }
132    let x = Expr::var(var);
133    let mono = if degree == 1 {
134        x
135    } else {
136        pow(x, const_(Rational::from(degree)))
137    };
138    if coeff.is_one() {
139        Some(mono)
140    } else {
141        Some(mul(const_(coeff), mono))
142    }
143}
144
145fn as_const(e: &Expr) -> Option<Rational> {
146    match e.kind() {
147        ExprKind::Const(c) => c.try_as_rational(),
148        _ => None,
149    }
150}
151
152fn coeff_monomial(e: &Expr) -> Option<(Rational, Expr)> {
153    match e.kind() {
154        ExprKind::Const(c) => c
155            .try_as_rational()
156            .map(|r| (r, const_(Rational::one()))),
157        ExprKind::Var(_) => Some((Rational::one(), e.clone())),
158        ExprKind::Mul(l, r) => {
159            if let Some(c) = as_const(l) {
160                let (_, m) = coeff_monomial(r)?;
161                return Some((c, m));
162            }
163            if let Some(c) = as_const(r) {
164                let (_, m) = coeff_monomial(l)?;
165                return Some((c, m));
166            }
167            None
168        }
169        ExprKind::Pow(base, exp) => {
170            let n = as_const(exp)?;
171            if let Some(k) = n.as_integer() {
172                if k >= 0 {
173                    return Some((Rational::one(), pow(base.clone(), const_(n))));
174                }
175            }
176            None
177        }
178        ExprKind::Div(num, den) => {
179            let cd = as_const(den)?;
180            if cd.is_zero() {
181                return None;
182            }
183            let (cn, m) = coeff_monomial(num)?;
184            Some((cn / cd, m))
185        }
186        _ => None,
187    }
188}
simsym/poly.rs

simsym/
poly.rs
