use crate::base::arena::Arena;
use crate::base::node::{ExprId, ExprNode, SymbolId};
use num_bigint::BigInt;
use num_rational::Ratio;
use num_traits::{One, Zero};
fn binom(n: u64, k: u64) -> BigInt {
if k > n {
return BigInt::zero();
}
if k == 0 || k == n {
return BigInt::one();
}
let k = k.min(n - k);
let mut result = BigInt::one();
for i in 0..k {
result = result * BigInt::from(n - i) / BigInt::from(i + 1);
}
result
}
fn contains_var(arena: &Arena, expr: ExprId, var_sym: SymbolId) -> bool {
crate::base::walk::has_free_symbol(arena, expr, var_sym)
}
fn as_i64(arena: &Arena, id: ExprId) -> Option<i64> {
let r = arena.as_num(id)?;
if !r.is_integer() {
return None;
}
let big = r.to_integer();
i64::try_from(&big).ok()
}
#[derive(Clone, Copy)]
struct TrigPower {
arg: ExprId,
exp: i64,
}
fn extract_power(
arena: &Arena,
factor: ExprId,
func: fn(&ExprNode) -> Option<ExprId>,
) -> Option<TrigPower> {
if let Some(arg) = func(arena.node(factor)) {
return Some(TrigPower { arg, exp: 1 });
}
match arena.node(factor) {
ExprNode::Pow(base, exp) => {
let arg = func(arena.node(*base))?;
Some(TrigPower {
arg,
exp: as_i64(arena, *exp)?,
})
}
_ => None,
}
}
fn sin_arg(node: &ExprNode) -> Option<ExprId> {
match node {
ExprNode::Sin(u) => Some(*u),
_ => None,
}
}
fn cos_arg(node: &ExprNode) -> Option<ExprId> {
match node {
ExprNode::Cos(u) => Some(*u),
_ => None,
}
}
fn tan_arg(node: &ExprNode) -> Option<ExprId> {
match node {
ExprNode::Tan(u) => Some(*u),
_ => None,
}
}
fn sinh_arg(node: &ExprNode) -> Option<ExprId> {
match node {
ExprNode::Sinh(u) => Some(*u),
_ => None,
}
}
fn cosh_arg(node: &ExprNode) -> Option<ExprId> {
match node {
ExprNode::Cosh(u) => Some(*u),
_ => None,
}
}
fn linear_rate(arena: &mut Arena, u: ExprId, var: ExprId, var_sym: SymbolId) -> Option<ExprId> {
if u == var {
return Some(arena.one);
}
if !contains_var(arena, u, var_sym) {
return None;
}
let a = crate::transforms::diff::diff(arena, u, var);
(a != arena.zero && !contains_var(arena, a, var_sym)).then_some(a)
}
const MAX_TRIG_POWER: i64 = 256;
fn within_power_bound(m: i64, n: i64) -> bool {
m.checked_abs().is_some_and(|a| a <= MAX_TRIG_POWER)
&& n.checked_abs().is_some_and(|a| a <= MAX_TRIG_POWER)
}
#[cfg(test)]
fn extract_sin_power(arena: &Arena, factor: ExprId, var: ExprId) -> Option<i64> {
extract_power(arena, factor, sin_arg)
.filter(|p| p.arg == var)
.map(|p| p.exp)
}
#[cfg(test)]
fn extract_cos_power(arena: &Arena, factor: ExprId, var: ExprId) -> Option<i64> {
extract_power(arena, factor, cos_arg)
.filter(|p| p.arg == var)
.map(|p| p.exp)
}
const MAX_LINEARIZE_DEGREE: u32 = 12;
pub(crate) fn linearize_sin_cos(arena: &mut Arena, u: ExprId, m: u32, n: u32) -> Option<ExprId> {
use std::collections::BTreeMap;
if m + n > MAX_LINEARIZE_DEGREE {
return None;
}
let half = Ratio::new(BigInt::one(), BigInt::from(2));
let mut terms: BTreeMap<u32, (Ratio<BigInt>, Ratio<BigInt>)> = BTreeMap::new();
terms.insert(0, (Ratio::one(), Ratio::zero()));
let zero = || (Ratio::<BigInt>::zero(), Ratio::<BigInt>::zero());
for step in 0..m + n {
let by_sin = step < m;
let mut next: BTreeMap<u32, (Ratio<BigInt>, Ratio<BigInt>)> = BTreeMap::new();
for (k, (a, b)) in terms {
let (ah, bh) = (&a * &half, &b * &half);
if by_sin {
let up = next.entry(k + 1).or_insert_with(zero);
up.1 += &ah;
up.0 -= &bh;
let down = next.entry(k.abs_diff(1)).or_insert_with(zero);
if k >= 1 {
down.1 -= &ah;
} else {
down.1 += &ah;
}
down.0 += &bh;
} else {
let up = next.entry(k + 1).or_insert_with(zero);
up.0 += &ah;
up.1 += &bh;
let down = next.entry(k.abs_diff(1)).or_insert_with(zero);
down.0 += &ah;
if k >= 1 {
down.1 += &bh;
} else {
down.1 -= &bh;
}
}
}
terms = next;
}
let mut out: Vec<ExprId> = Vec::new();
for (k, (a, b)) in terms {
if k == 0 {
if !a.is_zero() {
out.push(arena.num_ratio(a));
}
continue;
}
let k_id = arena.int(i64::from(k));
let ku = arena.mul(&[k_id, u]);
for (c, is_sin) in [(a, false), (b, true)] {
if c.is_zero() {
continue;
}
let g = if is_sin { arena.sin(ku) } else { arena.cos(ku) };
let c_id = arena.num_ratio(c);
out.push(arena.mul(&[c_id, g]));
}
}
Some(match out.len() {
0 => arena.zero,
1 => out[0],
_ => arena.add(&out),
})
}
fn trig_exp_polynomial(arena: &mut Arena, e: ExprId, var: ExprId, var_sym: SymbolId) -> bool {
let mut arguments: Vec<ExprId> = Vec::new();
let mut stack = vec![e];
let mut seen: rustc_hash::FxHashSet<ExprId> = rustc_hash::FxHashSet::default();
while let Some(id) = stack.pop() {
if !seen.insert(id) || !contains_var(arena, id, var_sym) {
continue;
}
match arena.node(id) {
ExprNode::Symbol(_) => {}
ExprNode::Add(ch) | ExprNode::Mul(ch) => stack.extend(ch.iter().copied()),
ExprNode::Pow(b, n) if as_i64(arena, *n).is_some_and(|k| k >= 0) => stack.push(*b),
ExprNode::Exp(u) | ExprNode::Sin(u) | ExprNode::Cos(u) => arguments.push(*u),
_ => return false,
}
}
arguments
.into_iter()
.all(|u| linear_rate(arena, u, var, var_sym).is_some())
}
pub(crate) fn linearize_trig_powers(
arena: &mut Arena,
factors: &[ExprId],
var: ExprId,
var_sym: SymbolId,
) -> Option<ExprId> {
let mut groups: Vec<(ExprId, u32, u32)> = Vec::new();
for &f in factors {
let found = match extract_power(arena, f, sin_arg) {
Some(p) => Some((p, true)),
None => extract_power(arena, f, cos_arg).map(|p| (p, false)),
};
let Some((p, is_sin)) = found else {
continue;
};
let e = u32::try_from(p.exp).ok()?;
match groups.iter_mut().find(|g| g.0 == p.arg) {
Some(g) if is_sin => g.1 += e,
Some(g) => g.2 += e,
None => groups.push(if is_sin { (p.arg, e, 0) } else { (p.arg, 0, e) }),
}
}
let &(u, m, n) = groups.iter().find(|g| g.1 + g.2 >= 2)?;
for &f in factors {
if !trig_exp_polynomial(arena, f, var, var_sym) {
return None;
}
}
let is_group_member = |arena: &Arena, f: ExprId| {
extract_power(arena, f, sin_arg)
.or_else(|| extract_power(arena, f, cos_arg))
.is_some_and(|p| p.arg == u)
};
let mut rest: Vec<ExprId> = factors
.iter()
.copied()
.filter(|&f| !is_group_member(arena, f))
.collect();
let linear = linearize_sin_cos(arena, u, m, n)?;
rest.push(linear);
Some(arena.mul(&rest))
}
fn over_rate(arena: &mut Arena, big_f: ExprId, rate: ExprId) -> ExprId {
if rate == arena.one {
big_f
} else {
arena.div(big_f, rate)
}
}
pub(crate) fn sin_pow_integrate(arena: &mut Arena, n: i64, var: ExprId) -> ExprId {
if n == 0 {
return var; }
if n == 1 {
let cos_x = arena.cos(var);
return arena.neg(cos_x);
}
if n == -1 {
let sin_x = arena.sin(var);
let cos_x = arena.cos(var);
let neg_one_id = arena.int(-1);
let csc_x = arena.pow(sin_x, neg_one_id); let cot_x = arena.mul(&[cos_x, csc_x]); let sum = arena.add(&[csc_x, cot_x]);
let abs_sum = arena.abs(sum);
let ln_val = arena.ln(abs_sum);
return arena.neg(ln_val);
}
if n < -1 {
let cos_x = arena.cos(var);
let sin_x = arena.sin(var);
let exp_id = arena.int(n + 1);
let sin_pow = arena.pow(sin_x, exp_id);
let inv = arena.rational(1, n + 1);
let first_term = arena.mul(&[inv, cos_x, sin_pow]);
let coeff = arena.rational(n + 2, n + 1);
let recursive = sin_pow_integrate(arena, n + 2, var);
let second_term = arena.mul(&[coeff, recursive]);
return arena.add(&[first_term, second_term]);
}
let cos_x = arena.cos(var);
let sin_x = arena.sin(var);
let sin_pow = if n - 1 == 1 {
sin_x
} else {
let n_minus_1 = arena.int(n - 1);
arena.pow(sin_x, n_minus_1)
};
let neg_inv_n = arena.rational(-1, n);
let first_term = arena.mul(&[neg_inv_n, cos_x, sin_pow]);
let coeff = arena.rational(n - 1, n);
let recursive = sin_pow_integrate(arena, n - 2, var);
let second_term = arena.mul(&[coeff, recursive]);
arena.add(&[first_term, second_term])
}
pub(crate) fn cos_pow_integrate(arena: &mut Arena, n: i64, var: ExprId) -> ExprId {
if n == 0 {
return var; }
if n == 1 {
return arena.sin(var);
}
if n == -1 {
let cos_x = arena.cos(var);
let neg_one_id = arena.int(-1);
let sec_x = arena.pow(cos_x, neg_one_id); let tan_x = arena.tan(var);
let sum = arena.add(&[sec_x, tan_x]);
let abs_sum = arena.abs(sum);
return arena.ln(abs_sum);
}
if n < -1 {
let sin_x = arena.sin(var);
let cos_x = arena.cos(var);
let exp_id = arena.int(n + 1);
let cos_pow = arena.pow(cos_x, exp_id);
let neg_inv = arena.rational(-1, n + 1);
let first_term = arena.mul(&[neg_inv, sin_x, cos_pow]);
let coeff = arena.rational(n + 2, n + 1);
let recursive = cos_pow_integrate(arena, n + 2, var);
let second_term = arena.mul(&[coeff, recursive]);
return arena.add(&[first_term, second_term]);
}
let sin_x = arena.sin(var);
let cos_x = arena.cos(var);
let cos_pow = if n - 1 == 1 {
cos_x
} else {
let n_minus_1 = arena.int(n - 1);
arena.pow(cos_x, n_minus_1)
};
let inv_n = arena.rational(1, n);
let first_term = arena.mul(&[inv_n, sin_x, cos_pow]);
let coeff = arena.rational(n - 1, n);
let recursive = cos_pow_integrate(arena, n - 2, var);
let second_term = arena.mul(&[coeff, recursive]);
arena.add(&[first_term, second_term])
}
pub(crate) fn sin_cos_integrate(arena: &mut Arena, m: i64, n: i64, var: ExprId) -> ExprId {
if m == 0 {
return cos_pow_integrate(arena, n, var);
}
if n == 0 {
return sin_pow_integrate(arena, m, var);
}
if m > 0 && m % 2 != 0 {
let cos_x = arena.cos(var);
return odd_power_substitution(arena, (m - 1) / 2, n, cos_x, true);
}
if n > 0 && n % 2 != 0 {
let sin_x = arena.sin(var);
return odd_power_substitution(arena, (n - 1) / 2, m, sin_x, false);
}
if n >= 2 && m < 0 {
return even_power_expansion(arena, n / 2, m, var, sin_pow_integrate);
}
if m >= 2 && n < 0 {
return even_power_expansion(arena, m / 2, n, var, cos_pow_integrate);
}
if m < 0 && n < 0 {
return both_negative_powers(arena, m, n, var);
}
debug_assert!(m % 2 == 0 && n % 2 == 0 && m >= 2 && n >= 2);
let sin_x = arena.sin(var);
let cos_x = arena.cos(var);
let mn = m + n;
let sin_pow = if m - 1 == 1 {
sin_x
} else {
let exp = arena.int(m - 1);
arena.pow(sin_x, exp)
};
let cos_pow = {
let exp = arena.int(n + 1);
arena.pow(cos_x, exp)
};
let neg_inv_mn = arena.rational(-1, mn);
let first_term = arena.mul(&[neg_inv_mn, sin_pow, cos_pow]);
let coeff = arena.rational(m - 1, mn);
let recursive = sin_cos_integrate(arena, m - 2, n, var);
let second_term = arena.mul(&[coeff, recursive]);
arena.add(&[first_term, second_term])
}
fn odd_power_substitution(arena: &mut Arena, k: i64, e: i64, u: ExprId, negate: bool) -> ExprId {
let mut terms: Vec<ExprId> = Vec::new();
for j in 0..=k {
let mut c = binom(k as u64, j as u64);
if (j % 2 != 0) != negate {
c = -c;
}
let power = e + 2 * j + 1;
let term = if power == 0 {
let abs_u = arena.abs(u);
let ln = arena.ln(abs_u);
let coeff = arena.intern_num(Ratio::from_integer(c));
let coeff = arena.intern(ExprNode::Num(coeff));
arena.mul(&[coeff, ln])
} else {
let ratio = Ratio::new(c, BigInt::from(power));
let coeff = arena.intern_num(ratio);
let coeff = arena.intern(ExprNode::Num(coeff));
let pow_id = arena.int(power);
let u_pow = arena.pow(u, pow_id);
arena.mul(&[coeff, u_pow])
};
terms.push(term);
}
arena.add(&terms)
}
fn even_power_expansion(
arena: &mut Arena,
h: i64,
e: i64,
var: ExprId,
single: fn(&mut Arena, i64, ExprId) -> ExprId,
) -> ExprId {
let mut terms: Vec<ExprId> = Vec::new();
for j in 0..=h {
let mut c = binom(h as u64, j as u64);
if j % 2 != 0 {
c = -c;
}
let integral = single(arena, e + 2 * j, var);
let coeff = arena.intern_num(Ratio::from_integer(c));
let coeff = arena.intern(ExprNode::Num(coeff));
terms.push(arena.mul(&[coeff, integral]));
}
arena.add(&terms)
}
fn both_negative_powers(arena: &mut Arena, m: i64, n: i64, var: ExprId) -> ExprId {
use std::collections::BTreeMap;
let mut open: BTreeMap<(i64, i64), BigInt> = BTreeMap::new();
let mut done: BTreeMap<(i64, i64), BigInt> = BTreeMap::new();
open.insert((m, n), BigInt::one());
while let Some(((a, b), c)) = open.pop_first() {
if a >= 0 || b >= 0 {
*done.entry((a, b)).or_insert_with(BigInt::zero) += c;
continue;
}
*open.entry((a + 2, b)).or_insert_with(BigInt::zero) += &c;
*open.entry((a, b + 2)).or_insert_with(BigInt::zero) += c;
}
let mut terms: Vec<ExprId> = Vec::new();
for ((a, b), c) in done {
let integral = sin_cos_integrate(arena, a, b, var);
let coeff = arena.intern_num(Ratio::from_integer(c));
let coeff = arena.intern(ExprNode::Num(coeff));
terms.push(arena.mul(&[coeff, integral]));
}
arena.add(&terms)
}
pub(crate) fn sinh_squared_integrate(arena: &mut Arena, var: ExprId) -> ExprId {
let two = arena.int(2);
let two_x = arena.mul(&[two, var]);
let sinh_2x = arena.sinh(two_x);
let quarter = arena.rational(1, 4);
let first_term = arena.mul(&[quarter, sinh_2x]);
let half = arena.rational(1, 2);
let second_term = arena.mul(&[half, var]);
arena.sub(first_term, second_term)
}
pub(crate) fn cosh_squared_integrate(arena: &mut Arena, var: ExprId) -> ExprId {
let two = arena.int(2);
let two_x = arena.mul(&[two, var]);
let sinh_2x = arena.sinh(two_x);
let quarter = arena.rational(1, 4);
let first_term = arena.mul(&[quarter, sinh_2x]);
let half = arena.rational(1, 2);
let second_term = arena.mul(&[half, var]);
arena.add(&[first_term, second_term])
}
pub(crate) fn try_trig_power_integral(
arena: &mut Arena,
expr: ExprId,
var: ExprId,
var_sym: SymbolId,
) -> Option<ExprId> {
let node = arena.node(expr).clone();
if let Some(p) = extract_power(arena, expr, sin_arg) {
if (0..2).contains(&p.exp) || !within_power_bound(p.exp, 0) {
return None;
}
let rate = linear_rate(arena, p.arg, var, var_sym)?;
let big_f = sin_pow_integrate(arena, p.exp, p.arg);
return Some(over_rate(arena, big_f, rate));
}
if let Some(p) = extract_power(arena, expr, cos_arg) {
if (0..2).contains(&p.exp) || !within_power_bound(p.exp, 0) {
return None;
}
let rate = linear_rate(arena, p.arg, var, var_sym)?;
let big_f = cos_pow_integrate(arena, p.exp, p.arg);
return Some(over_rate(arena, big_f, rate));
}
if let Some(p) = extract_power(arena, expr, sinh_arg)
&& p.exp == 2
{
let rate = linear_rate(arena, p.arg, var, var_sym)?;
let big_f = sinh_squared_integrate(arena, p.arg);
return Some(over_rate(arena, big_f, rate));
}
if let Some(p) = extract_power(arena, expr, cosh_arg)
&& p.exp == 2
{
let rate = linear_rate(arena, p.arg, var, var_sym)?;
let big_f = cosh_squared_integrate(arena, p.arg);
return Some(over_rate(arena, big_f, rate));
}
if let Some(p) = extract_power(arena, expr, tan_arg)
&& p.exp < 0
&& within_power_bound(p.exp, -p.exp)
{
let rate = linear_rate(arena, p.arg, var, var_sym)?;
let big_f = sin_cos_integrate(arena, p.exp, -p.exp, p.arg);
return Some(over_rate(arena, big_f, rate));
}
if let ExprNode::Mul(ref children) = node {
let mut sin_exp: i64 = 0;
let mut cos_exp: i64 = 0;
let mut other_factors: Vec<ExprId> = Vec::new();
let mut found_tan = false;
let mut arg: Option<ExprId> = None;
let mut same_arg = |p: TrigPower| -> bool { *arg.get_or_insert(p.arg) == p.arg };
for &child in children.iter() {
if let Some(p) = extract_power(arena, child, sin_arg) {
if !same_arg(p) {
return None;
}
sin_exp = sin_exp.checked_add(p.exp)?;
} else if let Some(p) = extract_power(arena, child, cos_arg) {
if !same_arg(p) {
return None;
}
cos_exp = cos_exp.checked_add(p.exp)?;
} else if let Some(p) = extract_power(arena, child, tan_arg) {
if !same_arg(p) {
return None;
}
sin_exp = sin_exp.checked_add(p.exp)?;
cos_exp = cos_exp.checked_sub(p.exp)?;
found_tan = true;
} else if contains_var(arena, child, var_sym) {
return None;
} else {
other_factors.push(child);
}
}
let u = arg?;
let dominated_by_basic = sin_exp == 0 && cos_exp == 1 || sin_exp == 1 && cos_exp == 0;
if dominated_by_basic && !found_tan {
return None;
}
if !within_power_bound(sin_exp, cos_exp) {
return None;
}
let rate = linear_rate(arena, u, var, var_sym)?;
let trig_result = sin_cos_integrate(arena, sin_exp, cos_exp, u);
let trig_result = over_rate(arena, trig_result, rate);
if other_factors.is_empty() {
return Some(trig_result);
}
other_factors.push(trig_result);
return Some(arena.mul(&other_factors));
}
None
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum EndValue {
Finite(ExprId),
Infinite,
}
fn leading_term(arena: &mut Arena, p: ExprId, t: ExprId) -> Option<(usize, ExprId)> {
let key = (std::ptr::from_ref::<Arena>(arena) as usize, p, t);
if let Some(found) = JUMP_MEMO.with(|m| {
m.borrow()
.as_ref()
.and_then(|m| m.leading.get(&key).copied())
}) {
return found;
}
let lead = leading_term_uncached(arena, p, t);
JUMP_MEMO.with(|m| {
if let Some(m) = m.borrow_mut().as_mut() {
m.leading.insert(key, lead);
}
});
lead
}
fn leading_term_uncached(arena: &mut Arena, p: ExprId, t: ExprId) -> Option<(usize, ExprId)> {
let mut coeffs = crate::transforms::solve::symbolic_poly_coeffs(arena, p, t)?;
while let Some(&lc) = coeffs.last() {
match crate::poly::algebraic::is_zero_checked(arena, lc) {
Some(true) => {
coeffs.pop();
}
Some(false) => return Some((coeffs.len() - 1, lc)),
None if !crate::base::walk::free_symbols(arena, lc).is_empty()
&& !arena.is_zero_structural(lc) =>
{
return Some((coeffs.len() - 1, lc));
}
None => return None,
}
}
None
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Placement {
Positive,
Negative,
Complex { re_sign: i8, im_at_least_one: bool },
}
fn placement(arena: &mut Arena, z: ExprId) -> Option<Placement> {
if let Ok(v) = crate::transforms::evalf::evalf_complex64(arena, z) {
let scale = v.norm();
if !scale.is_finite() || scale == 0.0 {
return None;
}
let tol = 1e-9 * scale;
if v.im.abs() <= tol {
return Some(if v.re > 0.0 {
Placement::Positive
} else {
Placement::Negative
});
}
let re_sign = if v.re > tol {
1
} else if v.re < -tol {
-1
} else {
0
};
if re_sign == 0 && (v.im.abs() - 1.0).abs() <= 1e-9 {
return None;
}
return Some(Placement::Complex {
re_sign,
im_at_least_one: v.im.abs() >= 1.0,
});
}
let mut facts = crate::base::assumptions::AssumptionCache::new();
if facts.query(arena, z, crate::base::assumptions::Props::POSITIVE) == Some(true) {
Some(Placement::Positive)
} else if facts.query(arena, z, crate::base::assumptions::Props::NEGATIVE) == Some(true) {
Some(Placement::Negative)
} else {
None
}
}
fn rational_end(arena: &mut Arena, r: ExprId, t: ExprId, positive: bool) -> Option<(i64, ExprId)> {
let (num, den) = arena.as_numer_denom_expr(r);
let (dn, ln) = leading_term(arena, num, t)?;
let (dd, ld) = leading_term(arena, den, t)?;
let d = i64::try_from(dn).ok()? - i64::try_from(dd).ok()?;
let mut lead = arena.div(ln, ld);
if !positive && d % 2 != 0 {
lead = arena.neg(lead);
}
Some((d, crate::transforms::eval::eval(arena, lead)))
}
fn log_end(
arena: &mut Arena,
p: ExprId,
t: ExprId,
lead: ExprId,
positive: bool,
) -> Option<ExprId> {
let principal = arena.ln(lead);
match placement(arena, lead) {
Some(Placement::Negative) => {}
Some(_) => return Some(principal),
None if ASSUME_REAL_PARAMETERS.with(std::cell::Cell::get)
&& rational_in_parameters(arena, lead) =>
{
let (num, den) = arena.as_numer_denom_expr(p);
if contains_var_id(arena, den, t) || !rational_in_parameters(arena, den) {
return None;
}
let coeffs = crate::transforms::solve::symbolic_poly_coeffs(arena, num, t)?;
return coeffs
.iter()
.all(|&k| rational_in_parameters(arena, k))
.then_some(principal);
}
None => return None,
}
let (num, den) = arena.as_numer_denom_expr(p);
if contains_var_id(arena, den, t) {
return None;
}
let coeffs = crate::transforms::solve::symbolic_poly_coeffs(arena, num, t)?;
let mut reals = crate::base::assumptions::AssumptionCache::new();
let mut undecided: Option<(usize, ExprId)> = None;
for (k, &c) in coeffs.iter().enumerate().rev() {
let c = arena.div(c, den);
let c = crate::transforms::eval::eval(arena, c);
match crate::transforms::realness::constant_realness_as_declared(arena, c, 30, &mut reals) {
Some(true) => continue,
None if ASSUME_REAL_PARAMETERS.with(std::cell::Cell::get)
&& rational_in_parameters(arena, c) =>
{
continue;
}
None if undecided.is_none() && only_real_parameters(arena, c) => {
undecided = Some((k, c));
continue;
}
None if undecided.is_some() => return None,
_ => {}
}
if undecided.is_some() {
return None;
}
let z = crate::transforms::evalf::evalf_complex64(arena, c).ok()?;
let im = z.im.abs();
if im.is_nan() || im <= 1e-12 * z.norm().max(1.0) {
return None;
}
let above = (z.im > 0.0) == (positive || k % 2 == 0);
if above {
return Some(principal);
}
let two_pi_i = {
let two = arena.int(2);
let pi = arena.pi();
let i = arena.i_unit();
arena.mul(&[two, pi, i])
};
return Some(arena.sub(principal, two_pi_i));
}
if let Some((k, c)) = undecided {
let im = arena.im(c);
let im = if positive || k % 2 == 0 {
im
} else {
arena.neg(im)
};
let sg = arena.sign(im);
let two = arena.int(2);
let sg_sq = arena.pow(sg, two);
let minus_sq = arena.neg(sg_sq);
let shift = arena.add(&[sg, minus_sq]);
let i = arena.i_unit();
let pi = arena.pi();
let correction = arena.mul(&[i, pi, shift]);
return Some(arena.add(&[principal, correction]));
}
Some(principal)
}
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
struct JumpKey {
arena: usize,
real_parameters: bool,
assumed_real: bool,
g: ExprId,
t: ExprId,
}
#[derive(Default)]
struct JumpMemoMaps {
jumps: rustc_hash::FxHashMap<JumpKey, Option<ExprId>>,
leading: rustc_hash::FxHashMap<(usize, ExprId, ExprId), Option<(usize, ExprId)>>,
}
thread_local! {
static JUMP_MEMO: std::cell::RefCell<Option<JumpMemoMaps>> =
const { std::cell::RefCell::new(None) };
}
pub(crate) struct JumpMemo(bool);
impl JumpMemo {
pub(crate) fn enter() -> Self {
JUMP_MEMO.with(|m| {
let mut m = m.borrow_mut();
let owner = m.is_none();
if owner {
*m = Some(JumpMemoMaps::default());
}
JumpMemo(owner)
})
}
}
impl Drop for JumpMemo {
fn drop(&mut self) {
if self.0 {
let _ = JUMP_MEMO.try_with(|m| {
if let Ok(mut m) = m.try_borrow_mut() {
*m = None;
}
});
}
}
}
fn memoized_jump(key: JumpKey, compute: impl FnOnce() -> Option<ExprId>) -> Option<ExprId> {
if let Some(found) =
JUMP_MEMO.with(|m| m.borrow().as_ref().and_then(|m| m.jumps.get(&key).copied()))
{
return found;
}
let jump = compute();
JUMP_MEMO.with(|m| {
if let Some(m) = m.borrow_mut().as_mut() {
m.jumps.insert(key, jump);
}
});
jump
}
fn only_real_parameters(arena: &Arena, c: ExprId) -> bool {
let symbols = crate::base::walk::free_symbols(arena, c);
let assumed = ASSUME_REAL_PARAMETERS.with(std::cell::Cell::get);
!symbols.is_empty()
&& symbols.into_iter().all(|s| match *arena.node(s) {
ExprNode::Symbol(sid) => {
assumed || crate::transforms::realness::symbol_declared_real(arena, sid)
}
_ => false,
})
}
fn rational_function_is_zero(arena: &mut Arena, e: ExprId) -> bool {
let e = polynomial_root_sums(arena, e);
if !rational_in_parameters(arena, e) {
return false;
}
let combined = crate::poly::polybridge::together(arena, e);
let (num, _) = crate::poly::polybridge::as_numer_denom(arena, combined);
let num = crate::transforms::expand::expand(arena, num);
if arena.is_zero_structural(num) {
return true;
}
let symbols = crate::base::walk::free_symbols(arena, e);
for point in 0..GENERIC_POINTS.len() {
let mut v = e;
for (k, &s) in symbols.iter().enumerate() {
let (p, q) = GENERIC_POINTS[(point + k) % GENERIC_POINTS.len()];
let value = arena.rational(p + i64::try_from(k).unwrap_or(0), q);
v = crate::transforms::subs::subs(arena, v, s, value);
}
let v = crate::transforms::eval::eval(arena, v);
if !arena.is_zero_structural(v) {
return false;
}
}
true
}
const GENERIC_POINTS: [(i64, i64); 3] = [(7_919, 1_031), (-6_421, 2_053), (4_259, 3_967)];
const MAX_NEWTON_DEGREE: usize = 64;
fn polynomial_root_sums(arena: &mut Arena, e: ExprId) -> ExprId {
let mut out = e;
for id in crate::base::walk::post_order_ids(arena, e) {
let ExprNode::RootSum(p, body, s) = *arena.node(id) else {
continue;
};
let (Some(a), Some(b)) = (
crate::transforms::solve::symbolic_poly_coeffs(arena, p, s),
crate::transforms::solve::symbolic_poly_coeffs(arena, body, s),
) else {
continue;
};
let n = a.len().saturating_sub(1);
if n == 0 || n > MAX_NEWTON_DEGREE || b.len() > MAX_NEWTON_DEGREE {
continue;
}
let lead = a[n];
let minus_one = arena.int(-1);
let inv_lead = arena.pow(lead, minus_one);
let r: Vec<ExprId> = (0..=n)
.map(|j| {
if j == 0 {
arena.one
} else {
arena.mul(&[a[n - j], inv_lead])
}
})
.collect();
let n_id = arena.int(i64::try_from(n).unwrap_or(0));
let mut sums: Vec<ExprId> = vec![n_id];
for k in 1..b.len() {
let mut terms: Vec<ExprId> = Vec::new();
if k <= n {
let k_id = arena.int(i64::try_from(k).unwrap_or(0));
terms.push(arena.mul(&[k_id, r[k]]));
}
for j in 1..k.min(n + 1) {
terms.push(arena.mul(&[r[j], sums[k - j]]));
}
let total = arena.add(&terms);
let s_k = arena.mul(&[minus_one, total]);
sums.push(crate::transforms::eval::eval(arena, s_k));
}
let terms: Vec<ExprId> = b
.iter()
.zip(&sums)
.map(|(&bk, &sk)| arena.mul(&[bk, sk]))
.collect();
let total = arena.add(&terms);
out = arena.subs_structural(out, id, total);
}
out
}
fn rational_in_parameters(arena: &Arena, c: ExprId) -> bool {
crate::base::walk::post_order_ids(arena, c)
.iter()
.all(|&id| match arena.node(id) {
ExprNode::Num(_) | ExprNode::Symbol(_) | ExprNode::Add(_) | ExprNode::Mul(_) => true,
ExprNode::Neg(_) => true,
ExprNode::Pow(_, e) => arena.as_num(*e).is_some_and(|r| r.is_integer()),
_ => false,
})
}
thread_local! {
static ASSUME_REAL_PARAMETERS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
static REAL_PARAMETER_JUMPS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
static PLAIN_LOGS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
pub(crate) struct RealParameterJumps(bool, bool);
impl RealParameterJumps {
pub(crate) fn enable(plain_logs: bool) -> Self {
let active = REAL_PARAMETER_JUMPS.with(|f| f.replace(true));
let plain = PLAIN_LOGS.with(std::cell::Cell::get);
PLAIN_LOGS.with(|f| f.set(plain_logs || (active && plain)));
Self(active, plain)
}
}
impl Drop for RealParameterJumps {
fn drop(&mut self) {
let _ = REAL_PARAMETER_JUMPS.try_with(|f| f.set(self.0));
let _ = PLAIN_LOGS.try_with(|f| f.set(self.1));
}
}
pub(crate) fn strict_jumps_active() -> bool {
STRICT_JUMPS.with(std::cell::Cell::get)
}
thread_local! {
static STRICT_JUMPS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
pub(crate) struct StrictJumps(bool);
impl StrictJumps {
pub(crate) fn enable() -> Self {
Self(STRICT_JUMPS.with(|f| f.replace(true)))
}
}
impl Drop for StrictJumps {
fn drop(&mut self) {
let _ = STRICT_JUMPS.try_with(|f| f.set(self.0));
}
}
pub(crate) fn plain_logs(arena: &mut Arena, g: ExprId, t_sym: SymbolId, all: bool) -> ExprId {
let mut out = g;
for id in crate::base::walk::post_order_ids(arena, g) {
if let ExprNode::Ln(inner) = *arena.node(id)
&& let ExprNode::Abs(p) = *arena.node(inner)
&& contains_var(arena, p, t_sym)
&& (all
|| crate::base::walk::free_symbols(arena, p)
.into_iter()
.any(|s| match *arena.node(s) {
ExprNode::Symbol(sid) => {
sid != t_sym
&& !crate::transforms::realness::symbol_declared_real(arena, sid)
}
_ => true,
}))
{
let plain = arena.ln(p);
out = arena.subs_structural(out, id, plain);
}
}
out
}
pub(crate) fn infinity_jump_real_parameters(
arena: &mut Arena,
g: ExprId,
t: ExprId,
t_sym: SymbolId,
) -> Option<ExprId> {
let key = JumpKey {
arena: std::ptr::from_ref::<Arena>(arena) as usize,
real_parameters: true,
assumed_real: true,
g,
t,
};
memoized_jump(key, || {
infinity_jump_real_parameters_uncached(arena, g, t, t_sym)
})
}
fn infinity_jump_real_parameters_uncached(
arena: &mut Arena,
g: ExprId,
t: ExprId,
t_sym: SymbolId,
) -> Option<ExprId> {
struct Restore(bool);
impl Drop for Restore {
fn drop(&mut self) {
let _ = ASSUME_REAL_PARAMETERS.try_with(|f| f.set(self.0));
}
}
let restore = Restore(ASSUME_REAL_PARAMETERS.with(|f| f.replace(true)));
let plus = end_value(arena, g, t, t_sym, true);
let minus = plus
.is_some()
.then(|| end_value(arena, g, t, t_sym, false))
.flatten();
drop(restore);
match (plus?, minus?) {
(EndValue::Finite(p), EndValue::Finite(m)) => {
let j = arena.sub(p, m);
let j = crate::transforms::eval::eval(arena, j);
Some(combined_signs(arena, j))
}
_ => Some(arena.zero),
}
}
fn contains_var_id(arena: &Arena, e: ExprId, t: ExprId) -> bool {
match *arena.node(t) {
ExprNode::Symbol(sid) => contains_var(arena, e, sid),
_ => true,
}
}
pub(crate) fn end_value(
arena: &mut Arena,
g: ExprId,
t: ExprId,
t_sym: SymbolId,
positive: bool,
) -> Option<EndValue> {
let terms = additive_terms(arena, g, t_sym);
let mut rational: Vec<ExprId> = Vec::new();
let mut finite: Vec<ExprId> = Vec::new();
let mut log_rate: Vec<ExprId> = Vec::new();
for term in terms {
if crate::base::walk::post_order_ids(arena, term).len() > MAX_END_TERM_NODES {
tracing::debug!("end_value: a term too large to follow to the end");
return None;
}
if !contains_var(arena, term, t_sym) {
finite.push(term);
continue;
}
let (c, h) = match arena.node(term).clone() {
ExprNode::Mul(children) => {
let (dep, consts): (Vec<ExprId>, Vec<ExprId>) = children
.iter()
.partition(|&&k| contains_var(arena, k, t_sym));
let c = arena.mul(&consts);
let h = if dep.len() == 1 {
dep[0]
} else {
arena.mul(&dep)
};
(c, h)
}
ExprNode::Neg(inner) => (arena.int(-1), inner),
_ => (arena.one, term),
};
match arena.node(h).clone() {
ExprNode::Ln(arg) => {
let (p, absolute) = match arena.node(arg) {
ExprNode::Abs(inner) => (*inner, true),
_ => (arg, false),
};
let (d, lead) = rational_end(arena, p, t, positive)?;
if d != 0 {
let d_id = arena.int(d);
log_rate.push(arena.mul(&[c, d_id]));
}
let value = if absolute {
let a = arena.abs(lead);
arena.ln(a)
} else {
log_end(arena, p, t, lead, positive)?
};
finite.push(arena.mul(&[c, value]));
}
ExprNode::Atan(q) => {
let (d, lead) = rational_end(arena, q, t, positive)?;
let value = match d.cmp(&0) {
std::cmp::Ordering::Greater => {
let pi = arena.pi();
match placement(arena, lead) {
Some(Placement::Complex { re_sign: 0, .. }) => return None,
Some(found) => {
let sign = match found {
Placement::Positive => 1,
Placement::Negative => -1,
Placement::Complex { re_sign, .. } => i64::from(re_sign),
};
let half = arena.rational(sign, 2);
arena.mul(&[half, pi])
}
None if only_real_parameters(arena, lead)
&& (crate::transforms::realness::constant_realness_as_declared(
arena,
lead,
30,
&mut crate::base::assumptions::AssumptionCache::new(),
) == Some(true)
|| (ASSUME_REAL_PARAMETERS.with(std::cell::Cell::get)
&& rational_in_parameters(arena, lead))) =>
{
let half = arena.rational(1, 2);
let sign = sign_of_real(arena, lead);
arena.mul(&[half, pi, sign])
}
None => return None,
}
}
std::cmp::Ordering::Equal => match placement(arena, lead) {
Some(Placement::Complex {
re_sign: 0,
im_at_least_one: true,
}) => return None,
Some(_) => arena.atan(lead),
None if ASSUME_REAL_PARAMETERS.with(std::cell::Cell::get)
&& rational_in_parameters(arena, lead) =>
{
arena.atan(lead)
}
None => return None,
},
std::cmp::Ordering::Less => arena.zero,
};
finite.push(arena.mul(&[c, value]));
}
ExprNode::RootSum(p, body, s) => {
let (rate, value) = root_sum_end(arena, p, body, s, t, t_sym, positive)?;
log_rate.push(arena.mul(&[c, rate]));
finite.push(arena.mul(&[c, value]));
}
_ if rational_end(arena, h, t, positive).is_some() => rational.push(term),
_ => {
let point = if positive {
arena.infinity()
} else {
arena.neg_infinity()
};
let v = arena.limit_expr(term, t, point).ok()?;
let v = crate::transforms::eval::eval(arena, v);
if contains_var(arena, v, t_sym)
|| crate::base::walk::has_unevaluated(arena, v)
|| crate::transforms::evalf::evalf_complex64(arena, v)
.ok()
.is_none_or(|z| !z.re.is_finite() || !z.im.is_finite())
{
return None;
}
finite.push(v);
}
}
}
if !rational.is_empty() {
let r = arena.add(&rational);
let (d, lead) = rational_end(arena, r, t, positive)?;
match d.cmp(&0) {
std::cmp::Ordering::Greater => return Some(EndValue::Infinite),
std::cmp::Ordering::Equal => finite.push(lead),
std::cmp::Ordering::Less => {}
}
}
if !log_rate.is_empty() {
let rate = arena.add(&log_rate);
let rate = crate::transforms::eval::eval(arena, rate);
match crate::poly::algebraic::is_zero_checked(arena, rate) {
Some(true) => {}
Some(false) => return Some(EndValue::Infinite),
None if rational_function_is_zero(arena, rate) => {}
None => return None,
}
}
let value = arena.add(&finite);
Some(EndValue::Finite(crate::transforms::eval::eval(
arena, value,
)))
}
const MAX_END_TERM_NODES: usize = 2_000;
fn additive_terms(arena: &mut Arena, g: ExprId, t_sym: SymbolId) -> Vec<ExprId> {
let mut out: Vec<ExprId> = Vec::new();
let mut work: Vec<ExprId> = vec![g];
while let Some(e) = work.pop() {
match arena.node(e).clone() {
ExprNode::Add(children)
if out.len() + work.len() + children.len() <= MAX_DISTRIBUTED_TERMS =>
{
work.extend(children.iter().copied());
}
ExprNode::Mul(children) => {
let dependent: Vec<usize> = (0..children.len())
.filter(|&k| contains_var(arena, children[k], t_sym))
.collect();
let sum = match dependent.as_slice() {
[k] => match arena.node(children[*k]) {
ExprNode::Add(parts) => Some((*k, parts.to_vec())),
_ => None,
},
_ => None,
};
match sum {
Some((k, parts))
if out.len() + work.len() + parts.len() <= MAX_DISTRIBUTED_TERMS =>
{
let mut rest: Vec<ExprId> = children.to_vec();
rest.remove(k);
let c = arena.mul(&rest);
for part in parts {
work.push(arena.mul(&[c, part]));
}
}
_ => out.push(e),
}
}
_ => out.push(e),
}
}
out
}
const MAX_DISTRIBUTED_TERMS: usize = 256;
fn root_sum_end(
arena: &mut Arena,
p: ExprId,
body: ExprId,
s: ExprId,
t: ExprId,
t_sym: SymbolId,
positive: bool,
) -> Option<(ExprId, ExprId)> {
let real_roots = match crate::poly::polybridge::expr_to_poly(arena, p, s) {
Some(p_poly) => {
if p_poly.degree().is_none_or(|d| d == 0) {
return None;
}
crate::poly::sturm::SturmChain::new(&p_poly).count_real_roots() != 0
}
None if ASSUME_REAL_PARAMETERS.with(std::cell::Cell::get) => {
let coeffs = crate::transforms::solve::symbolic_poly_coeffs(arena, p, s)?;
if coeffs.len() < 2 || !coeffs.iter().all(|&k| rational_in_parameters(arena, k)) {
return None;
}
true
}
None => return None,
};
let ExprNode::Mul(children) = arena.node(body).clone() else {
return None;
};
let logs: Vec<usize> = (0..children.len())
.filter(|&k| contains_var(arena, children[k], t_sym))
.collect();
let [k] = logs.as_slice() else {
return None;
};
let ExprNode::Ln(arg) = *arena.node(children[*k]) else {
return None;
};
let coeffs = crate::transforms::solve::symbolic_poly_coeffs(arena, arg, t)?;
let [beta, one] = coeffs.as_slice() else {
return None;
};
if *one != arena.one || contains_var(arena, *beta, t_sym) {
return None;
}
let beta = *beta;
let mut rest: Vec<ExprId> = children.to_vec();
rest.remove(*k);
let c = arena.mul(&rest);
let rate = arena.intern(ExprNode::RootSum(p, c, s));
if positive {
return Some((rate, arena.zero));
}
let im_beta = arena.im(beta);
let sign_im = arena.sign(im_beta);
let sum = if real_roots {
let one = arena.one;
let minus = arena.neg(sign_im);
let lower = arena.add(&[one, minus]);
let signed = arena.mul(&[c, sign_im, lower]);
let below = arena.intern(ExprNode::RootSum(p, signed, s));
let sum = arena.add(&[rate, below]);
crate::transforms::eval::eval(arena, sum)
} else {
let signed = arena.mul(&[c, sign_im]);
arena.intern(ExprNode::RootSum(p, signed, s))
};
let i = arena.i_unit();
let pi = arena.pi();
Some((rate, arena.mul(&[i, pi, sum])))
}
pub(crate) fn infinity_jump(
arena: &mut Arena,
g: ExprId,
t: ExprId,
t_sym: SymbolId,
) -> Option<ExprId> {
let key = JumpKey {
arena: std::ptr::from_ref::<Arena>(arena) as usize,
real_parameters: false,
assumed_real: ASSUME_REAL_PARAMETERS.with(std::cell::Cell::get),
g,
t,
};
memoized_jump(key, || infinity_jump_uncached(arena, g, t, t_sym))
}
fn infinity_jump_uncached(
arena: &mut Arena,
g: ExprId,
t: ExprId,
t_sym: SymbolId,
) -> Option<ExprId> {
let undeclared = crate::base::walk::free_symbols(arena, g)
.into_iter()
.any(|s| match *arena.node(s) {
ExprNode::Symbol(sid) => {
sid != t_sym && !crate::transforms::realness::symbol_declared_real(arena, sid)
}
_ => true,
});
if undeclared {
return None;
}
let plus = end_value(arena, g, t, t_sym, true)?;
let minus = end_value(arena, g, t, t_sym, false)?;
match (plus, minus) {
(EndValue::Finite(p), EndValue::Finite(m)) => {
let j = arena.sub(p, m);
let j = crate::transforms::eval::eval(arena, j);
Some(combined_signs(arena, j))
}
_ => Some(arena.zero),
}
}
fn combined_signs(arena: &mut Arena, j: ExprId) -> ExprId {
let has_sign = crate::base::walk::post_order_ids(arena, j)
.iter()
.any(|&id| matches!(arena.node(id), ExprNode::Sign(_)));
if !has_sign {
return j;
}
let expanded = crate::transforms::expand::expand(arena, j);
let expanded = crate::transforms::eval::eval(arena, expanded);
if printed_size(arena, expanded) < printed_size(arena, j) {
expanded
} else {
j
}
}
fn printed_size(arena: &Arena, e: ExprId) -> usize {
let mut size: rustc_hash::FxHashMap<ExprId, usize> = rustc_hash::FxHashMap::default();
for id in crate::base::walk::post_order_ids(arena, e) {
let mut s: usize = 1;
arena
.node(id)
.for_each_child(|c| s = s.saturating_add(size.get(&c).copied().unwrap_or(1)));
size.insert(id, s);
}
size.get(&e).copied().unwrap_or(1)
}
fn sign_of_real(arena: &mut Arena, z: ExprId) -> ExprId {
use crate::base::assumptions::{AssumptionCache, Props};
let factors: Vec<ExprId> = match arena.node(z) {
ExprNode::Mul(children) => children.to_vec(),
_ => vec![z],
};
let mut facts = AssumptionCache::new();
let mut kept: Vec<ExprId> = Vec::new();
let mut negative = false;
for f in factors {
if facts.query(arena, f, Props::POSITIVE) == Some(true) {
continue;
}
if facts.query(arena, f, Props::NEGATIVE) == Some(true) {
negative = !negative;
continue;
}
kept.push(f);
}
let s = if kept.is_empty() {
arena.one
} else {
let core = arena.mul(&kept);
arena.sign(core)
};
if negative { arena.neg(s) } else { s }
}
pub(crate) fn continuous_through_tan(
arena: &mut Arena,
g: ExprId,
t: ExprId,
t_sym: SymbolId,
w: ExprId,
) -> Option<ExprId> {
let g = if REAL_PARAMETER_JUMPS.with(std::cell::Cell::get) {
let all = PLAIN_LOGS.with(std::cell::Cell::get);
plain_logs(arena, g, t_sym, all)
} else {
g
};
let jump = match infinity_jump(arena, g, t, t_sym) {
Some(j) => j,
None if REAL_PARAMETER_JUMPS.with(std::cell::Cell::get) => {
infinity_jump_real_parameters(arena, g, t, t_sym)?
}
None => return None,
};
let tan_w = arena.tan(w);
let back = arena.subs_structural(g, t, tan_w);
Some(add_tan_floor(arena, back, w, jump))
}
pub(crate) fn add_tan_floor(arena: &mut Arena, big_f: ExprId, w: ExprId, jump: ExprId) -> ExprId {
add_tan_floor_in(arena, big_f, w, jump, None)
}
pub(crate) fn add_tan_floor_in(
arena: &mut Arena,
big_f: ExprId,
w: ExprId,
jump: ExprId,
var: Option<SymbolId>,
) -> ExprId {
let pi = arena.pi();
let (mut kept, atan_coeffs) = split_atan_tan_terms(arena, big_f, w, var);
let mut floor_coeff: Vec<ExprId> = vec![jump];
for c in atan_coeffs {
let minus_one = arena.int(-1);
floor_coeff.push(arena.mul(&[minus_one, c, pi]));
}
let coeff = arena.add(&floor_coeff);
let coeff = crate::transforms::eval::eval(arena, coeff);
if crate::poly::algebraic::is_zero_checked(arena, coeff) != Some(true) {
let minus_one = arena.int(-1);
let inv_pi = arena.pow(pi, minus_one);
let w_over_pi = arena.mul(&[w, inv_pi]);
let half = arena.rational(1, 2);
let arg = arena.add(&[w_over_pi, half]);
let fl = arena.floor(arg);
kept.push(arena.mul(&[coeff, fl]));
}
arena.add(&kept)
}
pub(crate) fn atan_tan_terms_to_argument(arena: &mut Arena, big_f: ExprId, w: ExprId) -> ExprId {
atan_tan_terms_to_argument_in(arena, big_f, w, None)
}
pub(crate) fn atan_tan_terms_to_argument_in(
arena: &mut Arena,
big_f: ExprId,
w: ExprId,
var: Option<SymbolId>,
) -> ExprId {
let (kept, atan_coeffs) = split_atan_tan_terms(arena, big_f, w, var);
if atan_coeffs.is_empty() {
return big_f;
}
arena.add(&kept)
}
fn split_atan_tan_terms(
arena: &mut Arena,
big_f: ExprId,
w: ExprId,
var: Option<SymbolId>,
) -> (Vec<ExprId>, Vec<ExprId>) {
let tan_w = arena.tan(w);
let atan_tan_w = arena.atan(tan_w);
if !crate::base::walk::contains(arena, big_f, atan_tan_w) {
let terms = match arena.node(big_f) {
ExprNode::Add(children) => children.to_vec(),
_ => vec![big_f],
};
return (terms, Vec::new());
}
let terms = terms_distributed(arena, big_f, w, var);
let mut kept: Vec<ExprId> = Vec::with_capacity(terms.len() + 1);
let mut coeffs: Vec<ExprId> = Vec::new();
for term in terms {
let c = if term == atan_tan_w {
Some(arena.one)
} else if let ExprNode::Mul(children) = arena.node(term).clone()
&& children.iter().filter(|&&k| k == atan_tan_w).count() == 1
{
let rest: Vec<ExprId> = children
.iter()
.copied()
.filter(|&k| k != atan_tan_w)
.collect();
Some(arena.mul(&rest))
} else {
None
};
match c {
Some(c) if !depends_on_variable(arena, c, w, var) => {
kept.push(arena.mul(&[c, w]));
coeffs.push(c);
}
_ => kept.push(term),
}
}
(kept, coeffs)
}
fn terms_distributed(
arena: &mut Arena,
big_f: ExprId,
w: ExprId,
var: Option<SymbolId>,
) -> Vec<ExprId> {
let mut out: Vec<ExprId> = Vec::new();
let mut work: Vec<ExprId> = vec![big_f];
while let Some(e) = work.pop() {
match arena.node(e).clone() {
ExprNode::Add(children)
if out.len() + work.len() + children.len() <= MAX_DISTRIBUTED_TERMS =>
{
work.extend(children.iter().copied());
}
ExprNode::Mul(children) => {
let dependent: Vec<usize> = (0..children.len())
.filter(|&k| depends_on_variable(arena, children[k], w, var))
.collect();
let sum = match dependent.as_slice() {
[k] => match arena.node(children[*k]) {
ExprNode::Add(parts) => Some((*k, parts.to_vec())),
_ => None,
},
_ => None,
};
match sum {
Some((k, parts))
if out.len() + work.len() + parts.len() <= MAX_DISTRIBUTED_TERMS =>
{
let mut rest: Vec<ExprId> = children.to_vec();
rest.remove(k);
let c = arena.mul(&rest);
for part in parts {
work.push(arena.mul(&[c, part]));
}
}
_ => out.push(e),
}
}
_ => out.push(e),
}
}
out
}
fn depends_on_variable(arena: &Arena, c: ExprId, w: ExprId, var: Option<SymbolId>) -> bool {
match var {
Some(v) => contains_var(arena, c, v),
None => free_symbols_meet(arena, c, w),
}
}
fn free_symbols_meet(arena: &Arena, c: ExprId, w: ExprId) -> bool {
let of_w = crate::base::walk::free_symbols(arena, w);
crate::base::walk::free_symbols(arena, c)
.iter()
.any(|s| of_w.contains(s))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::base::arena::Arena;
fn jump_of(g: &str) -> Option<String> {
let ctx = crate::api::context::Context::new();
let g = ctx.parse(g).unwrap().id();
ctx.with_arena_mut(|a| {
let t = a.symbol("t");
let ExprNode::Symbol(t_sym) = *a.node(t) else {
unreachable!()
};
infinity_jump(a, g, t, t_sym).map(|j| a.display(j).to_string())
})
}
#[test]
fn jumps_at_infinity_of_rational_antiderivatives() {
assert_eq!(
jump_of("2/3*sqrt(3)*atan(1/3*sqrt(3)*t)").as_deref(),
Some("2/3*sqrt(3)*pi")
);
assert_eq!(
jump_of("1/4*sqrt(2)*(2*atan(1/4*sqrt(2)*(t^3 + 7*t)) + 2*atan(1/4*sqrt(2)*t))")
.as_deref(),
Some("sqrt(2)*pi")
);
assert_eq!(
jump_of("-2/5*ln(abs(t + 2)) + 1/5*ln(t^2 + 1) + 1/5*atan(t)").as_deref(),
Some("1/5*pi")
);
let j = jump_of("(-3)^(-1/2)*(-ln(sqrt(3)*I + t) + ln(-sqrt(3)*I + t))").unwrap();
let ctx = crate::api::context::Context::new();
let v = ctx.parse(&j).unwrap().eval_complex64().unwrap();
let want = 2.0 * std::f64::consts::PI / 3f64.sqrt();
assert!(
(v.re - want).abs() < 1e-14 && v.im.abs() < 1e-14,
"{j} = {v}"
);
assert_eq!(jump_of("t").as_deref(), Some("0"));
assert_eq!(jump_of("ln(abs(t))").as_deref(), Some("0"));
assert_eq!(jump_of("atan(t/a)"), None);
}
#[test]
fn jump_of_a_root_sum_over_one_half_plane() {
let mut a = Arena::new();
let t = a.symbol("t");
let s = a.symbol("s");
let ExprNode::Symbol(t_sym) = *a.node(t) else {
unreachable!()
};
let four = a.int(4);
let s4 = a.pow(s, four);
let one = a.one;
let p = a.add(&[s4, one]);
let t_minus_s = a.sub(t, s);
let ln = a.ln(t_minus_s);
let quarter = a.rational(-1, 4);
let body = a.mul(&[quarter, s, ln]);
let g = a.intern(ExprNode::RootSum(p, body, s));
let j = infinity_jump(&mut a, g, t, t_sym).unwrap();
let v = crate::transforms::evalf::evalf_complex64(&a, j).unwrap();
let want = std::f64::consts::PI / 2f64.sqrt();
assert!(
(v.re - want).abs() < 1e-14 && v.im.abs() < 1e-14,
"{} = {v}",
a.display(j)
);
}
fn root_sum_jump(p: &str, q: &str) -> num_complex::Complex64 {
let ctx = crate::api::context::Context::new();
let p = ctx.parse(p).unwrap().id();
let q = ctx.parse(q).unwrap().id();
ctx.with_arena_mut(|a| {
let t = a.symbol("t");
let s = a.symbol("s");
let ExprNode::Symbol(t_sym) = *a.node(t) else {
unreachable!()
};
let t_minus_s = a.sub(t, s);
let ln = a.ln(t_minus_s);
let body = a.div(ln, q);
let g = a.intern(ExprNode::RootSum(p, body, s));
let j = infinity_jump(a, g, t, t_sym).unwrap();
crate::transforms::evalf::evalf_complex64(a, j).unwrap()
})
}
#[test]
fn jump_of_a_root_sum_with_real_roots() {
let j = root_sum_jump("s^4 - s^2 - 2", "4*s^3 - 2*s");
let want = -std::f64::consts::PI / 3.0;
assert!((j.re - want).abs() < 1e-14 && j.im.abs() < 1e-14, "{j}");
let j = root_sum_jump("s^3 - s^2 + s - 1", "3*s^2 - 2*s + 1");
let want = -std::f64::consts::FRAC_PI_2;
assert!(
(j.re - want).abs() < 1e-14 && (j.im - want).abs() < 1e-14,
"{j}"
);
}
#[test]
fn jump_of_a_root_sum_whose_own_rate_is_not_zero() {
let ctx = crate::api::context::Context::new();
let j = ctx.with_arena_mut(|a| {
let t = a.symbol("t");
let s = a.symbol("s");
let ExprNode::Symbol(t_sym) = *a.node(t) else {
unreachable!()
};
let t_minus_s = a.sub(t, s);
let ln = a.ln(t_minus_s);
let two = a.int(2);
let three = a.int(3);
let s2 = a.pow(s, two);
let s3 = a.pow(s, three);
let minus_two = a.int(-2);
let one = a.one;
let cubic = a.add(&[s3, minus_two]);
let quad = a.add(&[s2, one]);
let q1 = a.mul(&[three, s2, quad]);
let body1 = a.div(ln, q1);
let g1 = a.intern(ExprNode::RootSum(cubic, body1, s));
let q2 = a.mul(&[two, s, cubic]);
let body2 = a.div(ln, q2);
let g2 = a.intern(ExprNode::RootSum(quad, body2, s));
let g = a.add(&[g1, g2]);
let j = infinity_jump(a, g, t, t_sym).unwrap();
crate::transforms::evalf::evalf_complex64(a, j).unwrap()
});
assert!(
(j.re + 0.909_316_057_595_450_5).abs() < 1e-13
&& (j.im + 0.254_963_612_453_376_55).abs() < 1e-13,
"{j}"
);
}
#[test]
fn jumps_with_real_parameters_of_undecided_sign() {
let ctx = crate::api::context::Context::new();
let a = ctx
.symbol_with("a", &[crate::base::assumptions::Assumption::Real])
.unwrap();
let jump = |src: &str| -> crate::api::expr::Ex {
let g = ctx.parse(src).unwrap();
let j = ctx.with_arena_mut(|ar| {
let t = ar.symbol("t");
let ExprNode::Symbol(t_sym) = *ar.node(t) else {
unreachable!()
};
infinity_jump(ar, g.id(), t, t_sym)
});
g.wrap(j.unwrap())
};
let j = jump("atan(a*t/sqrt(a^2 + 2))");
assert_eq!(j.to_string(), "sign(a)*pi");
let j = jump("ln(t + sqrt(a)) - ln(t - sqrt(a))");
for (v, want) in [(4, 0.0), (-4, -2.0 * std::f64::consts::PI)] {
let z = j.subs(&a, &ctx.int(v)).eval_complex64().unwrap();
assert!(
z.re.abs() < 1e-14 && (z.im - want).abs() < 1e-14,
"{j} at a = {v}: {z}"
);
}
}
#[test]
fn atan_tan_terms_become_the_argument() {
let ctx = crate::api::context::Context::new();
let f = ctx
.parse("-2/5*ln(abs(tan(x) + 2)) + 1/5*atan(tan(x))")
.unwrap()
.id();
let (with_floor, plain) = ctx.with_arena_mut(|a| {
let x = a.symbol("x");
let pi = a.pi();
let fifth = a.rational(1, 5);
let jump = a.mul(&[fifth, pi]);
let with_floor = add_tan_floor(a, f, x, jump);
let plain = atan_tan_terms_to_argument(a, f, x);
(
a.display(with_floor).to_string(),
a.display(plain).to_string(),
)
});
assert_eq!(with_floor, "1/5*x - 2/5*ln(abs(tan(x) + 2))");
assert_eq!(plain, "1/5*x - 2/5*ln(abs(tan(x) + 2))");
}
#[test]
fn sin_zero() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_pow_integrate(&mut a, 0, x);
assert_eq!(result, x);
}
#[test]
fn sin_first() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_pow_integrate(&mut a, 1, x);
let s = a.display(result).to_string();
assert!(s.contains("cos"), "expected −cos(x), got: {s}");
}
#[test]
fn sin_squared() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_pow_integrate(&mut a, 2, x);
let s = a.display(result).to_string();
assert!(
s.contains("cos") && s.contains("sin"),
"expected reduction formula terms, got: {s}"
);
}
#[test]
fn sin_cubed() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_pow_integrate(&mut a, 3, x);
let s = a.display(result).to_string();
assert!(s.contains("cos"), "expected cos terms, got: {s}");
}
#[test]
fn sin_fourth() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_pow_integrate(&mut a, 4, x);
let s = a.display(result).to_string();
assert!(!s.is_empty(), "got: {s}");
}
#[test]
fn sin_neg2_is_neg_cot() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_pow_integrate(&mut a, -2, x);
let s = a.display(result).to_string();
assert!(
!s.contains("Integral"),
"expected evaluated result, got unevaluated: {s}"
);
assert!(
s.contains("cos") && s.contains("sin"),
"expected cos/sin terms for −cot(x), got: {s}"
);
}
#[test]
fn cos_zero() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = cos_pow_integrate(&mut a, 0, x);
assert_eq!(result, x);
}
#[test]
fn cos_first() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = cos_pow_integrate(&mut a, 1, x);
let s = a.display(result).to_string();
assert!(s.contains("sin"), "expected sin(x), got: {s}");
}
#[test]
fn cos_squared() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = cos_pow_integrate(&mut a, 2, x);
let s = a.display(result).to_string();
assert!(
s.contains("cos") && s.contains("sin"),
"expected reduction formula terms, got: {s}"
);
}
#[test]
fn cos_cubed() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = cos_pow_integrate(&mut a, 3, x);
let s = a.display(result).to_string();
assert!(s.contains("sin"), "expected sin terms, got: {s}");
}
#[test]
fn cos_neg1_is_ln_sec_tan() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = cos_pow_integrate(&mut a, -1, x);
let s = a.display(result).to_string();
assert!(
!s.contains("Integral"),
"expected evaluated result, got unevaluated: {s}"
);
assert!(
s.contains("ln"),
"expected ln term for ln|sec+tan|, got: {s}"
);
}
#[test]
fn cos_neg2_is_tan() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = cos_pow_integrate(&mut a, -2, x);
let s = a.display(result).to_string();
assert!(
!s.contains("Integral"),
"expected evaluated result, got unevaluated: {s}"
);
assert!(
s.contains("sin") && s.contains("cos"),
"expected sin/cos terms for tan(x), got: {s}"
);
}
#[test]
fn sin1_cos1() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_cos_integrate(&mut a, 1, 1, x);
let s = a.display(result).to_string();
assert!(s.contains("cos"), "expected cos term, got: {s}");
}
#[test]
fn sin2_cos1() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_cos_integrate(&mut a, 2, 1, x);
let s = a.display(result).to_string();
assert!(s.contains("sin"), "expected sin term, got: {s}");
}
#[test]
fn sin1_cos2() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_cos_integrate(&mut a, 1, 2, x);
let s = a.display(result).to_string();
assert!(s.contains("cos"), "expected cos term, got: {s}");
}
#[test]
fn sin3_cos2() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_cos_integrate(&mut a, 3, 2, x);
let s = a.display(result).to_string();
assert!(s.contains("cos"), "expected cos terms, got: {s}");
}
#[test]
fn sin2_cos3() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_cos_integrate(&mut a, 2, 3, x);
let s = a.display(result).to_string();
assert!(s.contains("sin"), "expected sin terms, got: {s}");
}
#[test]
fn sin2_cos2_both_even() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_cos_integrate(&mut a, 2, 2, x);
let s = a.display(result).to_string();
assert!(!s.is_empty(), "got: {s}");
}
#[test]
fn sin_cos_m_zero_delegates() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_cos_integrate(&mut a, 0, 3, x);
let direct = cos_pow_integrate(&mut a, 3, x);
assert_eq!(result, direct);
}
#[test]
fn sin_cos_n_zero_delegates() {
let mut a = Arena::new();
let x = a.symbol("x");
let result = sin_cos_integrate(&mut a, 3, 0, x);
let direct = sin_pow_integrate(&mut a, 3, x);
assert_eq!(result, direct);
}
#[test]
fn try_detect_sin_squared() {
let mut a = Arena::new();
let x = a.symbol("x");
let var_sym = match a.node(x) {
ExprNode::Symbol(sid) => *sid,
_ => unreachable!(),
};
let sin_x = a.sin(x);
let two = a.int(2);
let expr = a.pow(sin_x, two);
let result = try_trig_power_integral(&mut a, expr, x, var_sym);
assert!(result.is_some(), "should detect sin²(x)");
let s = a.display(result.unwrap()).to_string();
assert!(
s.contains("cos") && s.contains("sin"),
"expected reduction formula, got: {s}"
);
}
#[test]
fn try_detect_cos_cubed() {
let mut a = Arena::new();
let x = a.symbol("x");
let var_sym = match a.node(x) {
ExprNode::Symbol(sid) => *sid,
_ => unreachable!(),
};
let cos_x = a.cos(x);
let three = a.int(3);
let expr = a.pow(cos_x, three);
let result = try_trig_power_integral(&mut a, expr, x, var_sym);
assert!(result.is_some(), "should detect cos³(x)");
let s = a.display(result.unwrap()).to_string();
assert!(s.contains("sin"), "expected sin terms, got: {s}");
}
#[test]
fn try_detect_sin_cos_product() {
let mut a = Arena::new();
let x = a.symbol("x");
let var_sym = match a.node(x) {
ExprNode::Symbol(sid) => *sid,
_ => unreachable!(),
};
let sin_x = a.sin(x);
let cos_x = a.cos(x);
let two = a.int(2);
let three = a.int(3);
let sin2 = a.pow(sin_x, two);
let cos3 = a.pow(cos_x, three);
let expr = a.mul(&[sin2, cos3]);
let result = try_trig_power_integral(&mut a, expr, x, var_sym);
assert!(result.is_some(), "should detect sin²·cos³");
}
#[test]
fn try_detect_bare_sin_cos_product() {
let mut a = Arena::new();
let x = a.symbol("x");
let var_sym = match a.node(x) {
ExprNode::Symbol(sid) => *sid,
_ => unreachable!(),
};
let sin_x = a.sin(x);
let cos_x = a.cos(x);
let expr = a.mul(&[sin_x, cos_x]);
let result = try_trig_power_integral(&mut a, expr, x, var_sym);
assert!(result.is_some(), "should detect sin(x)·cos(x)");
}
#[test]
fn try_detect_with_constant_factor() {
let mut a = Arena::new();
let x = a.symbol("x");
let var_sym = match a.node(x) {
ExprNode::Symbol(sid) => *sid,
_ => unreachable!(),
};
let sin_x = a.sin(x);
let cos_x = a.cos(x);
let two = a.int(2);
let three_exp = a.int(3);
let sin2 = a.pow(sin_x, two);
let cos3 = a.pow(cos_x, three_exp);
let three = a.int(3);
let expr = a.mul(&[three, sin2, cos3]);
let result = try_trig_power_integral(&mut a, expr, x, var_sym);
assert!(
result.is_some(),
"should detect 3·sin²·cos³ with constant factor"
);
}
#[test]
fn try_returns_none_for_non_trig() {
let mut a = Arena::new();
let x = a.symbol("x");
let var_sym = match a.node(x) {
ExprNode::Symbol(sid) => *sid,
_ => unreachable!(),
};
let two = a.int(2);
let expr = a.pow(x, two);
let result = try_trig_power_integral(&mut a, expr, x, var_sym);
assert!(result.is_none(), "x² is not a trig power integral");
}
#[test]
fn try_returns_none_for_mixed_non_trig() {
let mut a = Arena::new();
let x = a.symbol("x");
let var_sym = match a.node(x) {
ExprNode::Symbol(sid) => *sid,
_ => unreachable!(),
};
let sin_x = a.sin(x);
let two = a.int(2);
let sin2 = a.pow(sin_x, two);
let expr = a.mul(&[x, sin2]);
let result = try_trig_power_integral(&mut a, expr, x, var_sym);
assert!(
result.is_none(),
"x·sin²(x) has a non-trig dependent factor"
);
}
#[test]
fn binom_values() {
assert_eq!(binom(0, 0), BigInt::from(1));
assert_eq!(binom(1, 0), BigInt::from(1));
assert_eq!(binom(1, 1), BigInt::from(1));
assert_eq!(binom(4, 2), BigInt::from(6));
assert_eq!(binom(5, 0), BigInt::from(1));
assert_eq!(binom(5, 5), BigInt::from(1));
assert_eq!(binom(5, 3), BigInt::from(10));
assert_eq!(binom(6, 3), BigInt::from(20));
assert_eq!(binom(3, 5), BigInt::from(0)); }
#[test]
fn extract_sin_power_bare() {
let mut a = Arena::new();
let x = a.symbol("x");
let sin_x = a.sin(x);
assert_eq!(extract_sin_power(&a, sin_x, x), Some(1));
}
#[test]
fn extract_sin_power_squared() {
let mut a = Arena::new();
let x = a.symbol("x");
let sin_x = a.sin(x);
let two = a.int(2);
let sin2 = a.pow(sin_x, two);
assert_eq!(extract_sin_power(&a, sin2, x), Some(2));
}
#[test]
fn extract_cos_power_bare() {
let mut a = Arena::new();
let x = a.symbol("x");
let cos_x = a.cos(x);
assert_eq!(extract_cos_power(&a, cos_x, x), Some(1));
}
#[test]
fn extract_cos_power_cubed() {
let mut a = Arena::new();
let x = a.symbol("x");
let cos_x = a.cos(x);
let three = a.int(3);
let cos3 = a.pow(cos_x, three);
assert_eq!(extract_cos_power(&a, cos3, x), Some(3));
}
#[test]
fn extract_returns_none_for_non_trig() {
let mut a = Arena::new();
let x = a.symbol("x");
let two = a.int(2);
let x2 = a.pow(x, two);
assert_eq!(extract_sin_power(&a, x2, x), None);
assert_eq!(extract_cos_power(&a, x2, x), None);
}
}