symplex 0.30.0

//! Compile symbolic expressions to callable numerical functions.
//!
//! [`Ex::compile`](crate::api::expr::Ex::compile) lowers one or more expressions into a compact stack-VM
//! program after a shared common-subexpression-elimination
//! pass.  The program is wrapped in [`CompiledFn`] (scalar) or
//! [`CompiledFnVec`] (vector-valued) handles that are `Clone + Send + Sync`
//! and evaluate without touching the arena or taking any locks.
//!
//! Every numerically evaluable [`ExprNode`] is supported, including the
//! special functions (Γ, ln Γ, ψ, erf/erfc and their inverses erfinv/erfcinv,
//! Lambert W, Beta, factorials,
//! binomials), Bessel functions and orthogonal polynomials with integer
//! order, integer sequences, piecewise expressions and boolean/relational
//! nodes (represented as `0.0`/`1.0`).  The special-function algorithms live
//! in the crate-internal `numeric_rt` module and are shared with
//! the Rust code generator, so `compile()` and `to_rust_fn()` agree bit for
//! bit on those routines.
//!
//! Symbol-free subexpressions are decided at compile time like the code
//! emitters decide them: a real constant is folded to its certified `f64`
//! value, and a constant that is not real (`atanh(9)`, `asin(2)`,
//! `(-8)^(1/3)`) is an error (`NotImplemented`, "the constant … is not
//! real"), since a real-valued function cannot represent it.
//!
//! One rule, no exceptions: a compiled function returns the value `evalf`
//! gives (principal branches over ℂ, SymPy's meaning), or NaN where that
//! value is not real.  So `x^(p/q)` of a negative `x` is NaN for every
//! non-integer `p/q` — odd denominators included (the principal
//! `(-8)^(1/3)` is `1 + √3·i`; before 0.30 the compiled value was the real
//! root −2, and a Cardano formula from `solve` compiled to a different root);
//! the real root is `real_root` (`sign(x)·|x|^(1/q)`), which compiles to −2.
//! `loggamma(x)` is NaN for a negative non-integer `x` (its value there is
//! `ln|Γ(x)| − iπ⌈−x⌉`); the real `ln|Γ(x)|` is `ln(abs(gamma(x)))`, lowered
//! to the overflow-safe `lgamma`.
//!
//! The compiler never recurses over the expression tree — it uses an
//! explicit work stack — so arbitrarily deep expressions are safe.

use crate::base::arena::Arena;
use crate::base::errors::SymplexError;
use crate::base::libfn::LibFn;
use crate::base::node::{ExprId, ExprNode, SymbolId};
use crate::output::codegen::numeric_rt as rt;
use num_traits::{One, Signed, ToPrimitive};
use rustc_hash::FxHashMap;
use smallvec::SmallVec;
use std::fmt;
use std::ops::Deref;
use std::sync::Arc;

/// Boxed, thread-safe closure type behind [`CompiledFn`]'s `Deref`.
type SharedFn = Arc<dyn Fn(&[f64]) -> f64 + Send + Sync>;

// ═══════════════════════════════════════════════════════════════════════════
// Public handle types
// ═══════════════════════════════════════════════════════════════════════════

/// A compiled scalar function `f(x₁, …, xₙ) → f64`.
///
/// Produced by [`Ex::compile`](crate::api::expr::Expr::compile).  The handle
/// is cheap to clone (it shares the underlying program), is `Send + Sync`,
/// and can be called directly like a closure thanks to a `Deref` to
/// `dyn Fn(&[f64]) -> f64`:
///
/// ```
/// use symplex::prelude::*;
///
/// let ctx = Context::new();
/// let x = ctx.symbol("x");
/// let f = (x.powi(2) + 1).compile(&["x"]).unwrap();
/// assert_eq!(f.arity(), 1);
/// assert_eq!(f(&[3.0]), 10.0);        // closure-style call
/// assert_eq!(f.call(&[3.0]), 10.0);   // explicit call
/// assert!(f.call(&[1.0, 2.0]).is_nan()); // wrong arity → NaN
/// assert!(f.try_call(&[1.0, 2.0]).is_err());
/// ```
///
/// Calling with the wrong number of arguments returns `NaN` from
/// [`call`](Self::call) (and the closure form); use
/// [`try_call`](Self::try_call) to get an error instead.
#[derive(Clone)]
pub struct CompiledFn {
    program: Arc<Program>,
    func: SharedFn,
}

impl CompiledFn {
    fn from_program(program: Program) -> Self {
        let program = Arc::new(program);
        let p = Arc::clone(&program);
        let func: SharedFn = Arc::new(move |args: &[f64]| p.run_scalar(args));
        Self { program, func }
    }

    /// Number of arguments the function expects.
    #[must_use]
    pub fn arity(&self) -> usize {
        self.program.arity
    }

    /// Evaluate the function.  Returns `NaN` if `args.len() != self.arity()`.
    #[must_use]
    pub fn call(&self, args: &[f64]) -> f64 {
        self.program.run_scalar(args)
    }

    /// Evaluate the function, returning an error on arity mismatch.
    pub fn try_call(&self, args: &[f64]) -> Result<f64, SymplexError> {
        self.program.check_arity("CompiledFn::try_call", args)?;
        Ok(self.program.run_scalar(args))
    }

    /// Number of VM instructions in the compiled program (a rough cost measure).
    #[must_use]
    pub fn instruction_count(&self) -> usize {
        self.program.code.len()
    }
}

impl Deref for CompiledFn {
    type Target = dyn Fn(&[f64]) -> f64 + Send + Sync;
    fn deref(&self) -> &Self::Target {
        &*self.func
    }
}

impl fmt::Debug for CompiledFn {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("CompiledFn")
            .field("arity", &self.program.arity)
            .field("instructions", &self.program.code.len())
            .field("locals", &self.program.n_locals)
            .finish()
    }
}

/// A compiled vector-valued function `F(x₁, …, xₙ) → (f₁, …, fₖ)`.
///
/// Produced by [`Ex::compile_many`](crate::api::expr::Expr::compile_many);
/// all outputs share one common-subexpression-elimination pass, which makes
/// this the right tool for gradients, Jacobians and other expression
/// families with heavy overlap.
///
/// ```
/// use symplex::prelude::*;
///
/// let ctx = Context::new();
/// let x = ctx.symbol("x");
/// let y = ctx.symbol("y");
/// let f = &x.sin() * &y;
/// let grad = Ex::compile_many(&[&f.diff(&x), &f.diff(&y)], &["x", "y"]).unwrap();
/// assert_eq!(grad.len(), 2);
/// let g = grad.call_vec(&[0.0, 2.0]);
/// assert_eq!(g, vec![2.0, 0.0]);
/// ```
#[derive(Clone)]
pub struct CompiledFnVec {
    program: Arc<Program>,
}

impl CompiledFnVec {
    /// Number of arguments the function expects.
    #[must_use]
    pub fn arity(&self) -> usize {
        self.program.arity
    }

    /// Number of output values.
    #[must_use]
    pub fn len(&self) -> usize {
        self.program.n_outputs
    }

    /// True when the function produces no outputs.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.program.n_outputs == 0
    }

    /// Evaluate into `out`.
    ///
    /// On arity mismatch, or if `out.len() != self.len()`, every element of
    /// `out` is set to `NaN` and no other work is done.
    pub fn call(&self, args: &[f64], out: &mut [f64]) {
        if args.len() != self.program.arity || out.len() != self.program.n_outputs {
            out.fill(f64::NAN);
            return;
        }
        self.program.run(args, out);
    }

    /// Evaluate into `out`, returning an error on arity or length mismatch.
    pub fn try_call(&self, args: &[f64], out: &mut [f64]) -> Result<(), SymplexError> {
        self.program.check_arity("CompiledFnVec::try_call", args)?;
        if out.len() != self.program.n_outputs {
            return Err(SymplexError::InvalidArgument {
                operation: "CompiledFnVec::try_call",
                reason: format!(
                    "output slice has length {}, expected {}",
                    out.len(),
                    self.program.n_outputs
                ),
            });
        }
        self.program.run(args, out);
        Ok(())
    }

    /// Evaluate and return the outputs as a freshly allocated vector
    /// (all `NaN` on arity mismatch).
    #[must_use]
    pub fn call_vec(&self, args: &[f64]) -> Vec<f64> {
        let mut out = vec![f64::NAN; self.program.n_outputs];
        self.call(args, &mut out);
        out
    }

    /// Number of VM instructions in the compiled program.
    #[must_use]
    pub fn instruction_count(&self) -> usize {
        self.program.code.len()
    }
}

