use rustc_hash::{FxHashMap, FxHashSet};
use smallvec::SmallVec;
use crate::base::arena::Arena;
use crate::base::node::SymbolId;
use crate::base::node::{ExprId, ExprNode};
pub(crate) fn post_order_ids(arena: &Arena, root: ExprId) -> Vec<ExprId> {
let mut result = Vec::new();
let mut visited: FxHashSet<ExprId> = FxHashSet::default();
let mut stack: Vec<(ExprId, bool)> = vec![(root, false)];
while let Some((id, children_pushed)) = stack.last_mut() {
let id = *id;
if visited.contains(&id) {
stack.pop();
continue;
}
if !*children_pushed {
*children_pushed = true;
let node = arena.node(id);
let mut child_buf: SmallVec<[ExprId; 6]> = SmallVec::new();
node.for_each_child(|c| child_buf.push(c));
for &child in child_buf.iter().rev() {
if !visited.contains(&child) {
stack.push((child, false));
}
}
} else {
stack.pop();
visited.insert(id);
result.push(id);
}
}
result
}
pub(crate) fn walk_and_rebuild(
arena: &mut Arena,
root: ExprId,
transform: &dyn Fn(&Arena, ExprId) -> Option<ExprId>,
) -> ExprId {
let mut cache: FxHashMap<ExprId, ExprId> = FxHashMap::default();
let post_order = post_order_ids(arena, root);
for &id in &post_order {
if let Some(replacement) = transform(arena, id) {
cache.insert(id, replacement);
continue;
}
if arena.node(id).is_atom() {
cache.insert(id, id);
continue;
}
let new_id = rebuild_with_cache(arena, id, &cache);
cache.insert(id, new_id);
}
cache.get(&root).copied().unwrap_or(root)
}
pub(crate) fn rebuild_with_cache(
arena: &mut Arena,
id: ExprId,
cache: &FxHashMap<ExprId, ExprId>,
) -> ExprId {
rebuild_with(arena, id, &|c| cache.get(&c).copied().unwrap_or(c))
}
pub(crate) fn rebuild_with<F: Fn(ExprId) -> ExprId>(
arena: &mut Arena,
id: ExprId,
get: &F,
) -> ExprId {
match arena.node(id).map_children(get) {
None => id,
Some(node) => arena.make(node),
}
}
pub(crate) fn contains(arena: &Arena, haystack: ExprId, needle: ExprId) -> bool {
if haystack == needle {
return true;
}
let mut visited: FxHashSet<ExprId> = FxHashSet::default();
let mut stack: Vec<ExprId> = vec![haystack];
while let Some(id) = stack.pop() {
if id == needle {
return true;
}
if visited.contains(&id) {
continue;
}
visited.insert(id);
arena.node(id).for_each_child(|c| stack.push(c));
}
false
}
#[derive(Clone, Debug)]
pub(crate) struct Binder {
pub(crate) var: ExprId,
pub(crate) scoped: SmallVec<[ExprId; 2]>,
pub(crate) outer: SmallVec<[ExprId; 2]>,
}
pub(crate) fn binder(arena: &Arena, id: ExprId) -> Option<Binder> {
arena
.node(id)
.binder_operands()
.filter(|b| matches!(arena.node(b.var), ExprNode::Symbol(_)))
}
pub(crate) fn rebuild_binder(
arena: &mut Arena,
id: ExprId,
var: ExprId,
scoped: &[ExprId],
outer: &[ExprId],
) -> ExprId {
match arena.node(id).with_binder_operands(var, scoped, outer) {
Some(node) => arena.make(node),
None => id,
}
}
pub(crate) fn var_slot(arena: &Arena, id: ExprId) -> Option<Binder> {
arena
.node(id)
.var_slot_operands()
.filter(|b| matches!(arena.node(b.var), ExprNode::Symbol(_)))
}
pub(crate) fn rebuild_var_slot(
arena: &mut Arena,
id: ExprId,
var: ExprId,
scoped: &[ExprId],
outer: &[ExprId],
) -> ExprId {
match arena.node(id).with_var_slot_operands(var, scoped, outer) {
Some(node) => arena.make(node),
None => id,
}
}
pub(crate) fn subs_node(arena: &mut Arena, body: ExprId, var: ExprId, point: ExprId) -> ExprId {
if let ExprNode::Symbol(sym) = *arena.node(var)
&& (point == var || !has_free_symbol(arena, body, sym))
{
return body;
}
arena.intern(ExprNode::Subs(body, var, point))
}
pub(crate) fn malformed_binder(arena: &Arena, root: ExprId, except: ExprId) -> Option<ExprId> {
let old: FxHashSet<ExprId> = post_order_ids(arena, except).into_iter().collect();
post_order_ids(arena, root).into_iter().find(|id| {
!old.contains(id)
&& arena
.node(*id)
.var_operand()
.is_some_and(|v| !matches!(arena.node(v), ExprNode::Symbol(_)))
})
}
pub(crate) fn root_of_implied_var(arena: &Arena, poly: ExprId) -> Option<ExprId> {
match free_symbols(arena, poly)[..] {
[var] => Some(var),
_ => None,
}
}
pub(crate) fn has_free_symbol(arena: &Arena, root: ExprId, sym: SymbolId) -> bool {
let mut stack: Vec<ExprId> = vec![root];
let mut visited: FxHashSet<ExprId> = FxHashSet::default();
while let Some(id) = stack.pop() {
if !visited.insert(id) {
continue;
}
match arena.node(id) {
ExprNode::Symbol(s) => {
if *s == sym {
return true;
}
}
node if node.is_atom() => {}
node => match binder(arena, id) {
Some(b) => {
stack.extend_from_slice(&b.outer);
if !matches!(arena.node(b.var), ExprNode::Symbol(s) if *s == sym) {
stack.extend_from_slice(&b.scoped);
}
}
None => node.for_each_child(|c| stack.push(c)),
},
}
}
false
}
pub(crate) fn has_free_var(arena: &Arena, root: ExprId, var: ExprId) -> bool {
match *arena.node(var) {
ExprNode::Symbol(sym) => has_free_symbol(arena, root, sym),
_ => contains(arena, root, var),
}
}
pub(crate) fn free_symbols(arena: &Arena, root: ExprId) -> Vec<ExprId> {
let mut scopes: Vec<SmallVec<[ExprId; 4]>> = vec![SmallVec::new()];
let mut result = Vec::new();
let mut visited: FxHashSet<(ExprId, u32)> = FxHashSet::default();
let mut seen_syms: FxHashSet<SymbolId> = FxHashSet::default();
let mut stack: Vec<(ExprId, u32)> = vec![(root, 0)];
fn extend_scope(scopes: &mut Vec<SmallVec<[ExprId; 4]>>, sc: u32, var: ExprId) -> u32 {
if scopes[sc as usize].contains(&var) {
return sc;
}
let mut new_scope = scopes[sc as usize].clone();
new_scope.push(var);
new_scope.sort_unstable_by_key(|&v| crate::base::arena::ArenaOrder(v));
if let Some(pos) = scopes.iter().position(|s| *s == new_scope) {
return pos as u32;
}
scopes.push(new_scope);
(scopes.len() - 1) as u32
}
while let Some((id, sc)) = stack.pop() {
if !visited.insert((id, sc)) {
continue;
}
match arena.node(id) {
ExprNode::Symbol(sid) => {
if !scopes[sc as usize].contains(&id) && seen_syms.insert(*sid) {
result.push(id);
}
}
node if node.is_atom() => {}
node => match binder(arena, id) {
Some(b) => {
for &c in &b.outer {
stack.push((c, sc));
}
let inner = extend_scope(&mut scopes, sc, b.var);
for &c in &b.scoped {
stack.push((c, inner));
}
}
None => node.for_each_child(|c| stack.push((c, sc))),
},
}
}
result
}
pub(crate) fn has_unevaluated(arena: &Arena, root: ExprId) -> bool {
let mut stack = vec![root];
let mut visited = FxHashSet::default();
while let Some(id) = stack.pop() {
if !visited.insert(id) {
continue;
}
match arena.node(id) {
ExprNode::Integral(..)
| ExprNode::DefiniteIntegral(..)
| ExprNode::Derivative(..)
| ExprNode::Limit(..)
| ExprNode::Series(..)
| ExprNode::LaplaceTransform(..)
| ExprNode::InverseLaplaceTransform(..)
| ExprNode::Residue(..)
| ExprNode::DSolve(..)
| ExprNode::ConditionSet(..)
| ExprNode::Subs(..)
| ExprNode::Sum(..)
| ExprNode::Product_(..) => return true,
node => {
for child in node.children() {
stack.push(child);
}
}
}
}
false
}
#[cfg(test)]
mod tests {
use super::*;
use crate::base::arena::Arena;
fn sym(a: &mut Arena, name: &str) -> ExprId {
a.symbol(name)
}
fn display(a: &Arena, id: ExprId) -> String {
a.display(id).to_string()
}
#[test]
fn post_order_atom() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let order = post_order_ids(&a, x);
assert_eq!(order, vec![x]);
}
#[test]
fn post_order_add() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let y = sym(&mut a, "y");
let sum = a.add(&[x, y]);
let order = post_order_ids(&a, sum);
assert!(
order.iter().position(|&id| id == x).unwrap()
< order.iter().position(|&id| id == sum).unwrap()
);
assert!(
order.iter().position(|&id| id == y).unwrap()
< order.iter().position(|&id| id == sum).unwrap()
);
}
#[test]
fn post_order_deduplicates() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let expr = a.add(&[x, x]);
let order = post_order_ids(&a, expr);
let x_count = order.iter().filter(|&&id| id == x).count();
assert_eq!(x_count, 1, "x should appear exactly once");
}
#[test]
fn walk_identity_is_noop() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let two = a.int(2);
let xp = a.pow(x, two);
let one = a.one;
let expr = a.add(&[xp, x, one]);
let result = walk_and_rebuild(&mut a, expr, &|_, _| None);
assert_eq!(result, expr, "identity walk should return same ExprId");
}
#[test]
fn walk_replace_leaf() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let y = sym(&mut a, "y");
let two = a.int(2);
let expr = a.pow(x, two);
let result = walk_and_rebuild(&mut a, expr, &|_, id| {
if id == x { Some(y) } else { None }
});
assert_eq!(display(&a, result), "y^2");
}
#[test]
fn walk_deep_no_overflow() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let y = sym(&mut a, "y");
let mut expr = x;
for _ in 0..1000 {
expr = a.sin(expr);
}
let result = walk_and_rebuild(&mut a, expr, &|_, id| {
if id == x { Some(y) } else { None }
});
assert_ne!(result, expr);
}
#[test]
fn contains_self() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
assert!(contains(&a, x, x));
}
#[test]
fn contains_child() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let y = sym(&mut a, "y");
let sum = a.add(&[x, y]);
assert!(contains(&a, sum, x));
assert!(contains(&a, sum, y));
}
#[test]
fn does_not_contain() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let y = sym(&mut a, "y");
let z = sym(&mut a, "z");
let sum = a.add(&[x, y]);
assert!(!contains(&a, sum, z));
}
#[test]
fn contains_deep() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let two = a.int(2);
let xp = a.pow(x, two);
let expr = a.sin(xp);
assert!(contains(&a, expr, x));
assert!(contains(&a, expr, xp));
}
#[test]
fn free_symbols_atom() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let syms = free_symbols(&a, x);
assert_eq!(syms.len(), 1);
}
#[test]
fn free_symbols_expression() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let y = sym(&mut a, "y");
let expr = a.add(&[x, y]);
let syms = free_symbols(&a, expr);
assert_eq!(syms.len(), 2);
}
#[test]
fn free_symbols_no_duplicates() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let expr = a.add(&[x, x]);
let syms = free_symbols(&a, expr);
assert_eq!(syms.len(), 1, "x should appear once even in 2*x");
}
#[test]
fn free_symbols_constant_has_none() {
let a = Arena::new();
let syms = free_symbols(&a, a.pi);
assert!(syms.is_empty(), "pi has no free symbols");
}
#[test]
fn free_symbols_sum_binds_index_but_not_bounds() {
let mut a = Arena::new();
let k = sym(&mut a, "k");
let n = sym(&mut a, "n");
let body = a.sin(k);
let s = a.intern(ExprNode::Sum(body, k, a.one, n));
assert_eq!(free_symbols(&a, s), vec![n]);
let p = a.intern(ExprNode::Product_(body, k, a.one, n));
assert_eq!(free_symbols(&a, p), vec![n]);
let s2 = a.intern(ExprNode::Sum(body, k, a.one, k));
assert_eq!(free_symbols(&a, s2), vec![k]);
}
#[test]
fn free_symbols_same_node_bound_and_free() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let three = a.int(3);
let s = a.intern(ExprNode::Sum(x, x, a.zero, three));
assert!(free_symbols(&a, s).is_empty());
let e = a.add(&[x, s]);
assert_eq!(free_symbols(&a, e), vec![x]);
let e2 = a.add(&[s, x]);
assert_eq!(free_symbols(&a, e2), vec![x]);
}
#[test]
fn free_symbols_rootof_and_rootsum_bind_polynomial_variable() {
let mut a = Arena::new();
let lam = sym(&mut a, "lambda");
let three = a.int(3);
let l3 = a.pow(lam, three);
let neg_lam = a.neg(lam);
let poly = a.add(&[l3, neg_lam, a.neg_one]);
let root = a.intern(ExprNode::RootOf(poly, lam, a.zero));
assert!(free_symbols(&a, root).is_empty(), "RootOf is a constant");
let n = sym(&mut a, "n");
let root_n = a.intern(ExprNode::RootOf(poly, lam, n));
assert_eq!(free_symbols(&a, root_n), vec![n]);
let t = sym(&mut a, "t");
let x = sym(&mut a, "x");
let t3 = a.pow(t, three);
let poly_t = a.add(&[t3, t, a.one]);
let tx = a.mul(&[t, x]);
let body = a.ln(tx);
let rs = a.intern(ExprNode::RootSum(poly_t, body, t));
assert_eq!(free_symbols(&a, rs), vec![x]);
}
#[test]
fn free_symbols_rootof_multivariate_polynomial_binds_its_variable_only() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let p = sym(&mut a, "p");
let five = a.int(5);
let x5 = a.pow(x, five);
let px = a.mul(&[p, x]);
let poly = a.add(&[x5, px, a.one]);
let root = a.intern(ExprNode::RootOf(poly, x, a.zero));
assert_eq!(free_symbols(&a, root), vec![p]);
}
#[test]
fn free_symbols_condition_set_binds_variable() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let y = sym(&mut a, "y");
let sx = a.sin(x);
let cond = a.gt(sx, y);
let cs = a.intern(ExprNode::ConditionSet(x, cond));
assert_eq!(free_symbols(&a, cs), vec![y]);
}
#[test]
fn free_symbols_integral_and_derivative_do_not_bind() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let sx = a.sin(x);
let integral = a.intern(ExprNode::Integral(sx, x));
assert_eq!(free_symbols(&a, integral), vec![x]);
let deriv = a.intern(ExprNode::Derivative(sx, x));
assert_eq!(free_symbols(&a, deriv), vec![x]);
}
#[test]
fn free_symbols_nested_binders() {
let mut a = Arena::new();
let i = sym(&mut a, "i");
let j = sym(&mut a, "j");
let n = sym(&mut a, "n");
let ij = a.mul(&[i, j]);
let inner = a.intern(ExprNode::Sum(ij, j, a.one, i));
assert_eq!(free_symbols(&a, inner), vec![i]);
let outer = a.intern(ExprNode::Sum(inner, i, a.one, n));
assert_eq!(free_symbols(&a, outer), vec![n]);
}
#[test]
fn free_symbols_limit_residue_and_transforms_bind_their_dummy() {
let mut a = Arena::new();
let (x, y, t, s) = (
sym(&mut a, "x"),
sym(&mut a, "y"),
sym(&mut a, "t"),
sym(&mut a, "s"),
);
let xy = a.mul(&[x, y]);
let lim = a.intern(ExprNode::Limit(xy, x, a.zero));
assert_eq!(free_symbols(&a, lim), vec![y]);
let res = a.intern(ExprNode::Residue(xy, x, y));
assert_eq!(free_symbols(&a, res), vec![y]);
let ty = a.mul(&[t, y]);
let lt = a.intern(ExprNode::LaplaceTransform(ty, t, s));
let by_id = |v: &ExprId| crate::base::arena::ArenaOrder(*v);
let mut got = free_symbols(&a, lt);
got.sort_unstable_by_key(by_id);
let mut want = vec![y, s];
want.sort_unstable_by_key(by_id);
assert_eq!(got, want);
let ilt = a.intern(ExprNode::InverseLaplaceTransform(ty, t, s));
let mut got = free_symbols(&a, ilt);
got.sort_unstable_by_key(by_id);
assert_eq!(got, want);
let lim_x = a.intern(ExprNode::Limit(xy, x, x));
let mut got = free_symbols(&a, lim_x);
got.sort_unstable_by_key(by_id);
let mut want = vec![x, y];
want.sort_unstable_by_key(by_id);
assert_eq!(got, want);
}
#[test]
fn has_free_symbol_agrees_with_free_symbols() {
let mut a = Arena::new();
let (x, k, n) = (sym(&mut a, "x"), sym(&mut a, "k"), sym(&mut a, "n"));
let ExprNode::Symbol(x_sym) = *a.node(x) else {
unreachable!()
};
let three = a.int(3);
let xk = a.mul(&[x, k]);
let five = a.int(5);
let x5 = a.pow(x, five);
let neg_x = a.neg(x);
let poly = a.add(&[x5, neg_x, a.one]);
let root = a.intern(ExprNode::RootOf(poly, x, a.zero));
let sum_x = a.intern(ExprNode::Sum(xk, x, a.zero, three));
let sum_k = a.intern(ExprNode::Sum(xk, k, a.zero, n));
let sum_x_to_x = a.intern(ExprNode::Sum(xk, x, a.zero, x));
let integral = a.intern(ExprNode::Integral(xk, x));
let shared = a.add(&[x, sum_x]);
for (e, free) in [
(root, false),
(sum_x, false),
(sum_k, true),
(sum_x_to_x, true),
(integral, true),
(shared, true),
] {
assert_eq!(has_free_symbol(&a, e, x_sym), free, "{}", display(&a, e));
assert_eq!(free_symbols(&a, e).contains(&x), free, "{}", display(&a, e));
}
}
#[test]
fn has_unevaluated_formal_nodes() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let integral = a.intern(ExprNode::Integral(x, x));
assert!(has_unevaluated(&a, integral));
let deriv = a.intern(ExprNode::Derivative(x, x));
assert!(has_unevaluated(&a, deriv));
let lim = a.intern(ExprNode::Limit(x, x, a.zero));
assert!(has_unevaluated(&a, lim));
let sum = a.add(&[integral, a.one]);
assert!(has_unevaluated(&a, sum));
}
#[test]
fn root_of_and_root_sum_are_not_unevaluated() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
let five = a.int(5);
let x5 = a.pow(x, five);
let poly = a.add(&[x5, x, a.one]);
let root = a.intern(ExprNode::RootOf(poly, x, a.zero));
assert!(
!has_unevaluated(&a, root),
"RootOf is a complete algebraic answer"
);
let t = sym(&mut a, "t");
let body = a.ln(t);
let rs = a.intern(ExprNode::RootSum(poly, body, t));
assert!(!has_unevaluated(&a, rs), "RootSum is not unevaluated");
}
#[test]
fn function_nodes_are_not_unevaluated() {
let mut a = Arena::new();
let x = sym(&mut a, "x");
for id in [
a.intern(ExprNode::Re(x)),
a.intern(ExprNode::Im(x)),
a.intern(ExprNode::Conjugate(x)),
a.intern(ExprNode::Arg(x)),
a.intern(ExprNode::Si(x)),
a.intern(ExprNode::Ci(x)),
a.intern(ExprNode::Ei(x)),
a.intern(ExprNode::Li(x)),
a.intern(ExprNode::Zeta(x)),
a.intern(ExprNode::Polygamma(a.one, x)),
a.intern(ExprNode::KroneckerDelta(x, a.one)),
] {
assert!(!has_unevaluated(&a, id), "{:?} is a function", a.node(id));
}
}
}