oximo_expr/

classify.rs

use crate::arena::{ArenaAccess, ExprArena, ExprId, ExprNode, UnaryOp};

/// Highest-degree polynomial class an expression belongs to, ignoring constant
/// folding. Used by backends to pick between linear, quadratic, and general
/// nonlinear translation paths.
///
/// Variants are ordered by increasing degree, so `max` of two classes yields the
/// dominating one (e.g. a model with a quadratic objective and a nonlinear
/// constraint is `Nonlinear`).
#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum ExprClass {
    Linear,
    Quadratic,
    Nonlinear,
}

/// Polynomial-degree bucket. `Higher` is a saturating sentinel for "anything
/// above quadratic". Both polynomial degree > 2 and transcendentals collapse
/// into it, since neither fits a QP solver's quadratic API.
#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
enum Degree {
    Zero,
    One,
    Two,
    Higher,
}

impl Degree {
    /// `+` on a sum: take the maximum, saturating at `Higher`.
    fn add(self, other: Degree) -> Degree {
        self.max(other)
    }

    /// `*` on a product: add ordinal degrees, saturating at `Higher`.
    fn mul(self, other: Degree) -> Degree {
        match (self, other) {
            (Degree::Higher, _) | (_, Degree::Higher) => Degree::Higher,
            (Degree::Zero, x) | (x, Degree::Zero) => x,
            (Degree::One, Degree::One) => Degree::Two,
            _ => Degree::Higher,
        }
    }

    /// `^n` on a power: multiply by `n`, saturating at `Higher`.
    fn pow(self, n: u32) -> Degree {
        match (self, n) {
            (_, 0) | (Degree::Zero, _) => Degree::Zero,
            (d, 1) => d,
            (Degree::One, 2) => Degree::Two,
            _ => Degree::Higher,
        }
    }
}

#[cfg(test)]
fn recursive_degree(arena: &(impl ArenaAccess + ?Sized), id: ExprId) -> Degree {
    match arena.get(id) {
        ExprNode::Const(_) | ExprNode::Param(_) => Degree::Zero,
        ExprNode::Var(_) | ExprNode::Linear { .. } => Degree::One,
        ExprNode::Unary(UnaryOp::Neg, inner) => recursive_degree(arena, *inner),
        ExprNode::Add(children) => {
            let mut d = Degree::Zero;
            for c in children {
                d = d.add(recursive_degree(arena, *c));
                if d == Degree::Higher {
                    return d;
                }
            }
            d
        }
        ExprNode::Mul(children) => {
            let mut d = Degree::Zero;
            for c in children {
                d = d.mul(recursive_degree(arena, *c));
                if d == Degree::Higher {
                    return d;
                }
            }
            d
        }
        ExprNode::Pow(base, exp) => {
            let ExprNode::Const(e) = arena.get(*exp) else { return Degree::Higher };
            if (*e - e.round()).abs() >= f64::EPSILON || *e < 0.0 {
                return Degree::Higher;
            }
            // Bucket the exponent into the only values `Degree::pow` treats
            // distinctly.
            let n = match e.round() {
                v if v < 0.5 => 0,
                v if v < 1.5 => 1,
                v if v < 2.5 => 2,
                _ => 3,
            };
            recursive_degree(arena, *base).pow(n)
        }
        // Transcendentals are always > quadratic. Division is too: `div_into`
        // folds the only degree-preserving case (constant denominator) before a
        // `Div` node is created, so any other `Div` has a non-constant
        // denominator.
        ExprNode::Div(_, _)
        | ExprNode::Atan2(_, _)
        | ExprNode::Min(_)
        | ExprNode::Max(_)
        | ExprNode::Unary(_, _) => Degree::Higher,
    }
}

struct DegreeFolder<'a, A: ?Sized>(&'a A);

enum DegreeOp {
    Add,
    Mul,
    Pow(u32),
}
struct DegreeState<'a> {
    rest: &'a [ExprId],
    value: Degree,
    op: DegreeOp,
}