impl fmt::Debug for CompiledFnVec {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("CompiledFnVec")
            .field("arity", &self.program.arity)
            .field("outputs", &self.program.n_outputs)
            .field("instructions", &self.program.code.len())
            .field("locals", &self.program.n_locals)
            .finish()
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Entry points (crate-internal)
// ═══════════════════════════════════════════════════════════════════════════

/// Compile a single expression into a [`CompiledFn`].
///
/// Runs `eval()` (cheap constant folding) and CSE before lowering.
pub(crate) fn compile(
    arena: &mut Arena,
    expr: ExprId,
    var_names: &[&str],
) -> Result<CompiledFn, SymplexError> {
    let program = compile_program(arena, &[expr], var_names)?;
    Ok(CompiledFn::from_program(program))
}

/// Compile a single expression **without** the `eval()` / CSE pre-passes.
///
/// Works on an immutable arena, so it can be used from contexts that only
/// hold `&Arena` (numeric evaluation of a `DefiniteIntegral` body inside
/// `evalf`).  The caller is expected to have constant-folded `expr`
/// already; nothing is interned.
pub(crate) fn compile_raw(
    arena: &Arena,
    expr: ExprId,
    var_names: &[&str],
) -> Result<CompiledFn, SymplexError> {
    check_params(var_names)?;
    // No constant folding here: `evalf` itself compiles integrand bodies
    // through this entry point.
    let mut em = Emitter::new(arena, var_names, FxHashMap::default());
    em.lower(expr)?;
    em.emit(Instruction::StoreOut(0));
    let code = em.finish()?;
    Ok(CompiledFn::from_program(Program {
        code,
        arity: var_names.len(),
        n_locals: 0,
        n_outputs: 1,
    }))
}

/// Compile several expressions into one [`CompiledFnVec`] with a shared CSE pass.
pub(crate) fn compile_many(
    arena: &mut Arena,
    exprs: &[ExprId],
    var_names: &[&str],
) -> Result<CompiledFnVec, SymplexError> {
    let program = compile_program(arena, exprs, var_names)?;
    Ok(CompiledFnVec {
        program: Arc::new(program),
    })
}

/// A [`CompiledFnVec`] with no outputs (no arena required).
pub(crate) fn compile_many_empty(var_names: &[&str]) -> Result<CompiledFnVec, SymplexError> {
    check_params(var_names)?;
    Ok(CompiledFnVec {
        program: Arc::new(Program {
            code: Vec::new(),
            arity: var_names.len(),
            n_locals: 0,
            n_outputs: 0,
        }),
    })
}

fn check_params(var_names: &[&str]) -> Result<(), SymplexError> {
    for (i, a) in var_names.iter().enumerate() {
        if var_names[..i].contains(a) {
            return Err(SymplexError::InvalidArgument {
                operation: "compile",
                reason: format!("duplicate parameter name '{a}'"),
            });
        }
    }
    Ok(())
}

/// Shared lowering pipeline: eval → CSE → bytecode.
fn compile_program(
    arena: &mut Arena,
    exprs: &[ExprId],
    var_names: &[&str],
) -> Result<Program, SymplexError> {
    check_params(var_names)?;

    // Pre-pass: fold exact constants (sin(0)→0, cos(π)→−1, 4^(1/2)→2, …).
    //
    // Two families are lowered as-is instead:
    // * `Piecewise` — `eval()` currently mis-simplifies piecewise nodes whose
    //   earlier conditions are symbolic (it selects the first literal-`True`
    //   branch);
    // * orthogonal polynomials — `eval()` expands them into monomials, which
    //   is numerically much worse than the three-term recurrence used by the
    //   runtime (1e-13 vs 1e-16 absolute error for P₁₁ on [−1, 1]).
    let evaled: Vec<ExprId> = exprs
        .iter()
        .map(|&e| {
            if needs_raw_lowering(arena, e) {
                e
            } else {
                crate::transforms::eval::eval(arena, e)
            }
        })
        .collect();

    // Constants are decided before CSE, which would otherwise split a real
    // constant around a shared non-real part.
    let (folded, constants) = fold_constants_before_cse(arena, &evaled, var_names)?;

    // Shared CSE across all outputs.
    let cse = crate::output::cse::cse_multi(arena, &folded);

    let mut locals: FxHashMap<SymbolId, usize> = FxHashMap::default();
    for (i, (name_id, _)) in cse.bindings.iter().enumerate() {
        if let ExprNode::Symbol(sid) = arena.node(*name_id) {
            locals.insert(*sid, i);
        }
    }

    let mut em = Emitter::new(arena, var_names, locals);
    em.constants = constants;
    em.fold_constants = true;
    for (i, (_, value)) in cse.bindings.iter().enumerate() {
        em.lower(*value)?;
        em.emit(Instruction::StoreLocal(i));
    }
    for (i, &e) in cse.exprs.iter().enumerate() {
        em.lower(e)?;
        em.emit(Instruction::StoreOut(i));
    }
    let code = em.finish()?;

    Ok(Program {
        code,
        arity: var_names.len(),
        n_locals: cse.bindings.len(),
        n_outputs: exprs.len(),
    })
}

/// Decide every maximal symbol-free compound subexpression of `exprs`
/// before CSE, with the rule of the code emitters
/// (`codegen::fold_constants_for_emission`): a real value from the
/// certified evaluator is folded (replaced by a placeholder symbol bound to
/// that `f64`, so CSE treats it as an atom); a formula the real VM already
/// evaluates to a finite value is kept; a constant that is not real
/// (`(-8)^(1/3)`, the principal cube root `1 + √3·i`) is refused with the
/// emitters' error; anything else keeps its formula and its parts are
/// examined.  Iterative (explicit stack).
///
/// Before, a non-real constant (`atanh(9)·x`, `asin(2) + x`) compiled to a
/// function returning NaN everywhere while `to_rust_fn`, `to_c_fn` and the
/// Python emitters refused it, and a real constant sharing a non-real part
/// with another (`|atanh(9)| + |atanh(9) + 1|·x`) was split by CSE around
/// the shared `atanh(9)` and compiled to NaN too.
fn fold_constants_before_cse(
    arena: &mut Arena,
    exprs: &[ExprId],
    var_names: &[&str],
) -> Result<(Vec<ExprId>, FxHashMap<SymbolId, f64>), SymplexError> {
    // Symbol-freeness of every node, and the names in use, bottom-up.
    let mut has_symbol: FxHashMap<ExprId, bool> = FxHashMap::default();
    let mut names: rustc_hash::FxHashSet<String> =
        var_names.iter().map(|s| (*s).to_string()).collect();
    for &root in exprs {
        for id in crate::base::walk::post_order_ids(arena, root) {
            if has_symbol.contains_key(&id) {
                continue;
            }
            let node = arena.node(id);
            if let ExprNode::Symbol(sid) = node {
                names.insert(arena.symbol_name(*sid).to_string());
            }
            let mut any = matches!(node, ExprNode::Symbol(_));
            node.for_each_child(|c| any |= has_symbol.get(&c).copied().unwrap_or(true));
            has_symbol.insert(id, any);
        }
    }
    let mut replacements: Vec<(ExprId, ExprId)> = Vec::new();
    let mut values: FxHashMap<SymbolId, f64> = FxHashMap::default();
    let mut seen = rustc_hash::FxHashSet::default();
    let mut stack: Vec<ExprId> = exprs.to_vec();
    let mut next_index = 0usize;
    while let Some(id) = stack.pop() {
        if !seen.insert(id) {
            continue;
        }
        let node = arena.node(id).clone();
        let atom = matches!(
            node,
            ExprNode::Num(_)
                | ExprNode::Symbol(_)
                | ExprNode::Pi
                | ExprNode::E
                | ExprNode::EulerGamma
                | ExprNode::Catalan
                | ExprNode::GoldenRatio
                | ExprNode::Infinity
                | ExprNode::NegInfinity
                | ExprNode::NaN
                | ExprNode::ComplexInfinity
                | ExprNode::ImaginaryUnit
        );
        if atom || is_bool_valued(arena, id) || has_symbol.get(&id).copied().unwrap_or(true) {
            node.for_each_child(|c| stack.push(c));
            continue;
        }
        match crate::transforms::evalf::evalf_f64(arena, id) {
            Ok(v) if v.is_finite() => {
                let name = loop {
                    let candidate = format!("__compile_const_{next_index}");
                    next_index += 1;
                    if !names.contains(&candidate) {
                        break candidate;
                    }
                };
                let sym = arena.symbol(&name);
                if let ExprNode::Symbol(sid) = arena.node(sym) {
                    values.insert(*sid, v);
                    replacements.push((id, sym));
                }
                continue;
            }
            _ => {}
        }
        let naive = compile_raw(arena, id, &[]).ok().map(|f| f.call(&[]));
        if naive.is_some_and(f64::is_finite) {
            continue;
        }
        match crate::transforms::evalf::evalf_complex64(arena, id) {
            Ok(z) if z.im != 0.0 && z.re.is_finite() && z.im.is_finite() => {
                return Err(SymplexError::NotImplemented(format!(
                    "compile: the constant `{}` is not real ({} {} {}i); \
                     a real-valued function cannot represent it",
                    arena.display(id),
                    z.re,
                    if z.im < 0.0 { '-' } else { '+' },
                    z.im.abs()
                )));
            }
            _ => node.for_each_child(|c| stack.push(c)),
        }
    }
    let folded = exprs
        .iter()
        .map(|&e| crate::transforms::subs::subs_map(arena, e, &replacements))
        .collect();
    Ok((folded, values))
}

// ═══════════════════════════════════════════════════════════════════════════
// Program / VM
// ═══════════════════════════════════════════════════════════════════════════

/// A lowered stack-machine program.
struct Program {
    code: Vec<Instruction>,
    arity: usize,
    n_locals: usize,
    n_outputs: usize,
}

impl Program {
    fn check_arity(&self, operation: &'static str, args: &[f64]) -> Result<(), SymplexError> {
        if args.len() != self.arity {
            return Err(SymplexError::InvalidArgument {
                operation,
                reason: format!("expected {} argument(s), got {}", self.arity, args.len()),
            });
        }
        Ok(())
    }

