use super::idmap::IdSet;
use crate::deriv::SideCondition;
use crate::errors::AlkahestError;
use crate::kernel::expr::PredicateKind;
use crate::kernel::{Domain, ExprData, ExprId, ExprPool};
use crate::logic::{satisfiable, Satisfiability};
use crate::simplify::{rules_for_config, simplify_with, SimplifyConfig};
use crate::DerivedExpr;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AssumptionError {
Contradiction,
}
impl std::fmt::Display for AssumptionError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Contradiction => write!(f, "assumption contradicts the current context"),
}
}
}
impl std::error::Error for AssumptionError {}
impl AlkahestError for AssumptionError {
fn code(&self) -> &'static str {
"E-SIMPLIFY-001"
}
fn remediation(&self) -> Option<&'static str> {
Some("remove the conflicting refinement or create a separate AssumptionContext")
}
}
#[derive(Debug, Clone, Default)]
pub struct AssumptionContext {
predicates: Vec<ExprId>,
facts: Vec<SideCondition>,
}
impl AssumptionContext {
pub fn new() -> Self {
Self::default()
}
pub fn predicates(&self) -> &[ExprId] {
&self.predicates
}
pub fn facts(&self) -> &[SideCondition] {
&self.facts
}
pub fn refine(&mut self, predicate: ExprId, pool: &ExprPool) -> Result<(), AssumptionError> {
let mut predicates = self.predicates.clone();
predicates.push(predicate);
let conjunction = if predicates.len() == 1 {
predicate
} else {
pool.predicate(PredicateKind::And, predicates.clone())
};
if matches!(satisfiable(conjunction, pool), Satisfiability::Unsat) {
return Err(AssumptionError::Contradiction);
}
self.predicates = predicates;
for fact in normalize_predicate(predicate, pool) {
push_unique(&mut self.facts, fact);
}
Ok(())
}
pub fn simplify(&self, expr: ExprId, pool: &ExprPool) -> DerivedExpr<ExprId> {
let mut facts = self.facts.clone();
collect_static_domain_facts(expr, pool, &mut facts);
let config = SimplifyConfig {
assumptions: facts,
..SimplifyConfig::default()
};
let rules = rules_for_config(&config);
simplify_with(expr, pool, &rules, config)
}
}
pub fn simplify_with_assumptions(
expr: ExprId,
pool: &ExprPool,
assumptions: &AssumptionContext,
) -> DerivedExpr<ExprId> {
assumptions.simplify(expr, pool)
}
use std::cell::RefCell;
use std::marker::PhantomData;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Sign {
Negative,
Zero,
Positive,
}
impl Sign {
pub fn to_i8(self) -> i8 {
match self {
Sign::Negative => -1,
Sign::Zero => 0,
Sign::Positive => 1,
}
}
fn of_ordering(ord: std::cmp::Ordering) -> Sign {
match ord {
std::cmp::Ordering::Less => Sign::Negative,
std::cmp::Ordering::Equal => Sign::Zero,
std::cmp::Ordering::Greater => Sign::Positive,
}
}
fn negate(self) -> Sign {
match self {
Sign::Negative => Sign::Positive,
Sign::Zero => Sign::Zero,
Sign::Positive => Sign::Negative,
}
}
}
struct AmbientFrame {
pool: usize,
predicates: Vec<ExprId>,
}
thread_local! {
static AMBIENT: RefCell<Vec<AmbientFrame>> = const { RefCell::new(Vec::new()) };
}
pub struct AssumptionScope {
_not_send: PhantomData<*const ()>,
}
impl Drop for AssumptionScope {
fn drop(&mut self) {
AMBIENT.with(|s| {
s.borrow_mut().pop();
});
}
}
pub fn enter_assumptions(assumptions: &AssumptionContext, pool: &ExprPool) -> AssumptionScope {
AMBIENT.with(|s| {
s.borrow_mut().push(AmbientFrame {
pool: pool as *const ExprPool as usize,
predicates: assumptions.predicates().to_vec(),
})
});
AssumptionScope {
_not_send: PhantomData,
}
}
pub fn ambient_assumptions_active(pool: &ExprPool) -> bool {
let addr = pool as *const ExprPool as usize;
AMBIENT.with(|s| s.borrow().iter().any(|f| f.pool == addr))
}
pub fn ambient_equalities(pool: &ExprPool) -> Vec<(ExprId, ExprId)> {
let addr = pool as *const ExprPool as usize;
let mut out: Vec<(ExprId, ExprId)> = Vec::new();
AMBIENT.with(|s| {
for frame in s.borrow().iter().filter(|f| f.pool == addr) {
for &p in &frame.predicates {
collect_equalities(p, pool, &mut out);
}
}
});
out
}
fn collect_equalities(predicate: ExprId, pool: &ExprPool, out: &mut Vec<(ExprId, ExprId)>) {
let ExprData::Predicate { kind, args } = pool.get(predicate) else {
return;
};
match kind {
PredicateKind::And => {
for &part in &args {
collect_equalities(part, pool, out);
}
}
PredicateKind::Eq if args.len() == 2 => {
let is_sym = |e: ExprId| matches!(pool.get(e), ExprData::Symbol { .. });
let is_num = |e: ExprId| {
matches!(
pool.get(e),
ExprData::Integer(_) | ExprData::Rational(_) | ExprData::Float(_)
)
};
if is_sym(args[0]) && is_num(args[1]) {
out.push((args[0], args[1]));
} else if is_num(args[0]) && is_sym(args[1]) {
out.push((args[1], args[0]));
}
}
_ => {}
}
}
pub fn assumed_sign(expr: ExprId, pool: &ExprPool) -> Option<Sign> {
assumed_sign_at(expr, pool, 0)
}
const MAX_SIGN_DEPTH: u32 = 24;
fn assumed_sign_at(expr: ExprId, pool: &ExprPool, depth: u32) -> Option<Sign> {
if depth > MAX_SIGN_DEPTH {
return None;
}
if let Some(s) = stated_sign(expr, pool) {
return Some(s);
}
match pool.get(expr) {
ExprData::Integer(n) => Some(Sign::of_ordering(n.0.cmp(&rug::Integer::ZERO))),
ExprData::Rational(r) => Some(Sign::of_ordering(r.0.cmp0())),
ExprData::Float(f) => {
if !f.inner.is_finite() {
return None;
}
Some(Sign::of_ordering(f.inner.cmp0()?))
}
ExprData::Symbol {
domain: Domain::Positive,
..
} => Some(Sign::Positive),
ExprData::Symbol { .. } => None,
ExprData::Mul(xs) => {
let mut sign = Sign::Positive;
for x in xs {
match assumed_sign_at(x, pool, depth + 1)? {
Sign::Zero => return Some(Sign::Zero),
Sign::Negative => sign = sign.negate(),
Sign::Positive => {}
}
}
Some(sign)
}
ExprData::Add(xs) => {
let mut sign: Option<Sign> = None;
for x in xs {
let s = assumed_sign_at(x, pool, depth + 1)?;
match (sign, s) {
(_, Sign::Zero) => {}
(None, s) => sign = Some(s),
(Some(a), b) if a == b => {}
_ => return None,
}
}
sign.or(Some(Sign::Zero))
}
ExprData::Pow { base, exp } => {
let base_sign = assumed_sign_at(base, pool, depth + 1);
if let ExprData::Integer(n) = pool.get(exp) {
if n.0.is_even() {
return match base_sign {
Some(Sign::Zero) => Some(Sign::Zero),
Some(_) => Some(Sign::Positive),
None => None,
};
}
return base_sign;
}
match base_sign {
Some(Sign::Positive) => Some(Sign::Positive),
_ => None,
}
}
ExprData::Func { name, args } => match name.as_str() {
"exp" | "cosh" => Some(Sign::Positive),
"abs" if args.len() == 1 => match assumed_sign_at(args[0], pool, depth + 1) {
Some(Sign::Zero) => Some(Sign::Zero),
Some(_) => Some(Sign::Positive),
None => None,
},
"sqrt" | "cbrt" if args.len() == 1 => assumed_sign_at(args[0], pool, depth + 1),
_ => None,
},
_ => None,
}
}
fn stated_sign(expr: ExprId, pool: &ExprPool) -> Option<Sign> {
if matches!(
pool.get(expr),
ExprData::Symbol {
domain: Domain::Positive,
..
}
) {
return Some(Sign::Positive);
}
let addr = pool as *const ExprPool as usize;
AMBIENT.with(|s| {
for frame in s.borrow().iter().rev().filter(|f| f.pool == addr) {
for &p in &frame.predicates {
if let Some(sign) = sign_from_predicate(p, expr, pool) {
return Some(sign);
}
}
}
None
})
}
fn sign_from_predicate(predicate: ExprId, subject: ExprId, pool: &ExprPool) -> Option<Sign> {
let ExprData::Predicate { kind, args } = pool.get(predicate) else {
return None;
};
if matches!(kind, PredicateKind::And) {
return args
.iter()
.find_map(|&part| sign_from_predicate(part, subject, pool));
}
if args.len() != 2 {
return None;
}
let (lhs, rhs) = (args[0], args[1]);
let kind = if lhs == subject && is_zero(rhs, pool) {
kind
} else if rhs == subject && is_zero(lhs, pool) {
match kind {
PredicateKind::Lt => PredicateKind::Gt,
PredicateKind::Le => PredicateKind::Ge,
PredicateKind::Gt => PredicateKind::Lt,
PredicateKind::Ge => PredicateKind::Le,
other => other,
}
} else {
return None;
};
match kind {
PredicateKind::Gt => Some(Sign::Positive),
PredicateKind::Lt => Some(Sign::Negative),
PredicateKind::Eq => Some(Sign::Zero),
_ => None,
}
}
fn normalize_predicate(predicate: ExprId, pool: &ExprPool) -> Vec<SideCondition> {
let ExprData::Predicate { kind, args } = pool.get(predicate) else {
return vec![];
};
match kind {
PredicateKind::And => args
.iter()
.flat_map(|&part| normalize_predicate(part, pool))
.collect(),
PredicateKind::Gt if args.len() == 2 && is_zero(args[1], pool) => {
vec![SideCondition::Positive(args[0])]
}
PredicateKind::Lt if args.len() == 2 && is_zero(args[0], pool) => {
vec![SideCondition::Positive(args[1])]
}
PredicateKind::Ne if args.len() == 2 && is_zero(args[1], pool) => {
vec![SideCondition::NonZero(args[0])]
}
PredicateKind::Ne if args.len() == 2 && is_zero(args[0], pool) => {
vec![SideCondition::NonZero(args[1])]
}
_ => vec![],
}
}
fn is_zero(expr: ExprId, pool: &ExprPool) -> bool {
match pool.get(expr) {
ExprData::Integer(value) => value.0 == 0,
ExprData::Rational(value) => value.0 == 0,
_ => false,
}
}
fn push_unique(facts: &mut Vec<SideCondition>, fact: SideCondition) {
if !facts.contains(&fact) {
facts.push(fact);
}
}
pub(crate) fn collect_static_domain_facts(
expr: ExprId,
pool: &ExprPool,
facts: &mut Vec<SideCondition>,
) {
let mut pending: Vec<ExprId> = Vec::new();
let mut expanded = IdSet::default();
let mut current = expr;
loop {
pool.with(current, |data| match data {
_ if !pool.has_static_domain_fact(current) => {}
ExprData::Integer(_) | ExprData::Rational(_) | ExprData::Float(_) => {}
ExprData::Symbol { domain, .. } => match domain {
Domain::Positive => {
push_unique(facts, SideCondition::Positive(current));
push_unique(facts, SideCondition::NonZero(current));
}
Domain::NonZero => push_unique(facts, SideCondition::NonZero(current)),
_ => {}
},
_ if !expanded.insert(current) => {}
ExprData::Add(args)
| ExprData::Mul(args)
| ExprData::Func { args, .. }
| ExprData::Predicate { args, .. } => pending.extend(args.iter().rev()),
ExprData::Pow { base, exp } => pending.extend([*exp, *base]),
ExprData::Piecewise { branches, default } => {
pending.push(*default);
for (condition, value) in branches.iter().rev() {
pending.extend([*value, *condition]);
}
}
ExprData::Forall { var, body } | ExprData::Exists { var, body } => {
pending.extend([*body, *var]);
}
ExprData::BigO(arg) => pending.push(*arg),
ExprData::RootSum { poly, var, body } => {
pending.extend([*body, *var, *poly]);
}
});
match pending.pop() {
Some(next) => current = next,
None => return,
}
}
}
#[cfg(test)]
pub(crate) fn collect_static_domain_facts_reference(
expr: ExprId,
pool: &ExprPool,
facts: &mut Vec<SideCondition>,
) {
let mut pending: Vec<ExprId> = Vec::new();
let mut current = expr;
loop {
pool.with(current, |data| match data {
ExprData::Symbol { domain, .. } => match domain {
Domain::Positive => {
push_unique(facts, SideCondition::Positive(current));
push_unique(facts, SideCondition::NonZero(current));
}
Domain::NonZero => push_unique(facts, SideCondition::NonZero(current)),
_ => {}
},
ExprData::Add(args)
| ExprData::Mul(args)
| ExprData::Func { args, .. }
| ExprData::Predicate { args, .. } => pending.extend(args.iter().rev()),
ExprData::Pow { base, exp } => pending.extend([*exp, *base]),
ExprData::Piecewise { branches, default } => {
pending.push(*default);
for (condition, value) in branches.iter().rev() {
pending.extend([*value, *condition]);
}
}
ExprData::Forall { var, body } | ExprData::Exists { var, body } => {
pending.extend([*body, *var]);
}
ExprData::BigO(arg) => pending.push(*arg),
ExprData::RootSum { poly, var, body } => {
pending.extend([*body, *var, *poly]);
}
ExprData::Integer(_) | ExprData::Rational(_) | ExprData::Float(_) => {}
});
match pending.pop() {
Some(next) => current = next,
None => return,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn positive_refinement_enables_conditional_rewrites() {
let pool = ExprPool::new();
let x = pool.symbol("x", Domain::Real);
let zero = pool.integer(0_i32);
let mut assumptions = AssumptionContext::new();
assumptions
.refine(pool.predicate(PredicateKind::Gt, vec![x, zero]), &pool)
.unwrap();
let squared = pool.pow(x, pool.integer(2_i32));
assert_eq!(
assumptions
.simplify(pool.func("sqrt", vec![squared]), &pool)
.value,
x
);
assert_eq!(
assumptions
.simplify(pool.func("exp", vec![pool.func("log", vec![x])]), &pool)
.value,
x
);
}
#[test]
fn nonzero_refinement_enables_cancellation() {
let pool = ExprPool::new();
let x = pool.symbol("x", Domain::Real);
let zero = pool.integer(0_i32);
let mut assumptions = AssumptionContext::new();
assumptions
.refine(pool.predicate(PredicateKind::Ne, vec![x, zero]), &pool)
.unwrap();
assert_eq!(
assumptions.simplify(pool.pow(x, zero), &pool).value,
pool.integer(1_i32)
);
let inverse = pool.pow(x, pool.integer(-1_i32));
assert_eq!(
assumptions
.simplify(pool.mul(vec![x, inverse]), &pool)
.value,
pool.integer(1_i32)
);
}
#[test]
fn contradiction_is_rejected_without_mutating_context() {
let pool = ExprPool::new();
let x = pool.symbol("x", Domain::Real);
let zero = pool.integer(0_i32);
let positive = pool.predicate(PredicateKind::Gt, vec![x, zero]);
let nonpositive = pool.predicate(PredicateKind::Le, vec![x, zero]);
let mut assumptions = AssumptionContext::new();
assumptions.refine(positive, &pool).unwrap();
assert_eq!(
assumptions.refine(nonpositive, &pool),
Err(AssumptionError::Contradiction)
);
assert_eq!(assumptions.predicates(), &[positive]);
}
#[test]
fn positive_refinement_enables_abs_of_positive() {
let pool = ExprPool::new();
let x = pool.symbol("x", Domain::Real);
let zero = pool.integer(0_i32);
let mut assumptions = AssumptionContext::new();
assumptions
.refine(pool.predicate(PredicateKind::Gt, vec![x, zero]), &pool)
.unwrap();
assert_eq!(
assumptions.simplify(pool.func("abs", vec![x]), &pool).value,
x
);
}
#[test]
fn abs_without_positive_fact_unchanged() {
let pool = ExprPool::new();
let x = pool.symbol("x", Domain::Real);
let expr = pool.func("abs", vec![x]);
assert_eq!(AssumptionContext::new().simplify(expr, &pool).value, expr);
}
#[test]
fn static_positive_domain_is_available_without_refinement() {
let pool = ExprPool::new();
let x = pool.symbol("x", Domain::Positive);
let squared = pool.pow(x, pool.integer(2_i32));
assert_eq!(
AssumptionContext::new()
.simplify(pool.func("sqrt", vec![squared]), &pool)
.value,
x
);
}
fn chebyshev(pool: &ExprPool, y: ExprId, n: usize) -> ExprId {
let two = pool.integer(2_i32);
let minus_one = pool.integer(-1_i32);
let (mut a, mut b) = (pool.integer(1_i32), y);
for _ in 1..n {
let c = pool.add(vec![
pool.mul(vec![two, y, b]),
pool.mul(vec![minus_one, a]),
]);
a = b;
b = c;
}
b
}
#[test]
fn static_domain_facts_on_a_shared_dag_are_linear() {
let pool = ExprPool::new();
let y = pool.symbol("y", Domain::Positive);
let e = chebyshev(&pool, y, 60);
let start = std::time::Instant::now();
let mut facts = Vec::new();
collect_static_domain_facts(e, &pool, &mut facts);
assert!(start.elapsed() < std::time::Duration::from_secs(2));
assert_eq!(
facts,
vec![SideCondition::Positive(y), SideCondition::NonZero(y)]
);
}
#[test]
fn static_domain_fact_flag_tracks_the_walk() {
let pool = ExprPool::new();
let x = pool.symbol("x", Domain::Real);
let p = pool.symbol("p", Domain::Positive);
let n = pool.symbol("n", Domain::NonZero);
let plain = pool.add(vec![x, pool.func("sin", vec![x])]);
assert!(!pool.has_static_domain_fact(plain));
assert!(pool.has_static_domain_fact(p));
assert!(pool.has_static_domain_fact(pool.pow(x, n)));
let cond = pool.pred_gt(p, pool.integer(0_i32));
let pw = pool.piecewise(vec![(cond, x)], x);
assert!(pool.has_static_domain_fact(pw));
let bound = pool.forall(n, x);
assert!(pool.has_static_domain_fact(bound));
for e in [plain, pw, bound] {
let (mut fast, mut slow) = (Vec::new(), Vec::new());
collect_static_domain_facts(e, &pool, &mut fast);
collect_static_domain_facts_reference(e, &pool, &mut slow);
assert_eq!(fast, slow);
}
}
#[test]
fn simplify_of_a_simplified_shared_dag_is_linear() {
let pool = ExprPool::new();
let x = pool.symbol("x", Domain::Real);
let (mut a, mut b) = (pool.integer(1_i32), x);
for _ in 1..40 {
let c = pool.add(vec![pool.mul(vec![x, b]), a]);
a = b;
b = c;
}
let once = crate::simplify::simplify(b, &pool).value;
let start = std::time::Instant::now();
let twice = crate::simplify::simplify(once, &pool);
assert!(start.elapsed() < std::time::Duration::from_secs(2));
assert_eq!(twice.value, once);
assert!(twice.log.is_empty());
}
}
#[cfg(test)]
fn tiny_float(pool: &ExprPool, negative: bool) -> ExprId {
let mut f = rug::Float::with_val(64, rug::Float::i_exp(1, -1400));
if negative {
f = -f;
}
assert_eq!(f.to_f64(), 0.0, "the premise: f64 flushes it to zero");
pool.intern(ExprData::Float(crate::kernel::BigFloat {
inner: f,
prec: 64,
}))
}
#[cfg(test)]
mod float_sign_tests {
use super::*;
#[test]
fn a_float_below_f64_range_is_not_zero() {
let pool = ExprPool::new();
assert_eq!(
assumed_sign(tiny_float(&pool, false), &pool),
Some(Sign::Positive)
);
assert_eq!(
assumed_sign(tiny_float(&pool, true), &pool),
Some(Sign::Negative)
);
assert_eq!(assumed_sign(pool.float(0.0, 53), &pool), Some(Sign::Zero));
}
}