use std::cell::Cell;
use num_bigint::BigInt;
use num_traits::Signed;
use crate::base::arena::Arena;
use crate::base::budget::{BudgetKind, Work};
use crate::base::node::{ExprId, ExprNode};
use crate::base::numeric::Q;
use crate::base::truth::{Evidence, Truth, Verdict};
use crate::poly::dense::Poly;
use crate::poly::generic::GenPoly;
use crate::poly::multipoly::{GrevLex, MultiPoly};
use crate::poly::paramfield::{ParamPolicy, PolyFraction};
use crate::poly::traits::{Field, Ring};
type MP = MultiPoly<GrevLex>;
type GP = GenPoly<PFrac>;
const MAX_PARAMS: usize = 4;
const MAX_ALL_PARAMS: usize = 8;
const MAX_EXPONENT: i64 = 48;
const MAX_DENOM_DEGREE: usize = 10;
const MAX_NUMER_DEGREE: usize = 24;
const OP_BUDGET: u64 = 60_000;
const MAX_COEFF_TERMS: usize = 400;
#[derive(Clone, Copy)]
struct AttemptBudget {
ops: Work,
ops_spent: bool,
gcds: Work,
abandoned: bool,
}
impl AttemptBudget {
const fn fresh() -> Self {
AttemptBudget {
ops: Work::new(BudgetKind::ParametricRational, OP_BUDGET),
ops_spent: false,
gcds: Work::new(BudgetKind::ParametricRational, u64::MAX),
abandoned: false,
}
}
const fn factored(&self) -> bool {
self.gcds.cap() != u64::MAX
}
}
thread_local! {
static BUDGET: Cell<AttemptBudget> = const { Cell::new(AttemptBudget::fresh()) };
}
fn with_budget<R>(f: impl FnOnce(&mut AttemptBudget) -> R) -> R {
BUDGET.with(|c| {
let mut b = c.get();
let r = f(&mut b);
c.set(b);
r
})
}
fn reset_budget() {
BUDGET.with(|c| c.set(AttemptBudget::fresh()));
}
fn budget_exceeded() -> bool {
let b = BUDGET.with(Cell::get);
b.ops_spent || b.abandoned
}
fn abandoned() -> bool {
BUDGET.with(Cell::get).abandoned
}
fn abandon() {
with_budget(|b| b.abandoned = true);
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct Budgeted;
impl ParamPolicy for Budgeted {
fn gcd(a: &MP, b: &MP) -> MP {
gcd_mp(a, b)
}
fn charge() {
let (over, factored) = with_budget(|b| {
if b.ops.charge(1, "param_rational").is_err() {
b.ops_spent = true;
}
(b.ops_spent, b.factored())
});
if over && factored {
abandon();
}
}
fn admit(num: &MP, den: &MP) -> bool {
let factored = BUDGET.with(Cell::get).factored();
if abandoned()
|| num.num_terms() + den.num_terms() > MAX_COEFF_TERMS
|| (factored
&& (max_coeff_bits(num) > MAX_FACTORED_COEFF_BITS
|| max_coeff_bits(den) > MAX_FACTORED_COEFF_BITS))
{
abandon();
return false;
}
true
}
fn abandoned() -> bool {
abandoned()
}
}
pub(crate) type PFrac = crate::poly::paramfield::PFrac<Budgeted>;
const GCD_COST_LIMIT: u64 = 200_000;
fn gcd_mp(a: &MP, b: &MP) -> MP {
let one = MP::from_int(a.num_vars(), 1);
if a.as_constant().is_some() || b.as_constant().is_some() || abandoned() {
return one;
}
if !gcd_affordable(a, b) {
abandon();
return one;
}
let cost = gcd_cost(a, b);
if with_budget(|w| w.gcds.charge(cost, "param_rational").is_err()) {
abandon();
return one;
}
let t = std::time::Instant::now();
let g = MP::gcd(a, b);
let ms = t.elapsed().as_millis();
if ms > 50 {
let degs = |p: &MP| {
(0..p.num_vars())
.map(|v| p.degree_in(v))
.collect::<Vec<_>>()
};
tracing::debug!(ms, a_degs = ?degs(a), b_degs = ?degs(b), a_tot = ?a.total_degree(), b_tot = ?b.total_degree(), "param_rational: slow gcd");
}
g
}
fn gcd_affordable(a: &MP, b: &MP) -> bool {
gcd_cost(a, b) <= GCD_COST_LIMIT
}
fn max_coeff_bits(p: &MP) -> u64 {
p.terms()
.map(|(_, c)| c.numer().bits().max(c.denom().bits()))
.max()
.unwrap_or(0)
}
const MAX_FACTORED_COEFF_BITS: u64 = 2048;
fn gcd_cost(a: &MP, b: &MP) -> u64 {
let mut cost: u64 = 1;
for v in 0..a.num_vars() {
let m = u64::from(a.degree_in(v).min(b.degree_in(v)));
cost = cost.saturating_mul(m + 1);
}
cost
}
fn deriv(p: &GP) -> GP {
let coeffs = p
.coeffs()
.iter()
.enumerate()
.skip(1)
.map(|(k, c)| c.scale_q(&Q::from_integer(BigInt::from(k))))
.collect();
GP::from_coeffs(coeffs)
}
fn div_exact(a: &GP, b: &GP) -> Option<GP> {
let (q, r) = a.try_div_rem(b)?;
r.is_zero().then_some(q)
}
fn to_q_poly(p: &GP) -> Option<Poly> {
let coeffs = p
.coeffs()
.iter()
.map(|c| c.as_const().cloned())
.collect::<Option<Vec<Q>>>()?;
Some(Poly::from_coeffs(coeffs))
}
fn from_q_poly(p: &Poly) -> GP {
GP::from_coeffs(p.coeffs().iter().cloned().map(PFrac::Const).collect())
}
fn mp_sqrt(p: &MP) -> Option<MP> {
let nv = p.num_vars();
let (lt_exp, lt_c) = p.leading_term()?;
if lt_exp.iter().any(|e| e % 2 == 1) {
return None;
}
let root_c = q_sqrt(lt_c)?;
let lead_exp: Vec<u32> = lt_exp.iter().map(|e| e / 2).collect();
let two_lead = MP::monomial(&root_c + &root_c, lead_exp.clone());
let mut s = MP::monomial(root_c, lead_exp);
for _ in 0..64 + 4 * p.num_terms() {
let r = p.sub(&s.mul(&s));
let Some((r_exp, r_c)) = r.leading_term() else {
return Some(s);
};
let (d_exp, d_c) = two_lead.leading_term()?;
if r_exp.iter().zip(d_exp).any(|(a, b)| a < b) {
return None;
}
let t_exp: Vec<u32> = r_exp.iter().zip(d_exp).map(|(a, b)| a - b).collect();
if t_exp.iter().zip(d_exp).all(|(a, b)| a == b) {
return None;
}
let t = MP::monomial(r_c / d_c, t_exp);
s = s.add(&t);
if s.num_vars() != nv {
return None;
}
}
None
}
pub(crate) fn polynomial_sqrt(arena: &mut Arena, e: ExprId) -> Option<ExprId> {
if arena.is_zero_structural(e) {
return Some(arena.zero());
}
let mut vars: Vec<ExprId> = crate::base::walk::free_symbols(arena, e)
.into_iter()
.collect();
vars.sort_by(|&a, &b| arena.canonical_cmp(a, b));
if vars.is_empty() {
let q = arena.as_num(e)?.clone();
let r = q_sqrt(&q)?;
return Some(arena.num_ratio(r));
}
let p = crate::poly::polybridge::expr_to_multipoly(arena, e, &vars)?;
let s = mp_sqrt(&p)?;
Some(crate::poly::polybridge::multipoly_to_expr(arena, &s, &vars))
}
fn factor_split(p: &MP) -> Option<Vec<(MP, u32)>> {
let nv = p.num_vars();
let mono = p.monomial_content();
let mut rest = p.div_exact(&MP::monomial(Q::one(), mono.clone()))?;
let mut basis: Vec<MP> = Vec::new();
if rest.as_constant().is_none() {
basis.push(rest.clone());
}
for _ in 0..16 {
let mut split: Option<(usize, MP)> = None;
'find: for (i, f) in basis.iter().enumerate() {
let deg_f = f.total_degree()?;
for v in f.variables_present() {
let g = MP::gcd(f, &f.partial_derivative(v));
if g.total_degree().is_some_and(|d| d > 0 && d < deg_f) {
split = Some((i, g));
break 'find;
}
}
for (j, h) in basis.iter().enumerate() {
if i != j {
let g = MP::gcd(f, h);
if g.total_degree().is_some_and(|d| d > 0 && d < deg_f) {
split = Some((i, g));
break 'find;
}
}
}
}
let Some((i, g)) = split else { break };
let f = basis.swap_remove(i);
let cof = f.div_exact(&g)?;
for piece in [g, cof] {
if piece.as_constant().is_none() && !basis.contains(&piece) {
basis.push(piece);
}
}
}
let mut factors: Vec<(MP, u32)> = Vec::new();
for (v, &e) in mono.iter().enumerate() {
if e > 0 {
factors.push((MP::var(nv, v), e));
}
}
for b in basis {
let mut k = 0u32;
while rest.as_constant().is_none()
&& let Some(q) = rest.div_exact(&b)
{
rest = q;
k += 1;
}
if k > 0 {
factors.push((b, k));
}
}
if factors.len() < 2 && factors.iter().all(|(_, k)| *k == 1) {
return None;
}
let c = rest.as_constant()?;
if !c.is_one() {
return None;
}
let product = factors
.iter()
.fold(MP::from_int(nv, 1), |acc, (s, k)| acc.mul(&s.pow(*k)));
(product == *p).then_some(factors)
}
fn q_sqrt(q: &Q) -> Option<Q> {
if q.is_negative() {
return None;
}
let n = int_sqrt(q.numer())?;
let d = int_sqrt(q.denom())?;
Some(Q::new(n, d))
}
fn int_sqrt(n: &BigInt) -> Option<BigInt> {
let r = num_integer::Roots::sqrt(n);
(&r * &r == *n).then_some(r)
}
fn split_square(delta: &PFrac, nv: usize) -> (PFrac, MP) {
let PolyFraction { numer: n, denom: d } = delta.parts(nv);
let p = n.mul(&d);
let mut sigma = PFrac::canon(MP::from_int(nv, 1), d);
let (lcm, pz) = p.clear_denominators();
let content = pz.integer_content();
if content.is_zero() {
return (PFrac::Const(Q::zero()), MP::from_int(nv, 1));
}
let negative = pz.leading_coeff().is_some_and(Signed::is_negative);
let mut p1 = pz.scale(&Q::new(BigInt::one(), content.clone()));
if negative {
p1 = p1.neg();
}
let mut w = &content * &lcm;
let mut s = BigInt::one();
for prime in [2u32, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37] {
let sq = BigInt::from(prime * prime);
while (&w % &sq).is_zero() {
w /= &sq;
s *= prime;
}
}
if let Some(r) = int_sqrt(&w) {
s *= r;
w = BigInt::one();
}
sigma = sigma.scale_q(&Q::new(s, lcm));
let mono = p1.monomial_content();
if mono.iter().any(|&e| e > 0) {
let half: Vec<u32> = mono.iter().map(|e| e / 2).collect();
let odd: Vec<u32> = mono.iter().map(|e| e % 2).collect();
if let Some(rest) = p1.div_exact(&MP::monomial(Q::one(), mono.clone())) {
p1 = rest.mul(&MP::monomial(Q::one(), odd));
sigma = sigma.mul(&PFrac::from_mp(MP::monomial(Q::one(), half)));
}
}
if let Some(root) = mp_sqrt(&p1) {
sigma = sigma.mul(&PFrac::from_mp(root));
p1 = MP::from_int(nv, 1);
}
let sign = if negative { -Q::one() } else { Q::one() };
let rho = p1.scale(&Q::from_integer(w)).scale(&sign);
(sigma, rho)
}
struct Ctx {
var: ExprId,
params: Vec<ExprId>,
}
impl Ctx {
fn nv(&self) -> usize {
self.params.len()
}
fn to_gp(&self, p: &MP) -> Option<GP> {
let nv = self.nv();
let deg = p.degree_in(0) as usize;
let mut buckets: Vec<Vec<(Vec<u32>, Q)>> = vec![Vec::new(); deg + 1];
for (e, c) in p.terms() {
buckets[e[0] as usize].push((e[1..].to_vec(), c.clone()));
}
let mut coeffs = Vec::with_capacity(deg + 1);
for terms in buckets {
let mp = if terms.is_empty() {
MP::zero(nv)
} else {
MP::from_terms(nv, terms)?
};
coeffs.push(PFrac::from_mp(mp));
}
Some(GP::from_coeffs(coeffs))
}
fn to_mp(&self, p: &GP) -> Option<MP> {
let nv = self.nv();
let mut l = MP::from_int(nv, 1);
for c in p.coeffs() {
if let PFrac::Frac(f) = c
&& f.1.as_constant().is_none()
{
l = MP::lcm(&l, &f.1);
}
}
let mut terms: Vec<(Vec<u32>, Q)> = Vec::new();
for (k, c) in p.coeffs().iter().enumerate() {
if c.is_zero() {
continue;
}
let PolyFraction { numer: n, denom: d } = c.parts(nv);
let scaled = n.mul(&l.div_exact(&d)?);
let k = u32::try_from(k).ok()?;
for (e, q) in scaled.terms() {
let mut exps = Vec::with_capacity(nv + 1);
exps.push(k);
exps.extend_from_slice(e);
terms.push((exps, q.clone()));
}
}
MP::from_terms(nv + 1, terms)
}
fn gcd(&self, a: &GP, b: &GP) -> Option<GP> {
if a.is_zero() {
return Some(b.make_monic());
}
if b.is_zero() {
return Some(a.make_monic());
}
let (am, bm) = (self.to_mp(a)?, self.to_mp(b)?);
let cost = gcd_cost(&am, &bm);
if cost > GCD_COST_LIMIT || with_budget(|w| w.gcds.charge(cost, "param_rational").is_err())
{
abandon();
return None;
}
let g = MP::gcd(&am, &bm);
Some(self.to_gp(&g)?.make_monic())
}
fn coprime(&self, a: &GP, b: &GP) -> Option<Verdict<Truth>> {
let (am, bm) = (self.to_mp(a)?, self.to_mp(b)?);
let (az, bz) = (am.clear_denominators().1, bm.clear_denominators().1);
if !az.is_zero() && !bz.is_zero() {
let bound = crate::poly::modgcd::main_gcd_degree_bound(
self.nv() + 1,
&integer_terms(&az),
&integer_terms(&bz),
);
match bound {
Some(0) => return Some(Verdict::new(Truth::True, Evidence::Exact)),
Some(_) => return Some(Verdict::new(Truth::False, Evidence::Probabilistic)),
None => {}
}
}
let g = self.gcd(a, b)?;
Some(Verdict::new(
Truth::from_bool(g.degree()? == 0),
Evidence::Exact,
))
}
fn mp_expr(&self, arena: &mut Arena, p: &MP) -> ExprId {
crate::poly::polybridge::multipoly_to_expr(arena, p, &self.params)
}
fn coeff_expr(&self, arena: &mut Arena, c: &PFrac) -> ExprId {
match c {
PFrac::Const(q) => arena.num_ratio(q.clone()),
PFrac::Frac(f) => {
let n = self.mp_expr(arena, &f.0);
if f.1.as_constant().is_some_and(|d| d.is_one()) {
n
} else {
let d = match factor_split(&f.1) {
Some(factors) => {
let mut ids = Vec::with_capacity(factors.len());
for (s, k) in factors {
let s_id = self.mp_expr(arena, &s);
ids.push(if k == 1 {
s_id
} else {
let k_id = arena.int(i64::from(k));
arena.pow(s_id, k_id)
});
}
arena.mul(&ids)
}
None => self.mp_expr(arena, &f.1),
};
arena.div(n, d)
}
}
}
}
fn poly_expr(&self, arena: &mut Arena, p: &GP) -> ExprId {
let mut terms = Vec::new();
for (k, c) in p.coeffs().iter().enumerate() {
if c.is_zero() {
continue;
}
let c_id = self.coeff_expr(arena, c);
let x_k = self.x_pow(arena, k);
terms.push(arena.mul(&[c_id, x_k]));
}
match terms.len() {
0 => arena.zero(),
1 => terms[0],
_ => arena.add(&terms),
}
}
fn x_pow(&self, arena: &mut Arena, k: usize) -> ExprId {
match k {
0 => arena.one(),
1 => self.var,
_ => {
let k_id = arena.int(i64::try_from(k).unwrap_or(i64::MAX));
arena.pow(self.var, k_id)
}
}
}
fn declared_real(&self, arena: &Arena, p: &GP) -> bool {
let nv = self.nv();
p.coeffs().iter().all(|c| {
let PolyFraction { numer: n, denom: d } = c.parts(nv);
let mut vars = n.variables_present();
vars.extend(d.variables_present());
vars.into_iter().all(|i| match arena.node(self.params[i]) {
ExprNode::Symbol(s) => crate::transforms::realness::symbol_declared_real(arena, *s),
_ => false,
})
})
}
}
fn parameters(arena: &Arena, expr: ExprId, var: ExprId) -> Option<Vec<ExprId>> {
let mut params = Vec::new();
let mut has_var = false;
for id in crate::base::walk::post_order_ids(arena, expr) {
match arena.node(id) {
ExprNode::Num(_) | ExprNode::Add(_) | ExprNode::Mul(_) | ExprNode::Neg(_) => {}
ExprNode::Pow(_, exp) => {
let ok = arena.as_num(*exp).is_some_and(|r| {
r.is_integer()
&& i64::try_from(r.to_integer()).is_ok_and(|n| n.abs() <= MAX_EXPONENT)
});
if !ok {
return None;
}
}
ExprNode::Symbol(_) if id == var => has_var = true,
ExprNode::Symbol(_) => params.push(id),
_ => return None,
}
}
params.sort_by(|&a, &b| arena.canonical_cmp(a, b));
params.dedup();
(has_var && !params.is_empty() && params.len() <= MAX_ALL_PARAMS).then_some(params)
}
fn integer_terms(p: &MP) -> Vec<(&[u32], &BigInt)> {
p.terms().map(|(e, c)| (e, c.numer())).collect()
}
fn numer_denom(arena: &mut Arena, expr: ExprId) -> (ExprId, ExprId) {
crate::poly::polybridge::fraction_parts(arena, expr)
}
fn syntactic_factors(arena: &Arena, d: ExprId, ctx: &Ctx) -> Vec<GP> {
let mut vars = vec![ctx.var];
vars.extend(&ctx.params);
let factors: Vec<ExprId> = match arena.node(d) {
ExprNode::Mul(ch) => ch.to_vec(),
_ => vec![d],
};
let mut out = Vec::new();
for f in factors {
let base = match arena.node(f) {
ExprNode::Pow(b, _) => *b,
_ => f,
};
if let Some(mp) = crate::poly::polybridge::expr_to_multipoly(arena, base, &vars)
&& let Some(p) = ctx.to_gp(&mp)
&& p.degree().is_some_and(|k| k >= 1)
{
out.push(p.make_monic());
}
}
out
}
pub(crate) fn integrate_param_rational(
arena: &mut Arena,
expr: ExprId,
var: ExprId,
) -> Option<ExprId> {
if matches!(arena.node(expr), ExprNode::Add(_)) {
return None;
}
let params = parameters(arena, expr, var)?;
tracing::debug!(integrand = %arena.display(expr), "param_rational: attempt");
reset_budget();
let (n_id, d_id) = numer_denom(arena, expr);
if d_id == arena.one() {
return None;
}
if params.len() > MAX_PARAMS
&& params
.iter()
.filter(|&&p| crate::base::walk::contains(arena, d_id, p))
.count()
> MAX_PARAMS
{
return None;
}
let ctx = Ctx { var, params };
let mut vars = vec![var];
vars.extend(&ctx.params);
let n_mp = crate::poly::polybridge::expr_to_multipoly(arena, n_id, &vars)?;
let d_mp = crate::poly::polybridge::expr_to_multipoly(arena, d_id, &vars)?;
let mut a = ctx.to_gp(&n_mp)?;
let mut d = ctx.to_gp(&d_mp)?;
if d.degree().is_none_or(|k| k == 0) || a.is_zero() {
return None;
}
if (a.degree().is_some_and(|k| k > MAX_NUMER_DEGREE)
|| d.degree().is_some_and(|k| k > MAX_DENOM_DEGREE))
&& let Some(result) = integrate_factored(arena, &ctx, &a, &d, d_id, expr)
{
return Some(result);
}
if a.degree().is_some_and(|k| k > MAX_NUMER_DEGREE)
|| d.degree().is_some_and(|k| k > 2 * MAX_DENOM_DEGREE)
{
return None;
}
let numerator_param_free = a.coeffs().iter().all(|c| c.as_const().is_some());
let g = ctx.gcd(&a, &d)?;
if g.degree().is_some_and(|k| k > 0) {
a = div_exact(&a, &g)?;
d = div_exact(&d, &g)?;
}
if d.degree()? > MAX_DENOM_DEGREE {
return None;
}
let lc = d.leading_coeff()?.clone();
let inv = lc.inv();
a = a.scale(&inv);
d = d.scale(&inv);
if d.degree()? == 0 {
let big_q = integrate_poly(&a);
return Some(ctx.poly_expr(arena, &big_q));
}
let cands = syntactic_factors(arena, d_id, &ctx);
let d_sqf = div_exact(&d, &ctx.gcd(&d, &deriv(&d))?)?;
if d_sqf.degree()? < 2
&& numerator_param_free
&& cands.len() <= 1
&& cands.iter().all(|c| c.degree() == Some(1))
&& !has_nested_fraction(arena, expr)
{
return None;
}
if let Some(k) = power_substitution_degree(&a, &d) {
return integrate_substituted(arena, &ctx, &a, &d, &cands, k, expr);
}
let result = integrate_gp(arena, &ctx, &a, &d, &cands, expr)?;
Some(crate::transforms::eval::eval(arena, result))
}
const MAX_FACTORED_DENOM_DEGREE: usize = 40;
const MAX_FACTORED_NUMER_DEGREE: usize = 48;
const MAX_FACTORED_PARAMS: usize = 4;
const FACTORED_GCD_BUDGET: u64 = 300_000;
fn factored_denominator(ctx: &Ctx, d: &GP, cands: &[GP]) -> Option<Vec<(GP, u32)>> {
let mut rest = d.make_monic();
let mut out: Vec<(GP, u32)> = Vec::new();
for c in cands {
if out.iter().any(|(f, _)| f == c) {
continue;
}
let mut k = 0u32;
while rest.degree() >= c.degree()
&& let Some(q) = div_exact(&rest, c)
{
rest = q;
k += 1;
}
if k > 0 {
out.push((c.clone(), k));
}
}
if rest.degree() != Some(0) || out.is_empty() {
return None;
}
let proved = |v: Verdict<Truth>| v.truth_at_least(Evidence::Exact).is_true();
for (i, (f, _)) in out.iter().enumerate() {
if !proved(ctx.coprime(f, &deriv(f))?) {
return None;
}
for (g, _) in &out[i + 1..] {
if !proved(ctx.coprime(f, g)?) {
return None;
}
}
}
Some(out)
}
fn factored_product(fz: &[(GP, u32)], shift: u32) -> GP {
let mut p = GP::one();
for (f, k) in fz {
for _ in shift..*k {
p = p.mul(f);
}
}
p
}
fn integrate_factored(
arena: &mut Arena,
ctx: &Ctx,
a: &GP,
d: &GP,
d_id: ExprId,
label: ExprId,
) -> Option<ExprId> {
if a.degree().is_some_and(|k| k > MAX_FACTORED_NUMER_DEGREE)
|| d.degree().is_some_and(|k| k > MAX_FACTORED_DENOM_DEGREE)
{
return None;
}
let cands = syntactic_factors(arena, d_id, ctx);
let mut fz = factored_denominator(ctx, d, &cands)?;
let mut denominator_params: Vec<usize> = Vec::new();
for (f, _) in &fz {
for v in ctx.to_mp(f)?.variables_present() {
if v > 0 && !denominator_params.contains(&v) {
denominator_params.push(v);
}
}
}
if denominator_params.len() > MAX_FACTORED_PARAMS {
return None;
}
with_budget(|b| {
let spent = b.gcds.spent();
b.gcds = Work::new(BudgetKind::ParametricRational, FACTORED_GCD_BUDGET);
if b.gcds.charge(spent, "param_rational").is_err() {
let _ = b.gcds.charge(b.gcds.left(), "param_rational");
}
});
let lc = d.leading_coeff()?.clone();
let mut a = a.scale(&lc.inv());
for (f, k) in &mut fz {
while *k > 0
&& let Some(q) = div_exact(&a, f)
{
a = q;
*k -= 1;
}
}
fz.retain(|(_, k)| *k > 0);
if fz.is_empty() {
let big_q = integrate_poly(&a);
return Some(ctx.poly_expr(arena, &big_q));
}
let sqf_degree: usize = fz.iter().map(|(f, _)| f.degree().unwrap_or(0)).sum();
if sqf_degree > MAX_DENOM_DEGREE {
return None;
}
let d = factored_product(&fz, 0);
if budget_exceeded() {
return None;
}
let result = integrate_gp_with(arena, ctx, &a, &d, &cands, Some(&fz), label);
tracing::debug!(
ops = BUDGET.with(Cell::get).ops.spent(),
gcd_spent = BUDGET.with(Cell::get).gcds.spent(),
ok = result.is_some(),
"param_rational: factored attempt"
);
let result = result?;
Some(crate::transforms::eval::eval(arena, result))
}
fn hermite_factored(a: &GP, fz: &[(GP, u32)]) -> Option<HermiteParts> {
let mut d_star = GP::one();
for (f, _) in fz {
d_star = d_star.mul(f);
}
let mut a = a.clone();
let mut out = Vec::new();
let mut m = 1u32;
let mut d_minus = factored_product(fz, m);
while d_minus.degree()? > 0 {
if budget_exceeded() {
return None;
}
let dm_p = deriv(&d_minus);
let d_minus2 = factored_product(fz, m + 1);
let d_minus_star = div_exact(&d_minus, &d_minus2)?;
let t = div_exact(&d_star.mul(&dm_p), &d_minus)?.neg();
let (b, c) = solve_bezout(&t, &d_minus_star, &a)?;
let bp = deriv(&b);
a = c.sub(&div_exact(&bp.mul(&d_star), &d_minus_star)?);
if !b.is_zero() {
out.push((b, d_minus.clone()));
}
d_minus = d_minus2;
m += 1;
}
Some((out, a, d_star))
}
fn has_nested_fraction(arena: &Arena, e: ExprId) -> bool {
let mut with_negative_power = rustc_hash::FxHashSet::default();
for id in crate::base::walk::post_order_ids(arena, e) {
let node = arena.node(id);
let here = matches!(node, ExprNode::Pow(_, exp)
if arena.as_num(*exp).is_some_and(|r| r.is_negative()));
let mut below = false;
node.for_each_child(|c| below |= with_negative_power.contains(&c));
if below && matches!(node, ExprNode::Add(_)) {
return true;
}
if here || below {
with_negative_power.insert(id);
}
}
false
}
fn power_substitution_degree(a: &GP, d: &GP) -> Option<usize> {
let mut g = 0usize;
for (e, c) in d.coeffs().iter().enumerate() {
if !c.is_zero() {
g = num_integer::gcd(g, e);
}
}
for (e, c) in a.coeffs().iter().enumerate() {
if !c.is_zero() {
g = num_integer::gcd(g, e + 1);
}
}
(g >= 2).then_some(g)
}
fn compress(p: &GP, k: usize, shift: usize) -> GP {
GP::from_coeffs(p.coeffs().iter().skip(shift).step_by(k).cloned().collect())
}
fn integrate_substituted(
arena: &mut Arena,
ctx: &Ctx,
a: &GP,
d: &GP,
cands: &[GP],
k: usize,
label: ExprId,
) -> Option<ExprId> {
let k_q = Q::from_integer(BigInt::from(k));
let a_u = compress(a, k, k - 1).scale(&PFrac::Const(k_q.recip()));
let d_u = compress(d, k, 0);
let cands_u: Vec<GP> = cands
.iter()
.filter(|c| {
c.coeffs()
.iter()
.enumerate()
.all(|(e, x)| x.is_zero() || e % k == 0)
})
.map(|c| compress(c, k, 0))
.collect();
let u = arena.symbol("__prr_u");
let ctx_u = Ctx {
var: u,
params: ctx.params.clone(),
};
let f_u = integrate_gp(arena, &ctx_u, &a_u, &d_u, &cands_u, label)?;
let k_id = arena.int(i64::try_from(k).ok()?);
let x_k = arena.pow(ctx.var, k_id);
let f_x = crate::transforms::subs::subs(arena, f_u, u, x_k);
Some(crate::transforms::eval::eval(arena, f_x))
}
fn integrate_poly(p: &GP) -> GP {
let mut coeffs = vec![PFrac::zero()];
for (k, c) in p.coeffs().iter().enumerate() {
coeffs.push(c.scale_q(&Q::new(BigInt::one(), BigInt::from(k + 1))));
}
GP::from_coeffs(coeffs)
}
fn integrate_gp(
arena: &mut Arena,
ctx: &Ctx,
a: &GP,
d: &GP,
cands: &[GP],
label: ExprId,
) -> Option<ExprId> {
integrate_gp_with(arena, ctx, a, d, cands, None, label)
}
fn integrate_gp_with(
arena: &mut Arena,
ctx: &Ctx,
a: &GP,
d: &GP,
cands: &[GP],
factored: Option<&[(GP, u32)]>,
label: ExprId,
) -> Option<ExprId> {
let (q, r) = a.try_div_rem(d)?;
let mut terms: Vec<ExprId> = Vec::new();
if !q.is_zero() {
let big_q = integrate_poly(&q);
terms.push(ctx.poly_expr(arena, &big_q));
}
if !r.is_zero() {
let (rational, h, d_star) = match factored {
Some(fz) => {
let parts = hermite_factored(&r, fz)?;
let nv = ctx.nv();
let too_big = parts.1.coeffs().iter().any(|c| {
let PolyFraction { numer: n, denom: d } = c.parts(nv);
n.num_terms() + d.num_terms() > MAX_COEFF_TERMS / 4
|| max_coeff_bits(&n).max(max_coeff_bits(&d)) > MAX_FACTORED_COEFF_BITS / 8
});
if too_big {
return None;
}
parts
}
None => hermite(ctx, &r, d)?,
};
if budget_exceeded() {
return None;
}
let basis = coprime_basis(ctx, &d_star, cands)?;
for (b, dm) in rational {
let (dm_id, lambda) = factored_expr(arena, ctx, &dm, &basis);
let b_id = ctx.poly_expr(arena, &b.scale(&lambda));
terms.push(arena.div(b_id, dm_id));
}
if !h.is_zero() {
terms.extend(log_part(arena, ctx, &h, &d_star, &basis, label)?);
}
}
if budget_exceeded() {
return None;
}
Some(match terms.len() {
0 => arena.zero(),
1 => terms[0],
_ => arena.add(&terms),
})
}
type HermiteParts = (Vec<(GP, GP)>, GP, GP);
fn hermite(ctx: &Ctx, a: &GP, d: &GP) -> Option<HermiteParts> {
let mut d_minus = ctx.gcd(d, &deriv(d))?;
let d_star = div_exact(d, &d_minus)?;
let mut a = a.clone();
let mut out = Vec::new();
while d_minus.degree()? > 0 {
if budget_exceeded() {
return None;
}
let dm_p = deriv(&d_minus);
let d_minus2 = ctx.gcd(&d_minus, &dm_p)?;
let d_minus_star = div_exact(&d_minus, &d_minus2)?;
let t = div_exact(&d_star.mul(&dm_p), &d_minus)?.neg();
let (b, c) = solve_bezout(&t, &d_minus_star, &a)?;
let bp = deriv(&b);
a = c.sub(&div_exact(&bp.mul(&d_star), &d_minus_star)?);
if !b.is_zero() {
out.push((b, d_minus.clone()));
}
d_minus = d_minus2;
}
Some((out, a, d_star))
}
fn solve_bezout(f: &GP, g: &GP, c: &GP) -> Option<(GP, GP)> {
let e = GP::extended_gcd(f, g);
if e.gcd.degree()? != 0 {
return None;
}
let inv = e.gcd.leading_coeff()?.inv();
let s0 = e.x.scale(&inv).mul(c);
let t0 = e.y.scale(&inv).mul(c);
let (q, s) = s0.try_div_rem(g)?;
let t = t0.add(&q.mul(f));
Some((s, t))
}
const MAX_SPLIT_TERMS: usize = 40;
fn split_parametric(ctx: &Ctx, e: &GP) -> Option<Vec<GP>> {
let mp = ctx.to_mp(e)?;
if mp.num_terms() > MAX_SPLIT_TERMS {
return None;
}
let (_, factors) = crate::poly::factor_zassenhaus::factor_multivariate(&mp)?;
let pieces = factors
.iter()
.filter(|(f, _)| f.degree_in(0) > 0)
.map(|(f, _)| ctx.to_gp(f).map(|g| g.make_monic()))
.collect::<Option<Vec<GP>>>()?;
(pieces.len() >= 2).then_some(pieces)
}
fn coprime_basis(ctx: &Ctx, d: &GP, cands: &[GP]) -> Option<Vec<GP>> {
let mut basis = vec![d.make_monic()];
for c in cands {
let mut next = Vec::with_capacity(basis.len() + 1);
for e in basis {
let g = ctx.gcd(&e, c)?;
match (g.degree(), e.degree()) {
(Some(dg), Some(de)) if dg > 0 && dg < de => {
let rest = div_exact(&e, &g)?;
next.push(g.make_monic());
next.push(rest.make_monic());
}
_ => next.push(e),
}
}
basis = next;
if budget_exceeded() {
return None;
}
}
let mut out = Vec::with_capacity(basis.len());
for e in basis {
match to_q_poly(&e) {
Some(p) if e.degree().is_some_and(|k| k >= 2) => {
let (_, factors) = p.factor_over_z();
let total: usize = factors
.iter()
.map(|(f, m)| f.degree().unwrap_or(0) * *m as usize)
.sum();
if factors.is_empty() || total != p.degree().unwrap_or(0) {
out.push(e);
} else {
for (f, _) in factors {
out.push(from_q_poly(&f).make_monic());
}
}
}
None if e.degree().is_some_and(|k| k >= 3) => match split_parametric(ctx, &e) {
Some(pieces) => out.extend(pieces),
None => out.push(e),
},
_ => out.push(e),
}
}
Some(out)
}
fn factored_expr(arena: &mut Arena, ctx: &Ctx, p: &GP, basis: &[GP]) -> (ExprId, PFrac) {
let mut rest = p.clone();
let mut factors: Vec<ExprId> = Vec::new();
let mut lambda = PFrac::one();
for e in basis {
let mut m = 0u32;
while rest.degree() >= e.degree()
&& let Some(q) = div_exact(&rest, e)
{
rest = q;
m += 1;
}
if m > 0 {
let l = e
.coeffs()
.iter()
.fold(MP::from_int(ctx.nv(), 1), |acc, c| match c {
PFrac::Frac(f) if f.1.as_constant().is_none() => {
if gcd_affordable(&acc, &f.1) {
MP::lcm(&acc, &f.1)
} else {
acc.mul(&f.1)
}
}
_ => acc,
});
let l = PFrac::from_mp(l);
let raw = e.scale(&l);
let e_id = ctx.poly_expr(arena, &raw);
lambda = lambda.mul(&l.pow_usize(m as usize));
factors.push(if m == 1 {
e_id
} else {
let m_id = arena.int(i64::from(m));
arena.pow(e_id, m_id)
});
}
}
if rest.degree() != Some(0) || !rest.coeff(0).is_one() || factors.is_empty() {
return (ctx.poly_expr(arena, p), PFrac::one());
}
let product = if factors.len() == 1 {
factors[0]
} else {
arena.mul(&factors)
};
(product, lambda)
}
fn ln_abs(arena: &mut Arena, u: ExprId) -> ExprId {
let a = arena.abs(u);
arena.ln(a)
}
fn log_part(
arena: &mut Arena,
ctx: &Ctx,
h: &GP,
d: &GP,
basis: &[GP],
label: ExprId,
) -> Option<Vec<ExprId>> {
let mut terms = Vec::new();
for e in basis {
if budget_exceeded() {
return None;
}
let cof = div_exact(d, e)?;
let eg = GP::extended_gcd(&cof, e);
if eg.gcd.degree()? != 0 {
return None;
}
let inv = eg.gcd.leading_coeff()?.inv();
let (_, a_e) = h.mul(&eg.x.scale(&inv)).try_div_rem(e)?;
if a_e.is_zero() {
continue;
}
match e.degree()? {
1 => {
let e_id = ctx.poly_expr(arena, e);
let c = ctx.coeff_expr(arena, &a_e.coeff(0));
let ln = ln_abs(arena, e_id);
terms.push(arena.mul(&[c, ln]));
}
2 => terms.extend(quadratic_log_part(arena, ctx, &a_e, e)?),
_ => terms.push(higher_log_part(arena, ctx, &a_e, e, label)?),
}
}
Some(terms)
}
fn sqrt_expr(arena: &mut Arena, ctx: &Ctx, sigma: &PFrac, rho: &MP) -> ExprId {
let s = ctx.coeff_expr(arena, sigma);
match rho.as_constant() {
Some(c) if c.is_one() => s,
Some(c) if (-&c).is_one() => {
let i = arena.i_unit();
arena.mul(&[s, i])
}
_ => {
let r = ctx.mp_expr(arena, rho);
let root = arena.sqrt(r);
arena.mul(&[s, root])
}
}
}
fn quadratic_log_part(arena: &mut Arena, ctx: &Ctx, a_e: &GP, e: &GP) -> Option<Vec<ExprId>> {
let nv = ctx.nv();
let (mu, nu) = (a_e.coeff(1), a_e.coeff(0));
let (big_b, big_c) = (e.coeff(1), e.coeff(0));
let half = Q::new(BigInt::one(), BigInt::from(2));
let half_mu = mu.scale_q(&half);
let kappa = nu.sub(&half_mu.mul(&big_b));
let e_id = ctx.poly_expr(arena, e);
if kappa.is_zero() {
let c = ctx.coeff_expr(arena, &half_mu);
let ln = ln_abs(arena, e_id);
return Some(vec![arena.mul(&[c, ln])]);
}
let delta = big_b
.mul(&big_b)
.sub(&big_c.scale_q(&Q::from_integer(BigInt::from(4))));
if delta.is_zero() {
return None;
}
let (sigma, rho) = split_square(&delta, nv);
if rho.as_constant().is_some_and(|c| c.is_one()) {
let mut terms = Vec::with_capacity(2);
for sign in [Q::one(), -Q::one()] {
let r = big_b.neg().add(&sigma.scale_q(&sign)).scale_q(&half);
let res = mu.mul(&r).add(&nu).div(&sigma.scale_q(&sign));
if res.is_zero() {
continue;
}
let lin = GP::from_coeffs(vec![r.neg(), PFrac::one()]);
let lin_id = ctx.poly_expr(arena, &lin);
let c = ctx.coeff_expr(arena, &res);
let ln = ln_abs(arena, lin_id);
terms.push(arena.mul(&[c, ln]));
}
return Some(terms);
}
let b_id = ctx.coeff_expr(arena, &big_b);
let two = arena.int(2);
let mut terms = Vec::with_capacity(2);
if !half_mu.is_zero() {
let c = ctx.coeff_expr(arena, &half_mu);
let ln = ln_abs(arena, e_id);
terms.push(arena.mul(&[c, ln]));
}
let two_kappa = ctx.coeff_expr(arena, &kappa.scale_q(&Q::from_integer(BigInt::from(2))));
if ctx.declared_real(arena, e) && nonnegative(arena, ctx, &delta.neg()) {
let (sigma_k, rho_k) = split_square(&delta.neg(), nv);
let k = sqrt_expr(arena, ctx, &sigma_k, &rho_k);
let two_x = arena.mul(&[two, ctx.var]);
let lin = arena.add(&[two_x, b_id]);
let arg = arena.div(lin, k);
let atan = arena.atan(arg);
let coeff = arena.div(two_kappa, k);
terms.push(arena.mul(&[coeff, atan]));
return Some(terms);
}
let s = sqrt_expr(arena, ctx, &sigma, &rho);
let neg_b = arena.neg(b_id);
let mut lns = Vec::with_capacity(2);
for sign in [1i64, -1] {
let sgn = arena.int(sign);
let ss = arena.mul(&[sgn, s]);
let num = arena.add(&[neg_b, ss]);
let r = arena.div(num, two);
let neg_r = arena.neg(r);
let lin = arena.add(&[ctx.var, neg_r]);
lns.push(ln_abs(arena, lin));
}
let neg_minus = arena.neg(lns[1]);
let diff = arena.add(&[lns[0], neg_minus]);
let kappa_id = ctx.coeff_expr(arena, &kappa);
let coeff = arena.div(kappa_id, s);
terms.push(arena.mul(&[coeff, diff]));
Some(terms)
}
fn nonnegative(arena: &mut Arena, ctx: &Ctx, c: &PFrac) -> bool {
use crate::base::assumptions::{AssumptionCache, Props};
if let Some(q) = c.as_const() {
return !q.is_negative();
}
let (_, rho) = split_square(c, ctx.nv());
if let Some(q) = rho.as_constant() {
return !q.is_negative();
}
let e = ctx.mp_expr(arena, &rho);
let mut cache = AssumptionCache::new();
cache.query(arena, e, Props::NONNEGATIVE) == Some(true)
}
fn higher_log_part(
arena: &mut Arena,
ctx: &Ctx,
a_e: &GP,
e: &GP,
label: ExprId,
) -> Option<ExprId> {
let ep = deriv(e);
let c = a_e.leading_coeff()?.div(ep.leading_coeff()?);
if a_e.degree() == ep.degree() && ep.scale(&c) == *a_e {
let e_id = ctx.poly_expr(arena, e);
let c_id = ctx.coeff_expr(arena, &c);
let ln = ln_abs(arena, e_id);
return Some(arena.mul(&[c_id, ln]));
}
let Some(e_q) = to_q_poly(e) else {
return root_sum_log_part(arena, ctx, a_e, e);
};
let mut terms = Vec::new();
for (j, cj) in a_e.coeffs().iter().enumerate() {
if cj.is_zero() {
continue;
}
let mut xj = vec![Q::zero(); j + 1];
xj[j] = Q::one();
let f = super::integrate_rational_function(
arena,
&Poly::from_coeffs(xj),
&e_q,
ctx.var,
label,
)?;
if crate::base::walk::has_unevaluated(arena, f) {
return None;
}
let c_id = ctx.coeff_expr(arena, cj);
terms.push(arena.mul(&[c_id, f]));
}
Some(match terms.len() {
0 => arena.zero(),
1 => terms[0],
_ => arena.add(&terms),
})
}
fn root_sum_log_part(arena: &mut Arena, ctx: &Ctx, a_e: &GP, e: &GP) -> Option<ExprId> {
let ep = deriv(e);
let eg = GP::extended_gcd(&ep, e);
if eg.gcd.degree()? != 0 {
return None;
}
let inv = eg.gcd.leading_coeff()?.inv();
let (_, c) = a_e.mul(&eg.x.scale(&inv)).try_div_rem(e)?;
if budget_exceeded() || c.is_zero() {
return None;
}
let r = arena.symbol("__rs_t");
let ctx_r = Ctx {
var: r,
params: ctx.params.clone(),
};
let (e_r, _) = factored_expr(arena, &ctx_r, e, std::slice::from_ref(e));
let c_r = ctx_r.poly_expr(arena, &c);
let neg_r = arena.neg(r);
let lin = arena.add(&[ctx.var, neg_r]);
let ln = arena.ln(lin);
let body = arena.mul(&[c_r, ln]);
Some(arena.intern(ExprNode::RootSum(e_r, body, r)))
}
#[cfg(test)]
mod tests {
use super::*;
fn mp(nv: usize, terms: &[(&[u32], i64)]) -> MP {
MP::from_terms(
nv,
terms
.iter()
.map(|(e, c)| (e.to_vec(), Q::from_integer(BigInt::from(*c))))
.collect(),
)
.expect("valid terms")
}
#[test]
fn pfrac_is_canonical() {
let num = mp(2, &[(&[2, 0], 1), (&[0, 2], -1)]);
let den = mp(2, &[(&[1, 0], 1), (&[0, 1], -1)]);
let q = PFrac::canon(num, den);
assert_eq!(q, PFrac::from_mp(mp(2, &[(&[1, 0], 1), (&[0, 1], 1)])));
let half_inv = PFrac::canon(mp(2, &[(&[0, 0], 1)]), mp(2, &[(&[1, 0], 2)]));
let inv = PFrac::canon(mp(2, &[(&[0, 0], 1)]), mp(2, &[(&[1, 0], 1)]));
assert_eq!(half_inv.add(&half_inv), inv);
assert!(inv.sub(&inv).is_zero());
assert_eq!(inv.mul(&inv.inv()), PFrac::one());
}
#[test]
fn multipoly_square_roots() {
let sq = mp(2, &[(&[2, 0], 1), (&[1, 1], -4), (&[0, 2], 4)]);
let root = mp_sqrt(&sq).expect("a square");
assert_eq!(root.mul(&root), sq);
assert!(mp_sqrt(&mp(2, &[(&[2, 0], 1), (&[0, 1], -4)])).is_none());
let delta = PFrac::canon(mp(1, &[(&[0], -4)]), mp(1, &[(&[2], 1)]));
let (sigma, rho) = split_square(&delta, 1);
assert_eq!(rho.as_constant(), Some(-Q::one()));
assert_eq!(sigma.mul(&sigma).neg(), delta);
}
}