    fn run_scalar(&self, args: &[f64]) -> f64 {
        if args.len() != self.arity {
            return f64::NAN;
        }
        let mut out = [f64::NAN];
        self.run(args, &mut out);
        out[0]
    }

    /// Execute the program.  Callers must have validated arities.
    fn run(&self, args: &[f64], out: &mut [f64]) {
        let mut stack: SmallVec<[f64; 32]> = SmallVec::new();
        let mut locals: SmallVec<[f64; 16]> = SmallVec::new();
        locals.resize(self.n_locals, 0.0);
        let code = &self.code;
        let mut pc = 0usize;

        macro_rules! pop {
            () => {
                stack.pop().unwrap_or(f64::NAN)
            };
        }
        macro_rules! un {
            ($f:expr) => {{
                let a = pop!();
                stack.push($f(a));
            }};
        }
        macro_rules! bin {
            ($f:expr) => {{
                let b = pop!();
                let a = pop!();
                stack.push($f(a, b));
            }};
        }

        while pc < code.len() {
            match code[pc] {
                Instruction::PushConst(v) => stack.push(v),
                Instruction::PushVar(i) => stack.push(args.get(i).copied().unwrap_or(f64::NAN)),
                Instruction::LoadLocal(i) => stack.push(locals.get(i).copied().unwrap_or(f64::NAN)),
                Instruction::StoreLocal(i) => {
                    let v = pop!();
                    if let Some(slot) = locals.get_mut(i) {
                        *slot = v;
                    }
                }
                Instruction::StoreOut(i) => {
                    let v = pop!();
                    if let Some(slot) = out.get_mut(i) {
                        *slot = v;
                    }
                }
                Instruction::Add => bin!(|a: f64, b: f64| a + b),
                Instruction::Mul => bin!(|a: f64, b: f64| a * b),
                Instruction::Div => bin!(|a: f64, b: f64| a / b),
                Instruction::Neg => un!(|a: f64| -a),
                Instruction::Pow => bin!(|a: f64, b: f64| a.powf(b)),
                Instruction::Powi(n) => un!(|a: f64| a.powi(n)),
                Instruction::Sqrt => un!(|a: f64| a.sqrt()),
                // The principal cube root: not real for a negative argument.
                Instruction::Cbrt => un!(|a: f64| if a < 0.0 { f64::NAN } else { a.cbrt() }),
                Instruction::ExpM1 => un!(|a: f64| a.exp_m1()),
                Instruction::Ln1p => un!(|a: f64| a.ln_1p()),
                Instruction::Sin => un!(|a: f64| a.sin()),
                Instruction::Cos => un!(|a: f64| a.cos()),
                Instruction::Tan => un!(|a: f64| a.tan()),
                Instruction::Exp => un!(|a: f64| a.exp()),
                Instruction::Ln => un!(|a: f64| a.ln()),
                Instruction::Abs => un!(|a: f64| a.abs()),
                Instruction::Asin => un!(|a: f64| a.asin()),
                Instruction::Acos => un!(|a: f64| a.acos()),
                Instruction::Atan => un!(|a: f64| a.atan()),
                Instruction::Sinh => un!(|a: f64| a.sinh()),
                Instruction::Cosh => un!(|a: f64| a.cosh()),
                Instruction::Tanh => un!(|a: f64| a.tanh()),
                Instruction::Asinh => un!(|a: f64| a.asinh()),
                Instruction::Acosh => un!(|a: f64| a.acosh()),
                Instruction::Atanh => un!(|a: f64| a.atanh()),
                // An undefined (NaN) argument stays undefined: `sign(sqrt(-1))`
                // was 0 and `heaviside(asin(2))` 0.5 before 0.30.
                Instruction::Sign => un!(|a: f64| if a > 0.0 {
                    1.0
                } else if a < 0.0 {
                    -1.0
                } else if a == 0.0 {
                    0.0
                } else {
                    f64::NAN
                }),
                Instruction::Heaviside => un!(|a: f64| if a > 0.0 {
                    1.0
                } else if a < 0.0 {
                    0.0
                } else if a == 0.0 {
                    0.5
                } else {
                    f64::NAN
                }),
                // Distributional: pointwise evaluation is 0 everywhere (see docs).
                Instruction::DiracDelta => un!(|a: f64| if a.is_nan() { f64::NAN } else { 0.0 }),
                Instruction::Atan2 => bin!(|y: f64, x: f64| y.atan2(x)),
                Instruction::Floor => un!(|a: f64| a.floor()),
                Instruction::Ceiling => un!(|a: f64| a.ceil()),
                // NaN-propagating (`f64::min` returns the other operand).
                Instruction::Min2 => bin!(|a: f64, b: f64| if a.is_nan() || b.is_nan() {
                    f64::NAN
                } else {
                    a.min(b)
                }),
                Instruction::Max2 => bin!(|a: f64, b: f64| if a.is_nan() || b.is_nan() {
                    f64::NAN
                } else {
                    a.max(b)
                }),
                Instruction::Gamma => un!(rt::gamma),
                Instruction::LogGamma => un!(rt::loggamma),
                Instruction::LnAbsGamma => un!(rt::lgamma),
                Instruction::Re => un!(|a: f64| a),
                Instruction::Im => un!(|a: f64| if a.is_nan() { f64::NAN } else { 0.0 }),
                Instruction::Digamma => un!(rt::digamma),
                Instruction::Erf => un!(rt::erf),
                Instruction::Erfc => un!(rt::erfc),
                Instruction::Erfinv => un!(rt::erfinv),
                Instruction::Erfcinv => un!(rt::erfcinv),
                Instruction::LambertW => un!(rt::lambert_w0),
                Instruction::LambertWm1 => un!(rt::lambert_wm1),
                Instruction::Factorial => un!(rt::factorial),
                Instruction::Binomial => bin!(rt::binomial),
                Instruction::Beta => bin!(rt::beta),
                Instruction::BesselJ(n) => un!(|a: f64| rt::bessel_j(n, a)),
                Instruction::BesselY(n) => un!(|a: f64| rt::bessel_y(n, a)),
                Instruction::BesselI(n) => un!(|a: f64| rt::bessel_i(n, a)),
                Instruction::BesselK(n) => un!(|a: f64| rt::bessel_k(n, a)),
                Instruction::LegendreP(n) => un!(|a: f64| rt::legendre_p(n, a)),
                Instruction::ChebyshevT(n) => un!(|a: f64| rt::chebyshev_t(n, a)),
                Instruction::ChebyshevU(n) => un!(|a: f64| rt::chebyshev_u(n, a)),
                Instruction::HermiteH(n) => un!(|a: f64| rt::hermite_h(n, a)),
                Instruction::LaguerreL(n) => un!(|a: f64| rt::laguerre_l(n, a)),
                Instruction::Fibonacci => un!(rt::fibonacci),
                Instruction::Lucas => un!(rt::lucas),
                Instruction::Harmonic => un!(rt::harmonic),
                Instruction::Factorial2 => un!(rt::factorial2),
                Instruction::RisingFactorial => bin!(rt::rising_factorial),
                Instruction::FallingFactorial => bin!(rt::falling_factorial),
                Instruction::Gt => bin!(|a: f64, b: f64| bool_f64(a > b)),
                Instruction::Ge => bin!(|a: f64, b: f64| bool_f64(a >= b)),
                Instruction::Eq => bin!(|a: f64, b: f64| bool_f64(a == b)),
                Instruction::Ne => bin!(|a: f64, b: f64| bool_f64(a != b)),
                Instruction::And => bin!(|a: f64, b: f64| bool_f64(a != 0.0 && b != 0.0)),
                Instruction::Or => bin!(|a: f64, b: f64| bool_f64(a != 0.0 || b != 0.0)),
                Instruction::Not => un!(|a: f64| bool_f64(a == 0.0)),
                Instruction::JumpIfZero(target) => {
                    let c = pop!();
                    if c == 0.0 || c.is_nan() {
                        pc = target;
                        continue;
                    }
                }
                Instruction::Jump(target) => {
                    pc = target;
                    continue;
                }
            }
            pc += 1;
        }
    }
}

#[inline(always)]
fn bool_f64(b: bool) -> f64 {
    if b { 1.0 } else { 0.0 }
}

/// A single stack-machine instruction.
#[derive(Clone, Copy, Debug, PartialEq)]
enum Instruction {
    PushConst(f64),
    /// Push argument `i`.
    PushVar(usize),
    /// Push CSE temporary `i`.
    LoadLocal(usize),
    /// Pop into CSE temporary `i`.
    StoreLocal(usize),
    /// Pop into output slot `i`.
    StoreOut(usize),
    Add,
    Mul,
    /// `second / top` (produced by a peephole on `Powi(-1); Mul`).
    Div,
    Neg,
    /// `second ^ top`.
    Pow,
    Powi(i32),
    Sqrt,
    /// The principal cube root: NaN for a negative argument.
    Cbrt,
    ExpM1,
    Ln1p,
    Sin,
    Cos,
    Tan,
    Exp,
    Ln,
    Abs,
    Asin,
    Acos,
    Atan,
    Sinh,
    Cosh,
    Tanh,
    Asinh,
    Acosh,
    Atanh,
    Sign,
    Heaviside,
    DiracDelta,
    /// `atan2(second, top)`.
    Atan2,
    Floor,
    Ceiling,
    Min2,
    Max2,
    Gamma,
    /// SymPy's `loggamma`: NaN left of 0 except at the poles (`+∞`).
    LogGamma,
    /// `ln|Γ(x)|` (`lgamma`), for the pattern `ln(abs(gamma(x)))`.
    LnAbsGamma,
    /// `re` of a real value.
    Re,
    /// `im` of a real value: 0 (NaN for NaN, which stands for a value that
    /// is not real or not defined).
    Im,
    Digamma,
    Erf,
    Erfc,
    Erfinv,
    Erfcinv,
    LambertW,
    /// The lower real branch `W₋₁` (`lambertw(x, -1)`).
    LambertWm1,
    Factorial,
    /// `binomial(second, top)`.
    Binomial,
    /// `beta(second, top)`.
    Beta,
    BesselJ(i32),
    BesselY(i32),
    BesselI(i32),
    BesselK(i32),
    LegendreP(i32),
    ChebyshevT(i32),
    ChebyshevU(i32),
    HermiteH(i32),
    LaguerreL(i32),
    Fibonacci,
    Lucas,
    Harmonic,
    Factorial2,
    RisingFactorial,
    FallingFactorial,
    Gt,
    Ge,
    Eq,
    Ne,
    And,
    Or,
    Not,
    /// Pop a boolean (0/1); jump to the absolute index if it is zero (or NaN).
    JumpIfZero(usize),
    Jump(usize),
}

// ═══════════════════════════════════════════════════════════════════════════
// Emitter (iterative lowering with an explicit work stack)
// ═══════════════════════════════════════════════════════════════════════════

/// Work item for the explicit-stack lowering loop.
#[derive(Clone, Copy)]
enum Task {
    /// Lower this expression (its value ends up on the VM stack).
    Node(ExprId),
    /// Emit one instruction.
    Emit(Instruction),
    /// Bind label `l` to the current code position.
    Label(usize),
    /// Emit a `JumpIfZero` to label `l` (patched at the end).
    JumpIfZero(usize),
    /// Emit a `Jump` to label `l` (patched at the end).
    Jump(usize),
    /// Lower a boolean-valued expression (pushes 0.0/1.0).
    Bool(ExprId),
}

struct Emitter<'a> {
    arena: &'a Arena,
    vars: FxHashMap<&'a str, usize>,
    locals: FxHashMap<SymbolId, usize>,
    code: Vec<Instruction>,
    labels: Vec<Option<usize>>,
    fixups: Vec<(usize, usize)>,
    /// Peephole optimisations must not reach back across this position
    /// (a label may point here).
    barrier: usize,
    work: Vec<Task>,
    /// Evaluate symbol-free subexpressions once, at compile time, with the
    /// certified evaluator (see [`Emitter::constant_value`]).
    fold_constants: bool,
    /// Does the subtree at an id contain a symbol (memo for folding)?
    has_symbol: FxHashMap<ExprId, bool>,
    /// Placeholder symbols of the constants folded before CSE
    /// ([`fold_constants_before_cse`]) and their values.
    constants: FxHashMap<SymbolId, f64>,
}