impl<'a, A: ArenaAccess + ?Sized> crate::fold::Folder for DegreeFolder<'a, A> {
    type Value = Degree;
    type State = DegreeState<'a>;

    fn start(&self, id: ExprId) -> std::ops::ControlFlow<Option<Degree>, Self::State> {
        use std::ops::ControlFlow::{Break, Continue};
        let (rest, op) = match self.0.get(id) {
            ExprNode::Const(_) | ExprNode::Param(_) => return Break(Some(Degree::Zero)),
            ExprNode::Var(_) | ExprNode::Linear { .. } => return Break(Some(Degree::One)),
            ExprNode::Add(c) => (c.as_slice(), DegreeOp::Add),
            ExprNode::Mul(c) => (c.as_slice(), DegreeOp::Mul),
            ExprNode::Unary(UnaryOp::Neg, c) => (std::slice::from_ref(c), DegreeOp::Add),
            ExprNode::Pow(base, exp) => {
                let ExprNode::Const(e) = self.0.get(*exp) else {
                    return Break(Some(Degree::Higher));
                };
                if (*e - e.round()).abs() >= f64::EPSILON || *e < 0.0 {
                    return Break(Some(Degree::Higher));
                }
                let n = match e.round() {
                    v if v < 0.5 => 0,
                    v if v < 1.5 => 1,
                    v if v < 2.5 => 2,
                    _ => 3,
                };
                (std::slice::from_ref(base), DegreeOp::Pow(n))
            }
            _ => return Break(Some(Degree::Higher)),
        };
        Continue(DegreeState { rest, value: Degree::Zero, op })
    }

    fn next(&self, state: &mut Self::State) -> std::ops::ControlFlow<Option<Degree>, ExprId> {
        use std::ops::ControlFlow::{Break, Continue};
        while state.value != Degree::Higher {
            let Some((&child, tail)) = state.rest.split_first() else { break };
            state.rest = tail;
            match self.start(child) {
                Break(Some(value)) => self.accept(state, value),
                _ => return Continue(child),
            }
        }
        Break(Some(state.value))
    }

    fn accept(&self, state: &mut Self::State, value: Degree) {
        state.value = match state.op {
            DegreeOp::Add => state.value.add(value),
            DegreeOp::Mul => state.value.mul(value),
            DegreeOp::Pow(n) => value.pow(n),
        };
    }
}

fn degree(arena: &(impl ArenaAccess + ?Sized), mut id: ExprId) -> Degree {
    while let ExprNode::Unary(UnaryOp::Neg, child) = arena.get(id) {
        id = *child;
    }
    crate::fold::fold(arena, id, DegreeFolder(arena)).expect("degree classification is total")
}

/// Classify an expression as Linear, Quadratic (polynomial degree <= 2 with at
/// least one degree-2 term), or Nonlinear (transcendentals, non-integer powers,
/// or polynomial degree > 2).
pub fn classify(arena: &ExprArena, id: ExprId) -> ExprClass {
    classify_access(arena, id)
}

