use alloc::vec;
use alloc::vec::Vec;
use puremp::Rational;
use crate::math::polynomial::{Polynomial, Var};
use crate::nlsat::icp::Rel;
fn zero() -> Rational {
Rational::from_integer(0.into())
}
fn one() -> Rational {
Rational::from_integer(1.into())
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct UPoly {
coeffs: Vec<Rational>,
}
impl UPoly {
fn zero() -> UPoly {
UPoly { coeffs: Vec::new() }
}
fn from_coeffs(mut coeffs: Vec<Rational>) -> UPoly {
while coeffs.last().is_some_and(|c| c.is_zero()) {
coeffs.pop();
}
UPoly { coeffs }
}
fn is_zero(&self) -> bool {
self.coeffs.is_empty()
}
fn degree(&self) -> usize {
self.coeffs.len().saturating_sub(1)
}
fn lead(&self) -> Rational {
self.coeffs.last().cloned().unwrap_or_else(zero)
}
fn eval(&self, x: &Rational) -> Rational {
let mut acc = zero();
for c in self.coeffs.iter().rev() {
acc = acc.mul(x).add(c);
}
acc
}
fn sign_at(&self, x: &Rational) -> i32 {
self.eval(x).signum()
}
fn deriv(&self) -> UPoly {
if self.coeffs.len() <= 1 {
return UPoly::zero();
}
let mut c = Vec::with_capacity(self.coeffs.len() - 1);
for (i, coeff) in self.coeffs.iter().enumerate().skip(1) {
c.push(coeff.mul(&Rational::from_integer((i as i64).into())));
}
UPoly::from_coeffs(c)
}
fn scale(&self, s: &Rational) -> UPoly {
if s.is_zero() {
return UPoly::zero();
}
UPoly::from_coeffs(self.coeffs.iter().map(|c| c.mul(s)).collect())
}
fn neg(&self) -> UPoly {
UPoly::from_coeffs(self.coeffs.iter().map(|c| c.neg()).collect())
}
fn sub(&self, other: &UPoly) -> UPoly {
let n = self.coeffs.len().max(other.coeffs.len());
let mut c = vec![zero(); n];
for (i, a) in self.coeffs.iter().enumerate() {
c[i] = c[i].add(a);
}
for (i, b) in other.coeffs.iter().enumerate() {
c[i] = c[i].sub(b);
}
UPoly::from_coeffs(c)
}
fn mul(&self, other: &UPoly) -> UPoly {
if self.is_zero() || other.is_zero() {
return UPoly::zero();
}
let mut c = vec![zero(); self.coeffs.len() + other.coeffs.len() - 1];
for (i, a) in self.coeffs.iter().enumerate() {
for (j, b) in other.coeffs.iter().enumerate() {
c[i + j] = c[i + j].add(&a.mul(b));
}
}
UPoly::from_coeffs(c)
}
fn rem(&self, divisor: &UPoly) -> UPoly {
debug_assert!(!divisor.is_zero());
let mut r = self.clone();
let d_deg = divisor.degree();
let d_lead = divisor.lead();
while !r.is_zero() && r.degree() >= d_deg {
let shift = r.degree() - d_deg;
let factor = r.lead().div(&d_lead);
let mut sub = vec![zero(); shift + divisor.coeffs.len()];
for (i, c) in divisor.coeffs.iter().enumerate() {
sub[i + shift] = c.mul(&factor);
}
r = r.sub(&UPoly::from_coeffs(sub));
}
r
}
fn gcd(&self, other: &UPoly) -> UPoly {
let mut a = self.clone();
let mut b = other.clone();
while !b.is_zero() {
let r = a.rem(&b);
a = b;
b = r;
}
if a.is_zero() {
a
} else {
let lead = a.lead();
a.scale(&lead.recip())
}
}
fn squarefree(&self) -> UPoly {
if self.degree() == 0 {
return self.clone();
}
let g = self.gcd(&self.deriv());
if g.degree() == 0 {
return self.clone();
}
self.div_exact(&g)
}
fn div_exact(&self, divisor: &UPoly) -> UPoly {
let mut r = self.clone();
let d_deg = divisor.degree();
let d_lead = divisor.lead();
let mut q = vec![zero(); self.degree().saturating_sub(d_deg) + 1];
while !r.is_zero() && r.degree() >= d_deg {
let shift = r.degree() - d_deg;
let factor = r.lead().div(&d_lead);
q[shift] = factor.clone();
let mut sub = vec![zero(); shift + divisor.coeffs.len()];
for (i, c) in divisor.coeffs.iter().enumerate() {
sub[i + shift] = c.mul(&factor);
}
r = r.sub(&UPoly::from_coeffs(sub));
}
UPoly::from_coeffs(q)
}
fn root_bound(&self) -> Rational {
if self.degree() == 0 {
return one();
}
let lead = self.lead().abs();
let mut m = zero();
for c in &self.coeffs[..self.coeffs.len() - 1] {
let ratio = c.abs().div(&lead);
if ratio > m {
m = ratio;
}
}
m.add(&one())
}
}
fn sturm_chain(p: &UPoly) -> Vec<UPoly> {
let mut chain = vec![p.clone(), p.deriv()];
while !chain.last().unwrap().is_zero() {
let n = chain.len();
let r = chain[n - 2].rem(&chain[n - 1]);
if r.is_zero() {
break;
}
chain.push(r.neg());
}
chain
}
fn variations(chain: &[UPoly], x: &Rational) -> usize {
let mut last = 0i32;
let mut count = 0;
for s in chain {
let sg = s.sign_at(x);
if sg == 0 {
continue;
}
if last != 0 && sg != last {
count += 1;
}
last = sg;
}
count
}
fn root_count(chain: &[UPoly], a: &Rational, b: &Rational) -> i64 {
variations(chain, a) as i64 - variations(chain, b) as i64
}
fn isolate_roots(p: &UPoly) -> Vec<(Rational, Rational)> {
if p.degree() == 0 {
return Vec::new();
}
let chain = sturm_chain(p);
let m = p.root_bound();
let two = Rational::from_integer(2.into());
let mut out = Vec::new();
let mut stack = vec![(m.neg(), m)];
let mut guard = 0;
while let Some((a, b)) = stack.pop() {
guard += 1;
if guard > 200_000 {
break;
}
let n = root_count(&chain, &a, &b);
if n <= 0 {
continue;
}
if n == 1 {
out.push((a, b));
continue;
}
let mut mid = a.add(&b).div(&two);
let step = b.sub(&a).div(&Rational::from_integer(1024.into()));
let mut tries = 0;
while p.sign_at(&mid) == 0 && tries < 2048 {
mid = mid.add(&step);
tries += 1;
}
stack.push((a, mid.clone()));
stack.push((mid, b));
}
out.sort_by(|x, y| x.0.cmp(&y.0));
out
}
fn to_upoly(p: &Polynomial, var: Var) -> UPoly {
let mut coeffs: Vec<Rational> = Vec::new();
for (c, m) in p.terms() {
let d = m.degree_of(var) as usize;
if coeffs.len() <= d {
coeffs.resize(d + 1, zero());
}
coeffs[d] = coeffs[d].add(c);
}
UPoly::from_coeffs(coeffs)
}
fn sign_satisfies(s: i32, rel: Rel) -> bool {
match rel {
Rel::Lt => s < 0,
Rel::Le => s <= 0,
Rel::Gt => s > 0,
Rel::Ge => s >= 0,
Rel::Eq => s == 0,
Rel::Ne => s != 0,
}
}
pub fn decide(constraints: &[(Polynomial, Rel)], var: Var) -> Option<bool> {
let mut ups: Vec<(UPoly, Rel)> = Vec::new();
for (p, rel) in constraints {
for v in p.vars() {
if v != var {
return None;
}
}
if p.degree_of(var) > 24 {
return None;
}
ups.push((to_upoly(p, var), *rel));
}
if ups.is_empty() {
return Some(true);
}
let mut prod = UPoly::from_coeffs(vec![one()]);
for (p, _) in &ups {
if !p.is_zero() {
prod = prod.mul(&p.squarefree());
}
}
let roots = if prod.degree() == 0 {
Vec::new()
} else {
isolate_roots(&prod.squarefree())
};
let mut open_samples: Vec<Rational> = Vec::new();
if roots.is_empty() {
open_samples.push(zero());
} else {
let two = Rational::from_integer(2.into());
open_samples.push(roots[0].0.sub(&one())); for w in roots.windows(2) {
let a = &w[0].1; let b = &w[1].0; open_samples.push(a.add(b).div(&two));
}
open_samples.push(roots.last().unwrap().1.add(&one())); }
for q in &open_samples {
if ups
.iter()
.all(|(p, rel)| sign_satisfies(p.sign_at(q), *rel))
{
return Some(true);
}
}
for (a, b) in &roots {
if ups.iter().all(|(p, rel)| {
let s = root_sign(p, a, b);
sign_satisfies(s, *rel)
}) {
return Some(true);
}
}
Some(false)
}
fn integer_roots(p: &Polynomial, var: Var) -> Option<Vec<i64>> {
let up = to_upoly(p, var);
if up.is_zero() || up.degree() == 0 {
return Some(Vec::new());
}
let mut denom = puremp::Int::from(1);
for c in &up.coeffs {
denom = lcm_int(&denom, c.denominator());
}
let int_coeffs: Vec<puremp::Int> = up
.coeffs
.iter()
.map(|c| {
c.mul(&Rational::from_integer(denom.clone()))
.numerator()
.clone()
})
.collect();
let mut cands: Vec<i64> = Vec::new();
let m = int_coeffs.iter().position(|c| !c.is_zero());
let Some(m) = m else {
return Some(Vec::new()); };
if m > 0 {
cands.push(0);
}
let a0 = int_coeffs[m].abs();
let Some(a0_i) = i64_of(&a0) else {
return None; };
if a0_i > 20_000_000 {
return None;
}
let mut d = 1i64;
while d * d <= a0_i {
if a0_i % d == 0 {
cands.push(d);
cands.push(-d);
cands.push(a0_i / d);
cands.push(-(a0_i / d));
}
d += 1;
}
let zero_r = zero();
let mut roots: Vec<i64> = cands
.into_iter()
.filter(|&c| {
let x = Rational::from_integer(puremp::Int::from(c));
up.eval(&x) == zero_r
})
.collect();
roots.sort_unstable();
roots.dedup();
Some(roots)
}
fn lcm_int(a: &puremp::Int, b: &puremp::Int) -> puremp::Int {
if a.is_zero() || b.is_zero() {
return puremp::Int::from(0);
}
let g = a.gcd(b);
a.div_exact(&g).mul(b).abs()
}
fn i64_of(n: &puremp::Int) -> Option<i64> {
n.to_i64()
}
pub fn decide_int(constraints: &[(Polynomial, Rel)], var: Var) -> Option<bool> {
for (p, _) in constraints {
for v in p.vars() {
if v != var {
return None;
}
}
}
if decide(constraints, var) == Some(false) {
return Some(false);
}
let mut candidates: Option<Vec<i64>> = None;
for (p, rel) in constraints {
if *rel == Rel::Eq && p.degree_of(var) >= 1 {
let roots = integer_roots(p, var)?;
candidates = Some(match candidates {
None => roots,
Some(prev) => prev.into_iter().filter(|r| roots.contains(r)).collect(),
});
}
}
const SAMPLE: i64 = 256;
let (cands, exhaustive) = match candidates {
Some(c) => (c, true),
None => ((-SAMPLE..=SAMPLE).collect::<Vec<i64>>(), false),
};
for c in cands {
let x = Rational::from_integer(puremp::Int::from(c));
if constraints.iter().all(|(p, rel)| {
sign_satisfies(
p.eval(&|v| if v == var { x.clone() } else { zero() })
.signum(),
*rel,
)
}) {
return Some(true);
}
}
if exhaustive {
return Some(false);
}
let bound = Rational::from_integer(puremp::Int::from(SAMPLE));
let x = Polynomial::var(var);
let mut hi = constraints.to_vec();
hi.push((x.sub(&Polynomial::constant(bound.clone())), Rel::Ge)); let mut lo = constraints.to_vec();
lo.push((x.add(&Polynomial::constant(bound)), Rel::Le)); if decide(&hi, var) == Some(false) && decide(&lo, var) == Some(false) {
return Some(false);
}
None
}
fn root_sign(p: &UPoly, a: &Rational, b: &Rational) -> i32 {
if p.is_zero() {
return 0;
}
let sf = p.squarefree();
let sa = sf.sign_at(a);
let sb = sf.sign_at(b);
if sa != sb {
return 0;
}
let ea = p.sign_at(a);
if ea != 0 { ea } else { p.sign_at(b) }
}
#[cfg(test)]
mod tests {
use super::*;
use crate::math::polynomial::Monomial;
fn r(n: i64) -> Rational {
Rational::from_integer(n.into())
}
fn term(coeff: i64, deg: u32) -> (Rational, Monomial) {
(r(coeff), Monomial::from_powers(&[(0, deg)]))
}
fn poly(terms: &[(i64, u32)]) -> Polynomial {
Polynomial::from_terms(terms.iter().map(|&(c, d)| term(c, d)).collect())
}
#[test]
fn square_eq_four_sat() {
let c = vec![(poly(&[(1, 2), (-4, 0)]), Rel::Eq)];
assert_eq!(decide(&c, 0), Some(true));
}
#[test]
fn square_eq_two_sat() {
let c = vec![(poly(&[(1, 2), (-2, 0)]), Rel::Eq)];
assert_eq!(decide(&c, 0), Some(true));
}
#[test]
fn square_eq_negative_unsat() {
let c = vec![(poly(&[(1, 2), (1, 0)]), Rel::Eq)];
assert_eq!(decide(&c, 0), Some(false));
}
#[test]
fn square_eq_nine_positive_sat() {
let c = vec![
(poly(&[(1, 2), (-9, 0)]), Rel::Eq),
(poly(&[(1, 1)]), Rel::Gt),
];
assert_eq!(decide(&c, 0), Some(true));
}
#[test]
fn contradictory_bounds_unsat() {
let c = vec![
(poly(&[(1, 2), (-9, 0)]), Rel::Eq),
(poly(&[(1, 1)]), Rel::Gt),
(poly(&[(1, 1)]), Rel::Lt),
];
assert_eq!(decide(&c, 0), Some(false));
}
#[test]
fn square_lt_four_sat() {
let c = vec![(poly(&[(1, 2), (-4, 0)]), Rel::Lt)];
assert_eq!(decide(&c, 0), Some(true));
}
#[test]
fn square_gt_and_between_unsat() {
let c = vec![
(poly(&[(1, 2), (-4, 0)]), Rel::Gt),
(poly(&[(1, 1), (-2, 0)]), Rel::Lt),
(poly(&[(1, 1), (2, 0)]), Rel::Gt),
];
assert_eq!(decide(&c, 0), Some(false));
}
#[test]
fn cubic_roots_constrained_unsat() {
let c = vec![
(poly(&[(1, 3), (-1, 1)]), Rel::Eq),
(poly(&[(1, 1)]), Rel::Gt),
(poly(&[(1, 1), (-1, 0)]), Rel::Lt),
];
assert_eq!(decide(&c, 0), Some(false));
}
#[test]
fn multivariate_declined() {
let p = Polynomial::from_terms(vec![(r(1), Monomial::from_powers(&[(0, 1), (1, 1)]))]);
assert_eq!(decide(&[(p, Rel::Eq)], 0), None);
}
#[test]
fn integer_square_roots() {
assert_eq!(
decide_int(&[(poly(&[(1, 2), (-9, 0)]), Rel::Eq)], 0),
Some(true)
);
assert_eq!(
decide_int(&[(poly(&[(1, 2), (-2, 0)]), Rel::Eq)], 0),
Some(false)
);
assert_eq!(
decide_int(
&[
(poly(&[(1, 2), (-9, 0)]), Rel::Eq),
(poly(&[(1, 1)]), Rel::Gt)
],
0
),
Some(true)
);
assert_eq!(
decide_int(
&[
(poly(&[(1, 2), (-9, 0)]), Rel::Eq),
(poly(&[(1, 1)]), Rel::Gt),
(poly(&[(1, 1), (-3, 0)]), Rel::Lt),
],
0
),
Some(false)
);
}
#[test]
fn integer_inherits_real_unsat() {
assert_eq!(
decide_int(&[(poly(&[(1, 2), (1, 0)]), Rel::Eq)], 0),
Some(false)
);
}
#[test]
fn integer_roots_with_zero_constant() {
let c = vec![
(poly(&[(1, 2), (-7, 1)]), Rel::Eq),
(poly(&[(1, 1)]), Rel::Gt),
];
assert_eq!(decide_int(&c, 0), Some(true));
}
#[test]
fn integer_trivial_equality_not_unsat() {
let zero_eq = Polynomial::zero(); assert_eq!(
decide_int(
&[(poly(&[(1, 2), (-4, 0)]), Rel::Eq), (zero_eq, Rel::Eq)],
0
),
Some(true)
);
}
}