impl<'a> Emitter<'a> {
    fn new(arena: &'a Arena, var_names: &[&'a str], locals: FxHashMap<SymbolId, usize>) -> Self {
        let vars = var_names.iter().enumerate().map(|(i, &n)| (n, i)).collect();
        Self {
            arena,
            vars,
            locals,
            code: Vec::new(),
            labels: Vec::new(),
            fixups: Vec::new(),
            barrier: 0,
            work: Vec::new(),
            fold_constants: false,
            has_symbol: FxHashMap::default(),
            constants: FxHashMap::default(),
        }
    }

    /// The `f64` value of a symbol-free compound subexpression, from the
    /// certified evaluator, when it is real.
    ///
    /// Lowered operation by operation, a real constant can pass through a
    /// complex intermediate the real VM cannot hold: `abs(atanh(9))` (1.576)
    /// compiled to NaN, `re(sqrt(-2))`, `zeta(3)` and `polygamma(1, 2)` did
    /// not compile at all.  Folding also rounds such constants once.  The
    /// maximal constants were decided before CSE ([`fold_constants_before_cse`]:
    /// a non-real one is refused there); this folds the parts of a formula
    /// kept there and of CSE bindings.
    /// An unevaluable constant is lowered operation by operation.
    fn constant_value(&mut self, id: ExprId) -> Option<f64> {
        let arena = self.arena;
        if !self.fold_constants || is_bool_valued(arena, id) {
            return None;
        }
        let compound = !matches!(
            arena.node(id),
            ExprNode::Num(_)
                | ExprNode::Symbol(_)
                | ExprNode::Pi
                | ExprNode::E
                | ExprNode::EulerGamma
                | ExprNode::Catalan
                | ExprNode::GoldenRatio
                | ExprNode::Infinity
                | ExprNode::NegInfinity
                | ExprNode::NaN
                | ExprNode::ComplexInfinity
                | ExprNode::ImaginaryUnit
        );
        if !compound || self.subtree_has_symbol(id) {
            return None;
        }
        crate::transforms::evalf::evalf_f64(arena, id)
            .ok()
            .filter(|v| v.is_finite())
    }

    /// Does the subtree at `id` contain a symbol?  Bottom-up over the
    /// post-order of `id`, memoised across calls (no recursion).
    fn subtree_has_symbol(&mut self, id: ExprId) -> bool {
        if let Some(&known) = self.has_symbol.get(&id) {
            return known;
        }
        let arena = self.arena;
        for n in crate::base::walk::post_order_ids(arena, id) {
            if self.has_symbol.contains_key(&n) {
                continue;
            }
            let node = arena.node(n);
            let mut any = matches!(node, ExprNode::Symbol(_));
            node.for_each_child(|c| any |= self.has_symbol.get(&c).copied().unwrap_or(true));
            self.has_symbol.insert(n, any);
        }
        self.has_symbol.get(&id).copied().unwrap_or(true)
    }

    fn new_label(&mut self) -> usize {
        self.labels.push(None);
        self.labels.len() - 1
    }

    fn emit(&mut self, inst: Instruction) {
        // Peephole: `x; y; Powi(-1); Mul` → `x; y; Div` (one rounding instead of two).
        if inst == Instruction::Mul
            && self.code.len() > self.barrier
            && self.code.last() == Some(&Instruction::Powi(-1))
        {
            self.code.pop();
            self.code.push(Instruction::Div);
            return;
        }
        self.code.push(inst);
    }

    /// Push tasks so that they execute in the given (forward) order.
    fn push_seq(&mut self, tasks: &[Task]) {
        for t in tasks.iter().rev() {
            self.work.push(*t);
        }
    }

    /// Lower `root`, appending code that leaves its value on the VM stack.
    fn lower(&mut self, root: ExprId) -> Result<(), SymplexError> {
        self.work.push(Task::Node(root));
        while let Some(task) = self.work.pop() {
            match task {
                Task::Emit(inst) => self.emit(inst),
                Task::Label(l) => {
                    self.labels[l] = Some(self.code.len());
                    self.barrier = self.code.len();
                }
                Task::JumpIfZero(l) => {
                    self.fixups.push((self.code.len(), l));
                    self.code.push(Instruction::JumpIfZero(usize::MAX));
                }
                Task::Jump(l) => {
                    self.fixups.push((self.code.len(), l));
                    self.code.push(Instruction::Jump(usize::MAX));
                }
                Task::Node(id) => self.lower_node(id)?,
                Task::Bool(id) => self.lower_bool(id)?,
            }
        }
        Ok(())
    }

    fn finish(mut self) -> Result<Vec<Instruction>, SymplexError> {
        for (pos, label) in self.fixups.drain(..) {
            let target = self.labels[label].ok_or_else(|| SymplexError::ComputationFailed {
                operation: "compile",
                reason: "internal error: unbound jump label".to_string(),
            })?;
            match &mut self.code[pos] {
                Instruction::JumpIfZero(t) | Instruction::Jump(t) => *t = target,
                _ => {}
            }
        }
        Ok(self.code)
    }

    fn unsupported(&self, what: &str) -> SymplexError {
        SymplexError::NotImplemented(format!("cannot compile `{what}` to a numerical function"))
    }

    /// Resolve an expression to a compile-time integer (for orders of Bessel
    /// functions and degrees of orthogonal polynomials).
    fn const_i32(&self, id: ExprId, what: &str) -> Result<i32, SymplexError> {
        if let Some(r) = self.arena.as_num(id)
            && r.is_integer()
            && let Some(n) = r.to_integer().to_i32()
        {
            return Ok(n);
        }
        if let ExprNode::Neg(inner) = self.arena.node(id)
            && let Some(r) = self.arena.as_num(*inner)
            && r.is_integer()
            && let Some(n) = r.to_integer().to_i32()
        {
            return Ok(-n);
        }
        Err(SymplexError::NotImplemented(format!(
            "{what} requires a constant integer order/degree, got `{}`",
            self.arena.display(id)
        )))
    }

    fn unary(&mut self, child: ExprId, inst: Instruction) {
        self.push_seq(&[Task::Node(child), Task::Emit(inst)]);
    }

    fn binary(&mut self, a: ExprId, b: ExprId, inst: Instruction) {
        self.push_seq(&[Task::Node(a), Task::Node(b), Task::Emit(inst)]);
    }

    /// Left fold: `c₀ op c₁ op c₂ …` with `empty` for zero children.
    fn nary(&mut self, children: &[ExprId], inst: Instruction, empty: f64) {
        if children.is_empty() {
            self.emit(Instruction::PushConst(empty));
            return;
        }
        let mut tasks: Vec<Task> = Vec::with_capacity(children.len() * 2);
        tasks.push(Task::Node(children[0]));
        for &c in &children[1..] {
            tasks.push(Task::Node(c));
            tasks.push(Task::Emit(inst));
        }
        self.push_seq(&tasks);
    }

    fn lower_node(&mut self, id: ExprId) -> Result<(), SymplexError> {
        let arena = self.arena;
        if let Some(v) = self.constant_value(id) {
            self.emit(Instruction::PushConst(v));
            return Ok(());
        }
        match arena.node(id).clone() {
            ExprNode::Num(nid) => {
                // Correctly rounded: `numer as f64 / denom as f64` was
                // inf/inf = NaN for 10^400/(10^400 + 1).
                let r = arena.num(nid);
                let v = crate::base::numeric::ratio_to_f64(r).unwrap_or(if r.is_negative() {
                    f64::NEG_INFINITY
                } else {
                    f64::INFINITY
                });
                self.emit(Instruction::PushConst(v));
            }
            ExprNode::Symbol(sid) => {
                if let Some(&v) = self.constants.get(&sid) {
                    self.emit(Instruction::PushConst(v));
                } else if let Some(&slot) = self.locals.get(&sid) {
                    self.emit(Instruction::LoadLocal(slot));
                } else {
                    let name = arena.symbol_name(sid);
                    match self.vars.get(name) {
                        Some(&idx) => self.emit(Instruction::PushVar(idx)),
                        None => {
                            return Err(SymplexError::FreeSymbol {
                                name: name.to_string(),
                            });
                        }
                    }
                }
            }
            ExprNode::Pi => self.emit(Instruction::PushConst(std::f64::consts::PI)),
            ExprNode::E => self.emit(Instruction::PushConst(std::f64::consts::E)),
            ExprNode::EulerGamma => self.emit(Instruction::PushConst(
                crate::output::codegen::numeric_rt::EULER_GAMMA_F64,
            )),
            ExprNode::Catalan => self.emit(Instruction::PushConst(
                crate::output::codegen::numeric_rt::CATALAN_F64,
            )),
            ExprNode::GoldenRatio => self.emit(Instruction::PushConst(
                crate::output::codegen::numeric_rt::GOLDEN_RATIO_F64,
            )),
            ExprNode::PhysicalConstant(_, value_id) => self.work.push(Task::Node(value_id)),
            ExprNode::Infinity => self.emit(Instruction::PushConst(f64::INFINITY)),
            ExprNode::NegInfinity => self.emit(Instruction::PushConst(f64::NEG_INFINITY)),
            ExprNode::NaN | ExprNode::ComplexInfinity => {
                self.emit(Instruction::PushConst(f64::NAN))
            }
            // This VM is real-valued: intrinsically complex nodes have no f64
            // value.  Use `eval_complex64` for complex evaluation.
            ExprNode::ImaginaryUnit => return Err(self.unsupported("ImaginaryUnit")),
            // Every value the VM holds is real (or NaN, standing for a value
            // that is not): `re` is the value itself, `im` is 0.  Needed by
            // `real_root` of a symbol not known to be real, a `Piecewise` on
            // `im(x) = 0`.
            ExprNode::Re(x) => self.unary(x, Instruction::Re),
            ExprNode::Im(x) => self.unary(x, Instruction::Im),
            ExprNode::Conjugate(_) => return Err(self.unsupported("conjugate")),
            ExprNode::Arg(_) => return Err(self.unsupported("arg")),
            ExprNode::Si(_) => return Err(self.unsupported("Si")),
            ExprNode::Ci(_) => return Err(self.unsupported("Ci")),
            ExprNode::Ei(_) => return Err(self.unsupported("Ei")),
            ExprNode::Li(_) => return Err(self.unsupported("li")),
            ExprNode::Zeta(_) => return Err(self.unsupported("zeta")),
            ExprNode::Polygamma(_, _) => return Err(self.unsupported("polygamma")),
            ExprNode::KroneckerDelta(_, _) => return Err(self.unsupported("KroneckerDelta")),

            ExprNode::Add(ref children) => {
                // exp(x) + (-1) [+ rest] → exp_m1(x) [+ rest]
                if children.len() >= 2 {
                    let exp_idx = children
                        .iter()
                        .position(|&c| matches!(arena.node(c), ExprNode::Exp(_)));
                    let neg_one_idx = children.iter().position(|&c| is_neg_one(arena, c));
                    if let (Some(ei), Some(ni)) = (exp_idx, neg_one_idx)
                        && ei != ni
                        && let ExprNode::Exp(inner) = arena.node(children[ei]).clone()
                    {
                        let mut tasks = vec![Task::Node(inner), Task::Emit(Instruction::ExpM1)];
                        for (i, &c) in children.iter().enumerate() {
                            if i != ei && i != ni {
                                tasks.push(Task::Node(c));
                                tasks.push(Task::Emit(Instruction::Add));
                            }
                        }
                        self.push_seq(&tasks);
                        return Ok(());
                    }
                }
                self.nary(children, Instruction::Add, 0.0);
            }
            ExprNode::Mul(ref children) => self.nary(children, Instruction::Mul, 1.0),
            ExprNode::Min(ref children) => self.nary(children, Instruction::Min2, f64::INFINITY),
            ExprNode::Max(ref children) => {
                self.nary(children, Instruction::Max2, f64::NEG_INFINITY)
            }
            ExprNode::Pow(base, exp) => {
                if let Some(r) = arena.as_num(exp) {
                    if r.is_integer()
                        && let Some(n) = r.to_integer().to_i32()
                    {
                        self.unary(base, Instruction::Powi(n));
                        return Ok(());
                    }
                    let one: num_bigint::BigInt = One::one();
                    if *r.numer() == one {
                        if *r.denom() == num_bigint::BigInt::from(2) {
                            self.unary(base, Instruction::Sqrt);
                            return Ok(());
                        }
                        if *r.denom() == num_bigint::BigInt::from(3) {
                            self.unary(base, Instruction::Cbrt);
                            return Ok(());
                        }
                    }
                }
                // Any other exponent, a non-integer rational with an odd
                // denominator included, is `powf`: NaN for a negative base,
                // whose principal power is not real.  (Before 0.30 an odd
                // denominator took the real root, `(-8)^(1/3)` = −2 where
                // `evalf` gives `1 + √3·i`; the real root is `real_root`.)
                self.binary(base, exp, Instruction::Pow);
            }
            ExprNode::Neg(x) => self.unary(x, Instruction::Neg),
            ExprNode::Floor(x) => self.unary(x, Instruction::Floor),
            ExprNode::Ceiling(x) => self.unary(x, Instruction::Ceiling),
            ExprNode::Sin(x) => self.unary(x, Instruction::Sin),
            ExprNode::Cos(x) => self.unary(x, Instruction::Cos),
            ExprNode::Tan(x) => self.unary(x, Instruction::Tan),
            ExprNode::Exp(x) => self.unary(x, Instruction::Exp),
            ExprNode::Ln(x) => {
                // ln|Γ(u)| → lgamma(u), which does not overflow (Γ(200) does);
                // CSE keeps the pattern whole (`output::cse`).
                if let Some(u) = crate::output::codegen::abs_gamma_arg(arena, x) {
                    self.unary(u, Instruction::LnAbsGamma);
                    return Ok(());
                }
                // ln(1 + y) → ln_1p(y)
                if let ExprNode::Add(ref ch) = arena.node(x).clone()
                    && ch.len() == 2
                {
                    if is_one(arena, ch[0]) {
                        self.unary(ch[1], Instruction::Ln1p);
                        return Ok(());
                    }
                    if is_one(arena, ch[1]) {
                        self.unary(ch[0], Instruction::Ln1p);
                        return Ok(());
                    }
                }
                self.unary(x, Instruction::Ln);
            }
            ExprNode::Abs(x) => self.unary(x, Instruction::Abs),
            ExprNode::Asin(x) => self.unary(x, Instruction::Asin),
            ExprNode::Acos(x) => self.unary(x, Instruction::Acos),
            ExprNode::Atan(x) => self.unary(x, Instruction::Atan),
            ExprNode::Atan2(y, x) => self.binary(y, x, Instruction::Atan2),
            ExprNode::Sinh(x) => self.unary(x, Instruction::Sinh),
            ExprNode::Cosh(x) => self.unary(x, Instruction::Cosh),
            ExprNode::Tanh(x) => self.unary(x, Instruction::Tanh),
            ExprNode::Asinh(x) => self.unary(x, Instruction::Asinh),
            ExprNode::Acosh(x) => self.unary(x, Instruction::Acosh),
            ExprNode::Atanh(x) => self.unary(x, Instruction::Atanh),
            ExprNode::Sign(x) => self.unary(x, Instruction::Sign),
            ExprNode::Heaviside(x) => self.unary(x, Instruction::Heaviside),
            ExprNode::DiracDelta(x) => self.unary(x, Instruction::DiracDelta),

            ExprNode::Gamma(x) => self.unary(x, Instruction::Gamma),
            ExprNode::LogGamma(x) => self.unary(x, Instruction::LogGamma),
            ExprNode::Digamma(x) => self.unary(x, Instruction::Digamma),
            ExprNode::Erf(x) => self.unary(x, Instruction::Erf),
            ExprNode::Erfc(x) => self.unary(x, Instruction::Erfc),
            ExprNode::LambertW(x) => self.unary(x, Instruction::LambertW),
            ExprNode::Beta(a, b) => self.binary(a, b, Instruction::Beta),
            ExprNode::Factorial(x) => self.unary(x, Instruction::Factorial),
            ExprNode::Binomial(n, k) => self.binary(n, k, Instruction::Binomial),

            ExprNode::Piecewise(ref branches) => {
                let end = self.new_label();
                let mut tasks: Vec<Task> = Vec::with_capacity(branches.len() * 5 + 2);
                for &(value, cond) in branches.iter() {
                    let next = self.new_label();
                    tasks.push(Task::Bool(cond));
                    tasks.push(Task::JumpIfZero(next));
                    tasks.push(Task::Node(value));
                    tasks.push(Task::Jump(end));
                    tasks.push(Task::Label(next));
                }
                tasks.push(Task::Emit(Instruction::PushConst(f64::NAN)));
                tasks.push(Task::Label(end));
                self.push_seq(&tasks);
            }

            ExprNode::BoolTrue
            | ExprNode::BoolFalse
            | ExprNode::Gt(_, _)
            | ExprNode::Ge(_, _)
            | ExprNode::Eq_(_, _)
            | ExprNode::Ne(_, _)
            | ExprNode::And(_)
            | ExprNode::Or(_)
            | ExprNode::Not(_) => self.work.push(Task::Bool(id)),

            ExprNode::Apply(sid, ref args) => match arena.lib_fn(sid) {
                Some(f) => self.lower_apply(f, args)?,
                None => {
                    return Err(SymplexError::NotImplemented(format!(
                        "cannot compile `Apply({}, …)` to a numerical function",
                        arena.symbol_name(sid)
                    )));
                }
            },

            ExprNode::Derivative(_, _) => return Err(self.unsupported("Derivative")),
            ExprNode::Integral(_, _) => return Err(self.unsupported("Integral")),
            ExprNode::DefiniteIntegral(_, _, _, _) => {
                return Err(self.unsupported("DefiniteIntegral"));
            }
            ExprNode::Sum(_, _, _, _) => return Err(self.unsupported("Sum")),
            ExprNode::Product_(_, _, _, _) => return Err(self.unsupported("Product")),
            ExprNode::Limit(_, _, _) => return Err(self.unsupported("Limit")),
            ExprNode::Series(_, _, _, _) => return Err(self.unsupported("Series")),
            ExprNode::LaplaceTransform(_, _, _) => {
                return Err(self.unsupported("LaplaceTransform"));
            }
            ExprNode::InverseLaplaceTransform(_, _, _) => {
                return Err(self.unsupported("InverseLaplaceTransform"));
            }
            ExprNode::Residue(_, _, _) => return Err(self.unsupported("Residue")),
            ExprNode::RootOf(_, _, _) => return Err(self.unsupported("RootOf")),
            ExprNode::RootSum(_, _, _) => return Err(self.unsupported("RootSum")),
            ExprNode::DSolve(_, _, _) => return Err(self.unsupported("DSolve")),
            ExprNode::ConditionSet(_, _) => return Err(self.unsupported("ConditionSet")),
            ExprNode::Subs(_, _, _) => return Err(self.unsupported("Subs")),
            ExprNode::EmptySet => return Err(self.unsupported("EmptySet")),
            ExprNode::UniversalSet => return Err(self.unsupported("UniversalSet")),
            ExprNode::Interval(_, _, _) => return Err(self.unsupported("Interval")),
            ExprNode::FiniteSet(_) => return Err(self.unsupported("FiniteSet")),
            ExprNode::SetUnion(_) => return Err(self.unsupported("SetUnion")),
            ExprNode::SetIntersection(_) => return Err(self.unsupported("SetIntersection")),
            ExprNode::SetComplement(_, _) => return Err(self.unsupported("SetComplement")),
        }
        Ok(())
    }

    /// Lower a boolean-valued node to 0.0/1.0.
    fn lower_bool(&mut self, id: ExprId) -> Result<(), SymplexError> {
        let arena = self.arena;
        match arena.node(id).clone() {
            ExprNode::BoolTrue => self.emit(Instruction::PushConst(1.0)),
            ExprNode::BoolFalse => self.emit(Instruction::PushConst(0.0)),
            ExprNode::Gt(a, b) => self.binary(a, b, Instruction::Gt),
            ExprNode::Ge(a, b) => self.binary(a, b, Instruction::Ge),
            ExprNode::Eq_(a, b) => self.binary(a, b, Instruction::Eq),
            ExprNode::Ne(a, b) => self.binary(a, b, Instruction::Ne),
            ExprNode::And(ref ch) => self.nary_bool(ch, Instruction::And, 1.0),
            ExprNode::Or(ref ch) => self.nary_bool(ch, Instruction::Or, 0.0),
            ExprNode::Not(x) => self.push_seq(&[Task::Bool(x), Task::Emit(Instruction::Not)]),
            // A numeric expression in boolean position: nonzero ⇒ true.
            _ => self.work.push(Task::Node(id)),
        }
        Ok(())
    }

    fn nary_bool(&mut self, children: &[ExprId], inst: Instruction, empty: f64) {
        if children.is_empty() {
            self.emit(Instruction::PushConst(empty));
            return;
        }
        let mut tasks: Vec<Task> = Vec::with_capacity(children.len() * 2);
        tasks.push(Task::Bool(children[0]));
        for &c in &children[1..] {
            tasks.push(Task::Bool(c));
            tasks.push(Task::Emit(inst));
        }
        self.push_seq(&tasks);
    }

    /// Lower library `Apply` nodes (Bessel functions, orthogonal polynomials,
    /// integer sequences, factorial variants, inverse error functions).
    ///
    /// The VM instruction for each function is [`vm_lowering`]; the source
    /// back ends use `codegen::rt_lowering`, and a test keeps the two tables
    /// in step.
    fn lower_apply(&mut self, f: LibFn, args: &[ExprId]) -> Result<(), SymplexError> {
        let name = f.name();
        let arity_err = |n: usize| {
            SymplexError::NotImplemented(format!(
                "cannot compile `{name}` with {} argument(s) (expected {n})",
                args.len()
            ))
        };
        if f == LibFn::LambertW {
            let inst = match crate::output::codegen::lambert_rt_helper(self.arena, args)? {
                "lambert_wm1" => Instruction::LambertWm1,
                _ => Instruction::LambertW,
            };
            self.unary(args[0], inst);
            return Ok(());
        }
        match vm_lowering(f) {
            // (order, x) families with compile-time integer order.
            VmLowering::Ordered(make) => {
                let [order, x] = args else {
                    return Err(arity_err(2));
                };
                let order = self.const_i32(*order, name)?;
                self.unary(*x, make(order));
            }
            VmLowering::Unary(inst) => {
                let [x] = args else {
                    return Err(arity_err(1));
                };
                self.unary(*x, inst);
            }
            VmLowering::Binary(inst) => {
                let [a, b] = args else {
                    return Err(arity_err(2));
                };
                self.binary(*a, *b, inst);
            }
            VmLowering::Unsupported => {
                return Err(SymplexError::NotImplemented(format!(
                    "cannot compile `Apply({name}, …)` to a numerical function"
                )));
            }
        }
        Ok(())
    }
}