pub(crate) fn classify_access(arena: &(impl ArenaAccess + ?Sized), id: ExprId) -> ExprClass {
    match degree(arena, id) {
        Degree::Zero | Degree::One => ExprClass::Linear,
        Degree::Two => ExprClass::Quadratic,
        Degree::Higher => ExprClass::Nonlinear,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn iterative_degree_matches_recursive_classification() {
        let (arena, ids) = crate::fold::test_arena();
        for id in ids {
            assert_eq!(degree(&arena, id), recursive_degree(&arena, id), "node {id:?}");
        }
    }
    use crate::arena::{ExprArena, ExprNode, VarId};
    use smallvec::smallvec;

    fn var(arena: &mut ExprArena, i: u32) -> ExprId {
        arena.push(ExprNode::Var(VarId(i)))
    }

    #[test]
    fn linear_var_sum() {
        let mut a = ExprArena::new();
        let x = var(&mut a, 0);
        let y = var(&mut a, 1);
        let sum = a.push(ExprNode::Add(smallvec![x, y]));
        assert_eq!(classify(&a, sum), ExprClass::Linear);
    }

    #[test]
    fn quadratic_mul_two_vars() {
        let mut a = ExprArena::new();
        let x = var(&mut a, 0);
        let y = var(&mut a, 1);
        let xy = a.push(ExprNode::Mul(smallvec![x, y]));
        assert_eq!(classify(&a, xy), ExprClass::Quadratic);
    }

    #[test]
    fn quadratic_pow_two() {
        let mut a = ExprArena::new();
        let x = var(&mut a, 0);
        let two = a.push(ExprNode::Const(2.0));
        let sq = a.push(ExprNode::Pow(x, two));
        assert_eq!(classify(&a, sq), ExprClass::Quadratic);
    }

    #[test]
    fn nonlinear_pow_three() {
        let mut a = ExprArena::new();
        let x = var(&mut a, 0);
        let three = a.push(ExprNode::Const(3.0));
        let cube = a.push(ExprNode::Pow(x, three));
        assert_eq!(classify(&a, cube), ExprClass::Nonlinear);
    }

    #[test]
    fn nonlinear_div() {
        let mut a = ExprArena::new();
        let x = var(&mut a, 0);
        let y = var(&mut a, 1);
        let q = a.push(ExprNode::Div(x, y));
        assert_eq!(classify(&a, q), ExprClass::Nonlinear);
    }

    #[test]
    fn nonlinear_sin() {
        let mut a = ExprArena::new();
        let x = var(&mut a, 0);
        let s = a.push(ExprNode::Unary(UnaryOp::Sin, x));
        assert_eq!(classify(&a, s), ExprClass::Nonlinear);
    }

    #[test]
    fn nonlinear_abs() {
        let mut a = ExprArena::new();
        let x = var(&mut a, 0);
        let s = a.push(ExprNode::Unary(UnaryOp::Abs, x));
        assert_eq!(classify(&a, s), ExprClass::Nonlinear);
    }

    #[test]
    fn nonlinear_triple_mul() {
        let mut arena = ExprArena::new();
        let x = var(&mut arena, 0);
        let y = var(&mut arena, 1);
        let z = var(&mut arena, 2);
        let prod = arena.push(ExprNode::Mul(smallvec![x, y, z]));
        assert_eq!(classify(&arena, prod), ExprClass::Nonlinear);
    }

    #[test]
    fn linear_promoted_by_const_mul() {
        let mut a = ExprArena::new();
        let x = var(&mut a, 0);
        let c = a.push(ExprNode::Const(3.0));
        let m = a.push(ExprNode::Mul(smallvec![c, x]));
        assert_eq!(classify(&a, m), ExprClass::Linear);
    }

    #[test]
    fn param_alone_is_linear() {
        let mut a = ExprArena::new();
        let p = a.new_param(4.0);
        let pn = a.param(p);
        assert_eq!(classify(&a, pn), ExprClass::Linear);
    }

    #[test]
    fn param_times_var_is_linear() {
        let mut a = ExprArena::new();
        let p = a.new_param(4.0);
        let pn = a.param(p);
        let x = var(&mut a, 0);
        let m = a.push(ExprNode::Mul(smallvec![pn, x]));
        assert_eq!(classify(&a, m), ExprClass::Linear);
    }

    #[test]
    fn param_times_var_squared_is_quadratic() {
        let mut a = ExprArena::new();
        let p = a.new_param(4.0);
        let pn = a.param(p);
        let x = var(&mut a, 0);
        let two = a.push(ExprNode::Const(2.0));
        let sq = a.push(ExprNode::Pow(x, two));
        let m = a.push(ExprNode::Mul(smallvec![pn, sq]));
        assert_eq!(classify(&a, m), ExprClass::Quadratic);
    }
}