use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use puremp::{Int, Rational};
use crate::ast::AstId;
fn zero() -> Rational {
Rational::from_integer(Int::from(0))
}
#[derive(Clone, Debug, Default)]
pub struct LinExpr {
coeffs: BTreeMap<AstId, Rational>,
constant: Rational,
}
impl LinExpr {
pub fn new() -> LinExpr {
LinExpr {
coeffs: BTreeMap::new(),
constant: zero(),
}
}
pub fn constant(c: Rational) -> LinExpr {
LinExpr {
coeffs: BTreeMap::new(),
constant: c,
}
}
pub fn var(v: AstId) -> LinExpr {
let mut e = LinExpr::new();
e.coeffs.insert(v, Rational::from_integer(Int::from(1)));
e
}
fn add_scaled(&mut self, other: &LinExpr, k: &Rational) {
for (v, c) in &other.coeffs {
let entry = self.coeffs.entry(*v).or_insert_with(zero);
*entry = &*entry + &(c * k);
if entry.is_zero() {
self.coeffs.remove(v);
}
}
self.constant = &self.constant + &(&other.constant * k);
}
pub fn add(&self, other: &LinExpr) -> LinExpr {
let mut r = self.clone();
r.add_scaled(other, &Rational::from_integer(Int::from(1)));
r
}
pub fn sub(&self, other: &LinExpr) -> LinExpr {
let mut r = self.clone();
r.add_scaled(other, &Rational::from_integer(Int::from(-1)));
r
}
pub fn scale(&self, k: &Rational) -> LinExpr {
let mut r = LinExpr::new();
r.add_scaled(self, k);
r
}
pub fn neg(&self) -> LinExpr {
self.scale(&Rational::from_integer(Int::from(-1)))
}
fn coeff(&self, v: AstId) -> Rational {
self.coeffs.get(&v).cloned().unwrap_or_else(zero)
}
pub fn coeff_of(&self, v: AstId) -> Rational {
self.coeff(v)
}
pub fn is_constant(&self) -> bool {
self.coeffs.is_empty()
}
pub fn as_constant(&self) -> Option<Rational> {
self.is_constant().then(|| self.constant.clone())
}
pub fn vars(&self) -> impl Iterator<Item = AstId> + '_ {
self.coeffs.keys().copied()
}
pub fn terms(&self) -> impl Iterator<Item = (AstId, &Rational)> + '_ {
self.coeffs.iter().map(|(&v, c)| (v, c))
}
pub fn const_term(&self) -> &Rational {
&self.constant
}
pub fn integer_equality_infeasible(&self) -> bool {
if self.coeffs.is_empty() {
return !self.constant.is_zero();
}
let mut l = Int::from(1);
for c in self.coeffs.values() {
l = l.lcm(c.denominator());
}
l = l.lcm(self.constant.denominator());
let scaled = |r: &Rational| -> Int {
let factor = l.div_trunc(r.denominator()); r.numerator() * &factor
};
let mut g = Int::from(0);
for c in self.coeffs.values() {
g = g.gcd(&scaled(c));
}
let b = scaled(&self.constant);
if g.is_zero() {
return !b.is_zero();
}
!b.rem_euclid(&g).is_zero()
}
pub fn integer_strict_tighten(&self) -> LinExpr {
let mut l = Int::from(1);
for c in self.coeffs.values() {
l = l.lcm(c.denominator());
}
l = l.lcm(self.constant.denominator());
let mut e = self.scale(&Rational::from_integer(l));
e.constant = &e.constant + &Rational::from_integer(Int::from(1));
e
}
pub fn integer_gcd_tighten_le(&self) -> Option<LinExpr> {
if self.coeffs.is_empty() {
return None;
}
let mut l = Int::from(1);
for c in self.coeffs.values() {
l = l.lcm(c.denominator());
}
l = l.lcm(self.constant.denominator());
let e = self.scale(&Rational::from_integer(l));
let mut g = Int::from(0);
for c in e.coeffs.values() {
g = g.gcd(c.numerator());
}
if g <= Int::from(1) {
return None; }
let mut out = e.scale(&Rational::from_integer(g).recip());
out.constant = Rational::from_integer(out.constant.ceil());
Some(out)
}
pub fn eval(&self, assign: &Assignment) -> Rational {
let mut acc = self.constant.clone();
for (v, c) in &self.coeffs {
let val = assign.get(v).cloned().unwrap_or_else(zero);
acc = &acc + &(c * &val);
}
acc
}
}
pub fn feasible_with_diseqs(constraints: &[Constraint], diseqs: &[LinExpr]) -> bool {
model_with_diseqs(constraints, diseqs).is_some()
}
pub type Assignment = BTreeMap<AstId, Rational>;
pub enum SolveOutcome {
Sat(Assignment),
Unsat,
Exhausted,
}
pub fn model_with_diseqs_budgeted(
constraints: &[Constraint],
diseqs: &[LinExpr],
budget: &mut u64,
) -> SolveOutcome {
if *budget == 0 {
return SolveOutcome::Exhausted;
}
*budget -= 1;
let model = match model_budgeted(constraints, budget) {
SolveOutcome::Sat(m) => m,
other => return other, };
let Some(d) = diseqs.iter().find(|d| d.eval(&model).is_zero()) else {
return SolveOutcome::Sat(model);
};
let mut lt = constraints.to_vec();
lt.push(Constraint::lt(d.clone()));
match model_with_diseqs_budgeted(<, diseqs, budget) {
SolveOutcome::Unsat => {
let mut gt = constraints.to_vec();
gt.push(Constraint::lt(d.neg())); model_with_diseqs_budgeted(>, diseqs, budget)
}
other => other, }
}
pub fn model_with_diseqs(constraints: &[Constraint], diseqs: &[LinExpr]) -> Option<Assignment> {
match diseqs.split_first() {
None => model(constraints),
Some((d, rest)) => {
let mut lt = constraints.to_vec();
lt.push(Constraint::lt(d.clone()));
if let Some(a) = model_with_diseqs(<, rest) {
return Some(a);
}
let mut gt = constraints.to_vec();
gt.push(Constraint::lt(d.neg())); model_with_diseqs(>, rest)
}
}
}
#[derive(Clone, Debug)]
struct Ineq {
expr: LinExpr,
strict: bool,
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Rel {
Le,
Lt,
Eq,
}
#[derive(Clone, Debug)]
pub struct Constraint {
pub expr: LinExpr,
pub rel: Rel,
}
impl Constraint {
pub fn le(expr: LinExpr) -> Constraint {
Constraint { expr, rel: Rel::Le }
}
pub fn lt(expr: LinExpr) -> Constraint {
Constraint { expr, rel: Rel::Lt }
}
pub fn eq(expr: LinExpr) -> Constraint {
Constraint { expr, rel: Rel::Eq }
}
fn to_ineqs(&self) -> Vec<Ineq> {
match self.rel {
Rel::Le => alloc::vec![Ineq {
expr: self.expr.clone(),
strict: false
}],
Rel::Lt => alloc::vec![Ineq {
expr: self.expr.clone(),
strict: true
}],
Rel::Eq => alloc::vec![
Ineq {
expr: self.expr.clone(),
strict: false
},
Ineq {
expr: self.expr.neg(),
strict: false
},
],
}
}
}
pub fn feasible(constraints: &[Constraint]) -> bool {
model(constraints).is_some()
}
pub fn model(constraints: &[Constraint]) -> Option<Assignment> {
let mut unbounded = u64::MAX;
match model_budgeted(constraints, &mut unbounded) {
SolveOutcome::Sat(a) => Some(a),
_ => None, }
}
pub fn model_budgeted(constraints: &[Constraint], budget: &mut u64) -> SolveOutcome {
let mut ineqs: Vec<Ineq> = constraints.iter().flat_map(|c| c.to_ineqs()).collect();
let mut history: Vec<(AstId, Vec<Ineq>)> = Vec::new();
while let Some(v) = ineqs
.iter()
.find_map(|i| i.expr.coeffs.keys().next().copied())
{
let mut upper = Vec::new(); let mut lower = Vec::new(); let mut rest = Vec::new(); for i in ineqs {
let c = i.expr.coeff(v);
if c.is_zero() {
rest.push(i);
} else if c > zero() {
upper.push(i);
} else {
lower.push(i);
}
}
let mut next = rest;
for u in &upper {
let au = u.expr.coeff(v); for l in &lower {
if *budget == 0 {
return SolveOutcome::Exhausted; }
*budget -= 1;
let al = l.expr.coeff(v); let mut e = u.expr.scale(&al.neg());
e.add_scaled(&l.expr, &au);
next.push(Ineq {
expr: e,
strict: u.strict || l.strict,
});
}
}
let mut mentioning = upper;
mentioning.extend(lower);
history.push((v, mentioning));
ineqs = next;
}
for i in &ineqs {
let k = &i.expr.constant;
let violated = if i.strict { *k >= zero() } else { *k > zero() };
if violated {
return SolveOutcome::Unsat;
}
}
let mut assign: Assignment = BTreeMap::new();
for (v, bounds) in history.into_iter().rev() {
let mut lo: Option<(Rational, bool)> = None; let mut hi: Option<(Rational, bool)> = None;
for i in &bounds {
let cv = i.expr.coeff(v); let mut r = i.expr.constant.clone();
for (u, c) in &i.expr.coeffs {
if *u != v {
let uv = assign.get(u).cloned().unwrap_or_else(zero);
r = &r + &(c * &uv);
}
}
let bound = &r.neg() / &cv;
if cv > zero() {
hi = Some(match hi {
Some((h, hs)) if h < bound || (h == bound && !hs) => (h, hs),
_ => (bound, i.strict),
});
} else {
lo = Some(match lo {
Some((l, ls)) if l > bound || (l == bound && !ls) => (l, ls),
_ => (bound, i.strict),
});
}
}
assign.insert(v, pick_between(lo, hi));
}
SolveOutcome::Sat(assign)
}
pub fn project(constraints: &[Constraint], x: AstId, budget: &mut u64) -> Option<Vec<Constraint>> {
let ineqs: Vec<Ineq> = constraints.iter().flat_map(|c| c.to_ineqs()).collect();
let mut upper = Vec::new(); let mut lower = Vec::new(); let mut rest = Vec::new(); for i in ineqs {
let c = i.expr.coeff(x);
if c.is_zero() {
rest.push(i);
} else if c > zero() {
upper.push(i);
} else {
lower.push(i);
}
}
for u in &upper {
let au = u.expr.coeff(x);
for l in &lower {
if *budget == 0 {
return None;
}
*budget -= 1;
let al = l.expr.coeff(x);
let mut e = u.expr.scale(&al.neg());
e.add_scaled(&l.expr, &au);
rest.push(Ineq {
expr: e,
strict: u.strict || l.strict,
});
}
}
Some(
rest.into_iter()
.map(|i| {
if i.strict {
Constraint::lt(i.expr)
} else {
Constraint::le(i.expr)
}
})
.collect(),
)
}
#[derive(Clone, Debug, PartialEq)]
pub enum OptOutcome {
Attained(Rational),
Bound(Rational),
Unbounded,
Exhausted,
}
pub fn optimize(
constraints: &[Constraint],
obj: &LinExpr,
z: AstId,
maximize: bool,
budget: &mut u64,
) -> OptOutcome {
let mut cons = constraints.to_vec();
cons.push(Constraint::eq(LinExpr::var(z).sub(obj)));
let mut ineqs: Vec<Ineq> = cons.iter().flat_map(|c| c.to_ineqs()).collect();
while let Some(v) = ineqs
.iter()
.find_map(|i| i.expr.coeffs.keys().find(|&&k| k != z).copied())
{
let mut upper = Vec::new();
let mut lower = Vec::new();
let mut next = Vec::new();
for i in ineqs {
let c = i.expr.coeff(v);
if c.is_zero() {
next.push(i);
} else if c > zero() {
upper.push(i);
} else {
lower.push(i);
}
}
for u in &upper {
let au = u.expr.coeff(v);
for l in &lower {
if *budget == 0 {
return OptOutcome::Exhausted;
}
*budget -= 1;
let al = l.expr.coeff(v);
let mut e = u.expr.scale(&al.neg());
e.add_scaled(&l.expr, &au);
next.push(Ineq {
expr: e,
strict: u.strict || l.strict,
});
}
}
ineqs = next;
}
let mut upper: Option<(Rational, bool)> = None; let mut lower: Option<(Rational, bool)> = None; for i in &ineqs {
let cz = i.expr.coeff(z);
if cz.is_zero() {
continue; }
let bound = &i.expr.constant.neg() / &cz;
if cz > zero() {
upper = Some(match upper {
Some((h, hs)) if h < bound || (h == bound && !hs) => (h, hs),
_ => (bound, i.strict),
});
} else {
lower = Some(match lower {
Some((l, ls)) if l > bound || (l == bound && !ls) => (l, ls),
_ => (bound, i.strict),
});
}
}
let chosen = if maximize { upper } else { lower };
match chosen {
None => OptOutcome::Unbounded,
Some((b, true)) => OptOutcome::Bound(b),
Some((b, false)) => OptOutcome::Attained(b),
}
}
fn pick_between(lo: Option<(Rational, bool)>, hi: Option<(Rational, bool)>) -> Rational {
let one = Rational::from_integer(Int::from(1));
let two = Rational::from_integer(Int::from(2));
match (lo, hi) {
(None, None) => zero(),
(Some((l, _)), None) => &l + &one,
(None, Some((h, _))) => &h - &one,
(Some((l, ls)), Some((h, hs))) => {
if l == h {
debug_assert!(!ls && !hs, "empty interval in a feasible system");
l
} else {
&(&l + &h) / &two }
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ast::AstId;
fn rat(n: i64) -> Rational {
Rational::from_integer(Int::from(n))
}
fn x() -> AstId {
AstId(1000)
}
fn y() -> AstId {
AstId(1001)
}
#[test]
fn contradictory_bounds_infeasible() {
let ge0 = Constraint::le(LinExpr::var(x()).neg());
let le_m1 = Constraint::le(LinExpr::var(x()).add(&LinExpr::constant(rat(1))));
assert!(!feasible(&[ge0, le_m1]));
}
#[test]
fn satisfiable_system() {
let x_ge0 = Constraint::le(LinExpr::var(x()).neg());
let y_ge0 = Constraint::le(LinExpr::var(y()).neg());
let sum = LinExpr::var(x())
.add(&LinExpr::var(y()))
.sub(&LinExpr::constant(rat(1)));
let sum_le1 = Constraint::le(sum);
assert!(feasible(&[x_ge0, y_ge0, sum_le1]));
}
#[test]
fn strict_cycle_infeasible() {
let xy = LinExpr::var(x()).sub(&LinExpr::var(y()));
let yx = LinExpr::var(y()).sub(&LinExpr::var(x()));
assert!(!feasible(&[Constraint::lt(xy), Constraint::lt(yx)]));
}
#[test]
fn equality_pins_a_value() {
let x_eq2 = Constraint::eq(LinExpr::var(x()).sub(&LinExpr::constant(rat(2))));
let x_le1 = Constraint::le(LinExpr::var(x()).sub(&LinExpr::constant(rat(1))));
assert!(!feasible(&[x_eq2.clone(), x_le1]));
let x_ge1 = Constraint::le(LinExpr::constant(rat(1)).sub(&LinExpr::var(x())));
assert!(feasible(&[x_eq2, x_ge1]));
}
fn eval(expr: &LinExpr, assign: &Assignment) -> Rational {
let mut acc = expr.constant.clone();
for (v, c) in &expr.coeffs {
let val = assign.get(v).cloned().unwrap_or_else(zero);
acc = &acc + &(c * &val);
}
acc
}
fn check_sat(cs: &[Constraint], assign: &Assignment) {
for c in cs {
let v = eval(&c.expr, assign);
let ok = match c.rel {
Rel::Le => v <= zero(),
Rel::Lt => v < zero(),
Rel::Eq => v == zero(),
};
assert!(ok, "constraint {:?} violated: lhs = {v:?}", c.rel);
}
}
#[test]
fn model_satisfies_bounded_system() {
let x_ge0 = Constraint::le(LinExpr::var(x()).neg());
let y_ge0 = Constraint::le(LinExpr::var(y()).neg());
let sum = LinExpr::var(x())
.add(&LinExpr::var(y()))
.sub(&LinExpr::constant(rat(1)));
let sum_le1 = Constraint::le(sum);
let eq = Constraint::eq(LinExpr::var(x()).sub(&LinExpr::var(y())));
let cs = [x_ge0, y_ge0, sum_le1, eq];
let m = model(&cs).expect("feasible");
check_sat(&cs, &m);
}
#[test]
fn model_respects_strict_bounds() {
let gt0 = Constraint::lt(LinExpr::var(x()).neg()); let lt1 = Constraint::lt(LinExpr::var(x()).sub(&LinExpr::constant(rat(1))));
let cs = [gt0, lt1];
let m = model(&cs).expect("feasible");
check_sat(&cs, &m);
let xv = m.get(&x()).cloned().unwrap_or_else(zero);
assert!(xv > zero() && xv < rat(1));
}
#[test]
fn model_with_disequality() {
let le5 = Constraint::le(LinExpr::var(x()).sub(&LinExpr::constant(rat(5))));
let ge5 = Constraint::le(LinExpr::constant(rat(5)).sub(&LinExpr::var(x())));
let ne5 = LinExpr::var(x()).sub(&LinExpr::constant(rat(5)));
assert!(model_with_diseqs(&[le5, ge5.clone()], core::slice::from_ref(&ne5)).is_none());
let le6 = Constraint::le(LinExpr::var(x()).sub(&LinExpr::constant(rat(6))));
let cs = [le6, ge5];
let m = model_with_diseqs(&cs, &[ne5]).expect("feasible");
check_sat(&cs, &m);
assert_ne!(m.get(&x()).cloned().unwrap_or_else(zero), rat(5));
}
#[test]
fn non_strict_boundary_is_feasible() {
let le0 = Constraint::le(LinExpr::var(x()));
let ge0 = Constraint::le(LinExpr::var(x()).neg());
assert!(feasible(&[le0.clone(), ge0.clone()]));
let lt0 = Constraint::lt(LinExpr::var(x()));
assert!(!feasible(&[lt0, ge0]));
}
}