/// The stack-VM instruction a library function lowers to.
enum VmLowering {
    /// `make(order)` applied to the second argument.
    Ordered(fn(i32) -> Instruction),
    Unary(Instruction),
    Binary(Instruction),
    /// No `f64` routine in the shared runtime.
    Unsupported,
}

/// Exhaustive over [`LibFn`]; agrees with `codegen::rt_lowering` on which
/// functions are implemented (checked by `vm_and_rt_lowering_agree`).
fn vm_lowering(f: LibFn) -> VmLowering {
    match f {
        LibFn::BesselJ => VmLowering::Ordered(Instruction::BesselJ),
        LibFn::BesselY => VmLowering::Ordered(Instruction::BesselY),
        LibFn::BesselI => VmLowering::Ordered(Instruction::BesselI),
        LibFn::BesselK => VmLowering::Ordered(Instruction::BesselK),
        LibFn::Legendre => VmLowering::Ordered(Instruction::LegendreP),
        LibFn::ChebyshevT => VmLowering::Ordered(Instruction::ChebyshevT),
        LibFn::ChebyshevU => VmLowering::Ordered(Instruction::ChebyshevU),
        LibFn::Hermite => VmLowering::Ordered(Instruction::HermiteH),
        LibFn::Laguerre => VmLowering::Ordered(Instruction::LaguerreL),
        LibFn::Fibonacci => VmLowering::Unary(Instruction::Fibonacci),
        LibFn::Lucas => VmLowering::Unary(Instruction::Lucas),
        LibFn::Harmonic => VmLowering::Unary(Instruction::Harmonic),
        LibFn::Factorial2 => VmLowering::Unary(Instruction::Factorial2),
        LibFn::ErfInv => VmLowering::Unary(Instruction::Erfinv),
        LibFn::ErfcInv => VmLowering::Unary(Instruction::Erfcinv),
        LibFn::RisingFactorial => VmLowering::Binary(Instruction::RisingFactorial),
        LibFn::FallingFactorial => VmLowering::Binary(Instruction::FallingFactorial),
        LibFn::Subfactorial
        | LibFn::Bernoulli
        | LibFn::Catalan
        | LibFn::Bell
        | LibFn::EulerNumber
        | LibFn::Stirling1
        | LibFn::Stirling2
        | LibFn::PartitionCount
        | LibFn::LambertW
        | LibFn::Erfi
        | LibFn::ExpInt
        | LibFn::Shi
        | LibFn::Chi
        | LibFn::FresnelS
        | LibFn::FresnelC
        | LibFn::LowerGamma
        | LibFn::UpperGamma
        | LibFn::PolyLog
        | LibFn::DirichletEta
        | LibFn::AiryAi
        | LibFn::AiryBi
        | LibFn::AiryAiPrime
        | LibFn::AiryBiPrime
        | LibFn::EllipticK
        | LibFn::EllipticE
        | LibFn::EllipticF
        | LibFn::EllipticPi
        | LibFn::Gegenbauer
        | LibFn::Jacobi
        | LibFn::AssocLegendre
        | LibFn::AssocLaguerre
        | LibFn::BetaInc
        | LibFn::BetaIncRegularized => VmLowering::Unsupported,
    }
}

/// Iterative check for nodes that must bypass the `eval()` pre-pass
/// (`Piecewise` and orthogonal-polynomial `Apply` nodes).
fn needs_raw_lowering(arena: &Arena, root: ExprId) -> bool {
    let mut stack = vec![root];
    let mut seen = rustc_hash::FxHashSet::default();
    while let Some(id) = stack.pop() {
        if !seen.insert(id) {
            continue;
        }
        let node = arena.node(id);
        match node {
            ExprNode::Piecewise(_) => return true,
            ExprNode::Apply(sid, _)
                if matches!(
                    arena.lib_fn(*sid),
                    Some(
                        LibFn::Legendre
                            | LibFn::ChebyshevT
                            | LibFn::ChebyshevU
                            | LibFn::Hermite
                            | LibFn::Laguerre
                    )
                ) =>
            {
                return true;
            }
            _ => {}
        }
        node.for_each_child(|c| stack.push(c));
    }
    false
}

/// Check if a node is the constant −1 (either `Num(-1)` or `Neg(Num(1))`).
fn is_neg_one(arena: &Arena, id: ExprId) -> bool {
    if let Some(r) = arena.as_num(id) {
        return *r == num_rational::Ratio::from_integer((-1).into());
    }
    if let ExprNode::Neg(inner) = arena.node(id)
        && let Some(r) = arena.as_num(*inner)
    {
        return r.is_one();
    }
    false
}

/// Check if a node is the constant 1.
fn is_one(arena: &Arena, id: ExprId) -> bool {
    arena.as_num(id).is_some_and(|r| r.is_one())
}

/// Is the node a truth value, relation or connective (lowered to 0.0/1.0)?
fn is_bool_valued(arena: &Arena, id: ExprId) -> bool {
    matches!(
        arena.node(id),
        ExprNode::BoolTrue
            | ExprNode::BoolFalse
            | ExprNode::Gt(..)
            | ExprNode::Ge(..)
            | ExprNode::Eq_(..)
            | ExprNode::Ne(..)
            | ExprNode::And(_)
            | ExprNode::Or(_)
            | ExprNode::Not(_)
    )
}

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

    fn close(a: f64, b: f64, tol: f64) -> bool {
        (a - b).abs() <= tol * (1.0 + b.abs())
    }

    #[test]
    fn lambdify_polynomial() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        let two = a.int(2);
        let three = a.int(3);
        let x2 = a.pow(x, two);
        let three_x = a.mul(&[three, x]);
        let one = a.int(1);
        let expr = a.add(&[x2, three_x, one]);
        // f(x) = x^2 + 3x + 1
        let f = compile(&mut a, expr, &["x"]).unwrap();
        assert!((f(&[2.0]) - 11.0).abs() < 1e-10); // 4 + 6 + 1 = 11
        assert!((f(&[0.0]) - 1.0).abs() < 1e-10);
        assert!((f(&[-1.0]) - (-1.0)).abs() < 1e-10); // 1 - 3 + 1 = -1
        assert_eq!(f.arity(), 1);
    }

    #[test]
    fn lambdify_trig() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        let expr = a.sin(x);
        let f = compile(&mut a, expr, &["x"]).unwrap();
        assert!((f(&[0.0]) - 0.0).abs() < 1e-10);
        assert!((f(&[std::f64::consts::FRAC_PI_2]) - 1.0).abs() < 1e-10);
    }

    #[test]
    fn lambdify_two_vars() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        let y = a.symbol("y");
        let xy = a.mul(&[x, y]);
        let one = a.int(1);
        let expr = a.add(&[xy, one]);
        let f = compile(&mut a, expr, &["x", "y"]).unwrap();
        assert!((f(&[3.0, 4.0]) - 13.0).abs() < 1e-10);
    }

    #[test]
    fn lambdify_exp_sin() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        let sin_x = a.sin(x);
        let expr = a.exp(sin_x);
        let f = compile(&mut a, expr, &["x"]).unwrap();
        let expected = 0.5f64.sin().exp();
        assert!((f(&[0.5]) - expected).abs() < 1e-14);
    }

    #[test]
    fn lambdify_pi_e() {
        let mut a = Arena::new();
        let pi = a.pi;
        let f = compile(&mut a, pi, &[]).unwrap();
        assert!((f(&[]) - std::f64::consts::PI).abs() < 1e-10);
        let e = a.e_const;
        let f2 = compile(&mut a, e, &[]).unwrap();
        assert!((f2(&[]) - std::f64::consts::E).abs() < 1e-10);
    }

    #[test]
    fn lambdify_complex_fails() {
        let mut a = Arena::new();
        let i = a.i_unit;
        let result = compile(&mut a, i, &[]);
        assert!(matches!(result, Err(SymplexError::NotImplemented(_))));
    }

    #[test]
    fn lambdify_free_symbol_is_error() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        let y = a.symbol("y");
        let expr = a.add(&[x, y]);
        match compile(&mut a, expr, &["x"]) {
            Err(SymplexError::FreeSymbol { name }) => assert_eq!(name, "y"),
            other => panic!("expected FreeSymbol, got {other:?}"),
        }
    }

    #[test]
    fn arity_mismatch_is_nan_or_error() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        let f = compile(&mut a, x, &["x"]).unwrap();
        assert!(f.call(&[]).is_nan());
        assert!(f(&[1.0, 2.0]).is_nan());
        assert!(f.try_call(&[]).is_err());
        assert_eq!(f.try_call(&[7.0]).unwrap(), 7.0);
    }

    #[test]
    fn duplicate_param_is_error() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        assert!(matches!(
            compile(&mut a, x, &["x", "x"]),
            Err(SymplexError::InvalidArgument { .. })
        ));
    }

    #[test]
    fn division_peephole() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        let y = a.symbol("y");
        let m1 = a.int(-1);
        let inv_y = a.pow(y, m1);
        let expr = a.mul(&[x, inv_y]);
        let f = compile(&mut a, expr, &["x", "y"]).unwrap();
        assert_eq!(f(&[1.0, 3.0]), 1.0 / 3.0);
    }

    #[test]
    fn piecewise_and_booleans() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        let zero = a.int(0);
        let one = a.int(1);
        let two = a.int(2);
        let gt = a.gt(x, zero);
        let neg_x = a.neg(x);
        let eq1 = a.eq_(x, one);
        let t = a.bool_true;
        // piecewise: x==1 → 2 ; x>0 → x ; else → -x
        let pw = a.piecewise(&[(two, eq1), (x, gt), (neg_x, t)]);
        let f = compile(&mut a, pw, &["x"]).unwrap();
        assert_eq!(f(&[1.0]), 2.0);
        assert_eq!(f(&[3.0]), 3.0);
        assert_eq!(f(&[-4.0]), 4.0);
        // no matching branch → NaN
        let pw2 = a.piecewise(&[(x, gt)]);
        let g = compile(&mut a, pw2, &["x"]).unwrap();
        assert_eq!(g(&[2.0]), 2.0);
        assert!(g(&[-2.0]).is_nan());
        // nested boolean logic
        let five = a.int(5);
        let lt5 = a.gt(five, x);
        let both = a.and(&[gt, lt5]);
        let not_both = a.not(both);
        let pw3 = a.piecewise(&[(one, not_both), (zero, t)]);
        let h = compile(&mut a, pw3, &["x"]).unwrap();
        assert_eq!(h(&[3.0]), 0.0);
        assert_eq!(h(&[7.0]), 1.0);
        assert_eq!(h(&[-1.0]), 1.0);
    }

    #[test]
    fn special_functions_compile() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        let g = a.gamma(x);
        let f = compile(&mut a, g, &["x"]).unwrap();
        assert!(close(f(&[0.5]), std::f64::consts::PI.sqrt(), 1e-15));
        let e = a.erf(x);
        let f = compile(&mut a, e, &["x"]).unwrap();
        assert!(close(f(&[1.0]), 0.84270079294971486934, 1e-15));
        let w = a.lambertw(x);
        let f = compile(&mut a, w, &["x"]).unwrap();
        assert!(close(f(&[1.0]), 0.56714329040978387300, 1e-15));
        let two = a.int(2);
        let j2 = a.besselj(two, x);
        let f = compile(&mut a, j2, &["x"]).unwrap();
        assert!(close(f(&[10.0]), 0.25463031368512062253, 1e-14));
        let sym_order = a.symbol("n");
        let jn = a.besselj(sym_order, x);
        assert!(matches!(
            compile(&mut a, jn, &["x", "n"]),
            Err(SymplexError::NotImplemented(_))
        ));
        let fac = a.factorial(x);
        let f = compile(&mut a, fac, &["x"]).unwrap();
        assert_eq!(f(&[5.0]), 120.0);
        let fib = a.fibonacci(x);
        let f = compile(&mut a, fib, &["x"]).unwrap();
        assert_eq!(f(&[20.0]), 6765.0);
        // Named `Apply` inverses of erf/erfc (mpmath: erfinv(0.5), erfcinv(0.1)).
        let ei = a.lib_apply(LibFn::ErfInv, &[x]);
        let f = compile(&mut a, ei, &["x"]).unwrap();
        assert!(close(f(&[0.5]), 0.47693627620446987338, 1e-15));
        assert_eq!(f(&[1.0]), f64::INFINITY);
        assert!(f(&[1.5]).is_nan());
        let eci = a.lib_apply(LibFn::ErfcInv, &[x]);
        let f = compile(&mut a, eci, &["x"]).unwrap();
        assert!(close(f(&[0.1]), 1.1630871536766740677, 1e-15));
    }

    /// The VM and the source back ends implement the same library functions
    /// with the same shapes.
    #[test]
    fn vm_and_rt_lowering_agree() {
        use crate::output::codegen::{RtLowering, rt_lowering};
        for &f in LibFn::ALL {
            let vm = match vm_lowering(f) {
                VmLowering::Ordered(_) => "ordered",
                VmLowering::Unary(_) => "unary",
                VmLowering::Binary(_) => "binary",
                VmLowering::Unsupported => "unsupported",
            };
            let rt = match rt_lowering(f) {
                RtLowering::Ordered(_) => "ordered",
                RtLowering::Unary(_) => "unary",
                RtLowering::Binary(_) => "binary",
                RtLowering::Unsupported => "unsupported",
            };
            assert_eq!(vm, rt, "{f}");
        }
    }

    #[test]
    fn compile_many_shares_cse() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        let s = a.sin(x);
        let c = a.cos(x);
        let two = a.int(2);
        let s2 = a.pow(s, two);
        let e1 = a.add(&[s2, c]);
        let e2 = a.mul(&[s, c]);
        let f = compile_many(&mut a, &[e1, e2], &["x"]).unwrap();
        assert_eq!(f.len(), 2);
        assert_eq!(f.arity(), 1);
        let out = f.call_vec(&[0.3]);
        let (s, c) = (0.3f64.sin(), 0.3f64.cos());
        assert!(close(out[0], s * s + c, 1e-15));
        assert!(close(out[1], s * c, 1e-15));
        let mut buf = [0.0; 1];
        f.call(&[0.3], &mut buf);
        assert!(buf[0].is_nan());
        assert!(f.try_call(&[0.3], &mut buf).is_err());
        let empty = compile_many(&mut a, &[], &["x"]).unwrap();
        assert!(empty.is_empty());
        assert_eq!(empty.call_vec(&[1.0]).len(), 0);
        let empty2 = compile_many_empty(&["x", "y"]).unwrap();
        assert_eq!(empty2.arity(), 2);
        assert!(empty2.is_empty());
    }

    #[test]
    fn deep_expression_does_not_overflow_stack() {
        let mut a = Arena::new();
        let x = a.symbol("x");
        let mut e = x;
        for _ in 0..20_000 {
            e = a.sin(e);
        }
        let f = compile(&mut a, e, &["x"]).unwrap();
        assert!(f(&[1.0]).is_finite());
    }
}