alkahest-cas 3.12.0

//! Per-call wall-clock and step budgets, cooperative cancellation, and a
//! deterministic seed for search-style workloads.
//!
//! # Motivation
//!
//! A fan-out loop trying 10k candidate rewrites/integrals/Gröbner bases
//! cannot afford one pathological candidate to hang the whole batch, and an
//! orchestrator that decides "this candidate isn't worth it" needs a way to
//! *stop it now* rather than waiting for an OS-level kill (`SIGKILL`, process
//! timeout). This module gives heavy engines a cheap, structured way to bail
//! out honestly — returning [`BudgetError`] — instead of running unbounded or
//! being killed with no diagnostic.
//!
//! # Model
//!
//! A [`Budget`] is *entered* with [`enter`], which pushes it onto a
//! thread-local stack and returns a [`BudgetGuard`]. The budget stays active
//! until the guard is dropped (including on panic-unwind), mirroring the
//! `try`/`finally` discipline of Python's `with alkahest.context(budget=...)`.
//! Budgets nest like `context(...)` blocks: only the *innermost* active frame
//! is consulted by [`check`] and [`seed`] — entering a new budget shadows the
//! outer one for the scope of the block rather than combining limits with it.
//!
//! Heavy call sites ([`crate::integrate::engine`]'s top-level entry and
//! recursion boundary, [`crate::simplify::engine`]'s per-pass loop) call
//! [`check`] at a handful of strategic points. `check` is cheap when no
//! budget is active and cancellation has not been requested — it is a single
//! atomic load plus (if a budget is active) an `Instant::now()` — so it is
//! safe to call unconditionally at those boundaries.
//!
//! # Time is not the only way to run away
//!
//! [`check`] bounds an algorithm that keeps *returning* to a checkpoint. It
//! cannot bound one that leaves for a single step which allocates without
//! bound: no deadline helps if the deadline is never consulted again, and the
//! process dies of `handle_alloc_error` — an abort, not an error any caller
//! can catch.
//!
//! [`check_growth`] is the size half of the same mechanism. A step that is
//! about to grow a structure by `units` says so *before* growing, and is
//! refused if `units` is more work than the budget allows. Because the
//! failure it prevents is an abort rather than a slow answer, it applies
//! [`DEFAULT_MAX_GROWTH_UNITS`] even when no [`Budget`] is active: a caller
//! cannot opt into a limit for an allocation it did not know a call would
//! make. `Budget { max_steps: Some(n), .. }` replaces that default, so the
//! `E-BUDGET-002` remediation ("raise `Budget(max_steps=...)`") is the true
//! remedy for both halves.
//!
//! Cancellation ([`request_cancel`] / [`is_cancelled`] / [`clear_cancel`]) is
//! a single process-wide flag, deliberately *not* scoped to a thread or a
//! [`Budget`] frame: it models "the orchestrator wants the current heavy
//! operation to stop right now", e.g. because a fan-out loop decided a
//! candidate has used enough wall time across every thread working on it.
//! Call [`clear_cancel`] before starting the next candidate.
//!
//! [`seed`] exposes the active budget's seed so RNG-consuming samplers (e.g.
//! randomized modular tests, homotopy continuation start systems) can be
//! seeded deterministically from the ambient budget instead of threading an
//! explicit seed parameter through every call — two runs entering the same
//! `Budget { seed: Some(7), .. }` observe the same [`seed`] at every call site
//! that consults it.
//!
//! # Errors
//!
//! [`BudgetError`] implements [`AlkahestError`] with stable codes:
//!
//! | Code           | Variant                | Cause             |
//! |----------------|-------------------------|-------------------|
//! | `E-BUDGET-001` | [`BudgetError::WallClock`] | wall-clock deadline elapsed |
//! | `E-BUDGET-002` | [`BudgetError::Steps`]     | step counter exceeded `max_steps`, or one step asked to grow past it |
//! | `E-BUDGET-003` | [`BudgetError::Cancelled`] | [`request_cancel`] was called |
//!
//! Memory ceilings are reported through [`BudgetTrip`] rather than
//! [`BudgetError`]: `BudgetError` is a public *exhaustive* enum, so growing it
//! a `Memory` variant is a major semver break. [`check_all`] returns
//! [`BudgetTrip`], which wraps a [`BudgetError`] or carries one of the two
//! memory refusals:
//!
//! | Code           | Variant                       | Cause             |
//! |----------------|-------------------------------|-------------------|
//! | `E-BUDGET-004` | [`BudgetTrip::Memory`]        | the active budget's `max_bytes` ceiling |
//! | `E-BUDGET-005` | [`BudgetTrip::AddressSpace`]  | the process is about to exhaust `RLIMIT_AS` |
//! | `E-BUDGET-006` | [`BudgetTrip::Oversized`]     | one result is larger than the machine, or than GMP, can hold |
//!
//! See [`mod@memory`] for why the memory ceiling is enforced at these
//! checkpoints rather than inside a fallible allocator (GMP does not have
//! one, and a Rust `panic!` may not cross a C frame).

use crate::errors::AlkahestError;
use std::cell::{Cell, RefCell};
use std::fmt;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};

pub mod memory;

pub use memory::{gmp_live_bytes, install as install_memory_accounting};

// ---------------------------------------------------------------------------
// Budget
// ---------------------------------------------------------------------------

/// A per-call resource budget: an optional wall-clock limit, an optional
/// cooperative step limit, and an optional determinism seed.
///
/// Every field is optional. `Budget::default()` never trips [`check`] on its
/// own — only [`request_cancel`] can stop a call entered with a default
/// budget. This is intentional: entering an (otherwise empty) budget is how
/// a caller opts a code path into consulting [`seed`] without also imposing
/// a wall/step limit.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Budget {
    /// Wall-clock limit for the guarded block, measured from [`enter`].
    pub wall: Option<Duration>,
    /// Maximum number of work units the guarded block may spend.
    ///
    /// Counted by [`check`], which charges one unit per call, and by
    /// [`check_growth`], which charges a step for the size of the structure it
    /// is about to grow. When `None`, [`check`] is uncounted and
    /// [`check_growth`] falls back to [`DEFAULT_MAX_GROWTH_UNITS`].
    pub max_steps: Option<u64>,
    /// Determinism seed available to callers via [`seed`].
    pub seed: Option<u64>,
}

impl Budget {
    /// An empty budget — no wall/step limit, no seed. Equivalent to
    /// `Budget::default()`.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the wall-clock limit.
    pub fn with_wall(mut self, wall: Duration) -> Self {
        self.wall = Some(wall);
        self
    }

    /// Set the step limit.
    pub fn with_max_steps(mut self, max_steps: u64) -> Self {
        self.max_steps = Some(max_steps);
        self
    }

    /// Set the determinism seed.
    pub fn with_seed(mut self, seed: u64) -> Self {
        self.seed = Some(seed);
        self
    }
}

struct Frame {
    start: Instant,
    wall: Option<Duration>,
    max_steps: Option<u64>,
    steps: Cell<u64>,
    seed: Option<u64>,
    /// Ceiling on GMP bytes held live by this block, from
    /// [`enter_with_memory`].
    max_bytes: Option<u64>,
    /// GMP live-byte total when this frame was pushed, so `max_bytes` measures
    /// *this block's* appetite rather than whatever the process was already
    /// holding.
    bytes_at_entry: u64,
}

thread_local! {
    static STACK: RefCell<Vec<Frame>> = const { RefCell::new(Vec::new()) };
}

/// RAII guard returned by [`enter`].
///
/// Pops the corresponding [`Budget`] frame from the current thread's active-
/// budget stack on drop — on every exit path, including panic-unwind — so a
/// `?`-propagated error or an early `return` can never leak a stale budget
/// frame into unrelated code that runs later on the same thread.
///
/// Not [`Send`]: the stack it pops from is thread-local, so a guard created
/// on one thread must not be dropped from another.
pub struct BudgetGuard {
    _not_send: std::marker::PhantomData<*const ()>,
}

impl Drop for BudgetGuard {
    fn drop(&mut self) {
        STACK.with(|s| {
            s.borrow_mut().pop();
        });
    }
}

/// Push `budget` onto the current thread's active-budget stack.
///
/// Returns a [`BudgetGuard`]; the budget stays active — visible to [`check`],
/// [`seed`], and [`is_active`] on this thread — until the guard is dropped.
/// Budgets nest like `with` blocks: entering a new one shadows the previous
/// one until it is popped, rather than combining limits with the outer
/// frame (matching the non-merging nesting semantics of
/// `alkahest.context(...)` on the Python side).
pub fn enter(budget: Budget) -> BudgetGuard {
    enter_with_memory(budget, None)
}

/// [`enter`], plus a ceiling on the GMP memory the guarded block may hold
/// live (see [`mod@memory`]); `None` imposes no ceiling.
///
/// # Why this is not a `Budget` field
///
/// [`Budget`] has only public fields and is not `#[non_exhaustive]`, so every
/// caller may build one with a struct literal — and adding a field to such a
/// struct is a major semver break (`cargo semver-checks`'
/// `constructible_struct_adds_field`). A free function is additive, so the
/// memory ceiling travels beside the budget instead of inside it. The Python
/// binding still spells it `Budget(max_bytes=...)`, because a `@dataclass`
/// with a defaulted trailing field *is* additive there.
pub fn enter_with_memory(budget: Budget, max_bytes: Option<u64>) -> BudgetGuard {
    if max_bytes.is_some() {
        memory::install();
    }
    let frame = Frame {
        start: Instant::now(),
        wall: budget.wall,
        max_steps: budget.max_steps,
        steps: Cell::new(0),
        seed: budget.seed,
        max_bytes,
        bytes_at_entry: memory::gmp_live_bytes(),
    };
    STACK.with(|s| s.borrow_mut().push(frame));
    BudgetGuard {
        _not_send: std::marker::PhantomData,
    }
}

/// Returns `true` if a [`Budget`] is currently active on this thread.
pub fn is_active() -> bool {
    STACK.with(|s| !s.borrow().is_empty())
}

/// The seed of the innermost active [`Budget`] on this thread, or `None` if
/// no budget is active or the active budget did not set one.
pub fn seed() -> Option<u64> {
    STACK.with(|s| s.borrow().last().and_then(|f| f.seed))
}

/// The `max_bytes` ceiling of the innermost active budget frame, or `None`.
pub fn max_bytes() -> Option<u64> {
    STACK.with(|s| s.borrow().last().and_then(|f| f.max_bytes))
}

/// GMP bytes held live *by the innermost active budget frame* — the total now,
/// less the total when that frame was entered.
///
/// Zero when no budget is active. Saturating: a block allocated before the
/// frame was entered and freed inside it would otherwise underflow.
pub fn bytes_used() -> u64 {
    STACK.with(|s| {
        s.borrow()
            .last()
            .map(|f| memory::gmp_live_bytes().saturating_sub(f.bytes_at_entry))
            .unwrap_or(0)
    })
}

// ---------------------------------------------------------------------------
// Cancellation — process-wide, not scoped to a thread or Budget frame
// ---------------------------------------------------------------------------

static CANCELLED: AtomicBool = AtomicBool::new(false);

/// Request cancellation of the current cooperative operation(s).
///
/// Checked by every [`check`] call on every thread — regardless of whether a
/// [`Budget`] is active — until [`clear_cancel`] is called. Intended for an
/// orchestrator thread to stop a heavy call it decided is no longer worth
/// running, without waiting for an OS-level kill.
pub fn request_cancel() {
    CANCELLED.store(true, Ordering::SeqCst);
}

/// Clear a previously requested cancellation.
///
/// Call this before starting the next candidate in a fan-out loop — a
/// cancellation request left set would otherwise trip [`check`]
/// (`E-BUDGET-003`) immediately for every subsequent candidate.
pub fn clear_cancel() {
    CANCELLED.store(false, Ordering::SeqCst);
}

/// Returns `true` if [`request_cancel`] has been called and not yet cleared
/// by [`clear_cancel`].
pub fn is_cancelled() -> bool {
    CANCELLED.load(Ordering::SeqCst)
}

// ---------------------------------------------------------------------------
// Cooperative check
// ---------------------------------------------------------------------------

/// Cooperative checkpoint: call this at a natural short-circuit point in a
/// heavy algorithm — top-level entry, a recursion/depth-guard boundary, once
/// per major rewrite pass.
///
/// Checks, in order:
/// 1. [`is_cancelled`] — the process-wide cancellation flag.
/// 2. The innermost active [`Budget`]'s wall-clock limit, if any.
/// 3. The innermost active [`Budget`]'s step counter, if any — incremented on
///    every call, compared against `max_steps` after incrementing.
///
/// Returns `Ok(())` with no side effect (no counter increment) if no
/// [`Budget`] is active and cancellation has not been requested, so `check`
/// is cheap to call unconditionally at hot-loop boundaries even when no
/// caller has opted into a budget.
pub fn check() -> Result<(), BudgetError> {
    if is_cancelled() {
        return Err(BudgetError::Cancelled);
    }
    STACK.with(|s| {
        let stack = s.borrow();
        let Some(frame) = stack.last() else {
            return Ok(());
        };
        if let Some(wall) = frame.wall {
            let elapsed = frame.start.elapsed();
            if elapsed >= wall {
                return Err(BudgetError::WallClock {
                    limit: wall,
                    elapsed,
                });
            }
        }
        if let Some(max_steps) = frame.max_steps {
            let taken = frame.steps.get() + 1;
            frame.steps.set(taken);
            if taken > max_steps {
                return Err(BudgetError::Steps {
                    limit: max_steps,
                    taken,
                });
            }
        }
        Ok(())
    })
}

/// Work units a single growth step may ask for when the active [`Budget`]
/// sets no `max_steps`, and when no budget is active at all.
///
/// 2²⁶ units. Sized to be unreachable by any well-formed problem and cheap to
/// reach by a runaway one: the loop this was introduced for asks for 2⁶⁴
/// coefficients of a log-polynomial whose own documented degree bound is 20,
/// and every legitimate accumulator this crate builds is smaller than a
/// megabyte. A `Vec<ExprId>` of this many entries is already 256 MiB, so a
/// caller who genuinely wants more is asking for something worth having to
/// say out loud via `Budget::with_max_steps`.
///
/// It is deliberately a *default* rather than a hard cap. Growth is refused,
/// never truncated: a bound that silently returned a short answer would be the
/// wrong answer, where a refusal is an honest one.
pub const DEFAULT_MAX_GROWTH_UNITS: u64 = 1 << 26;

/// Cooperative checkpoint for a step that is about to do `units` of work *at
/// once* — typically grow a collection by that many elements.
///
/// [`check`] answers "may I keep going?", which bounds an algorithm that keeps
/// coming back to ask. This answers "may I spend this much before I come
/// back?", which is the only question a single unbounded allocation ever gets
/// to ask. Call it **before** the allocation, with the number of units it
/// would add; a refusal costs nothing, where discovering the same thing
/// afterwards costs the process.
///
/// Checks, in order:
/// 1. Everything [`check`] checks — cancellation, wall clock, step counter —
///    so a growth checkpoint is also an ordinary one and needs no companion
///    `check` call.
/// 2. `units` against the innermost active [`Budget`]'s `max_steps`, or
///    against [`DEFAULT_MAX_GROWTH_UNITS`] if it set none *or if no budget is
///    active at all*.
///
/// Point 2 is the one place in this module that does something without an
/// active budget, and it is deliberate: the failure mode is
/// `handle_alloc_error`, which aborts the process without unwinding, so there
/// is no later point at which any caller could have intervened. A caller who
/// wants more must ask for it (`Budget::with_max_steps`), which is exactly
/// what [`BudgetError::Steps`]'s remediation already tells them to do.
///
/// The comparison is per request and holds no state, so it is safe to call
/// from a loop condition and cannot latch a refusal onto unrelated later work.
pub fn check_growth(units: u64) -> Result<(), BudgetError> {
    check()?;
    let limit = STACK
        .with(|s| s.borrow().last().and_then(|f| f.max_steps))
        .unwrap_or(DEFAULT_MAX_GROWTH_UNITS);
    if units > limit {
        return Err(BudgetError::Steps {
            limit,
            taken: units,
        });
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Memory checkpoints
// ---------------------------------------------------------------------------

/// How often [`check_memory`] pays for a `/proc/self/statm` read.
///
/// The per-frame `max_bytes` ceiling is a pair of atomic loads and is checked
/// every time; the address-space probe is a file read (a few microseconds), so
/// it runs on every `PROBE_INTERVAL`-th call — often enough that a pivot loop
/// cannot climb a whole [`memory::reserve_bytes`] between probes, rarely
/// enough not to show up in a profile.
const PROBE_INTERVAL: u32 = 16;

/// How many probe intervals' worth of *observed* growth the address-space
/// guard keeps in reserve, on top of [`memory::reserve_bytes`].
///
/// A flat reserve alone is a bet that no probe interval can cross it, and that
/// bet is wrong for a solve whose matrix entries double: under a 900 MB
/// `ulimit -v` the m = 4 multinomial was seen to add 26 MB between two
/// consecutive probes. Scaling the reserve by the growth actually observed
/// makes the guard tighten exactly when the workload accelerates, and stay out
/// of the way when it does not — which matters, because the import alone maps
/// ~600 MB of address space, so a large flat reserve would refuse work that
/// fits.
const GROWTH_RESERVE_FACTOR: u64 = 4;

/// How recent the previous probe must be for the difference to count as
/// *growth* in [`GROWTH_RESERVE_FACTOR`]'s sense.
///
/// The growth term is meant to measure one probe interval of one computation.
/// The previous sample is per thread, though, and a thread's previous probe
/// can belong to a call that finished long ago — in a long-lived process
/// (a pytest run, a notebook kernel) minutes and thousands of unrelated calls
/// earlier. Everything the process mapped in between (thread stacks, malloc
/// arenas) then read as one interval's growth: a full `pytest tests/` run
/// under `ulimit -v 16G` saw 8.4 GB of it, a 34 GB "reserve", and a refusal at
/// 15.6 of 17.2 GB mapped — which cleared again at the very next probe, when
/// the difference was back to zero. That transient trip is what stopped a
/// `gcd` inside `q_zeilberger` and left the search running on unreduced
/// arithmetic. A stale sample now starts a fresh history instead; the flat
/// [`memory::reserve_bytes`] floor still applies to that first probe.
const GROWTH_WINDOW: Duration = Duration::from_secs(1);

/// The address-space reserve for a probe that sees `used` bytes mapped now,
/// given this thread's previous probe `prev = (bytes mapped, when)`.
fn address_space_reserve(limit: u64, used: u64, prev: Option<(u64, Instant)>, now: Instant) -> u64 {
    let growth = match prev {
        Some((prev_used, at)) if now.saturating_duration_since(at) <= GROWTH_WINDOW => {
            used.saturating_sub(prev_used)
        }
        _ => 0,
    };
    memory::reserve_bytes(limit).max(growth.saturating_mul(GROWTH_RESERVE_FACTOR))
}

thread_local! {
    static PROBE_TICK: Cell<u32> = const { Cell::new(0) };
    /// Address space mapped at the previous probe, and when, for the growth
    /// term above. `None` means "no history yet on this thread".
    static LAST_VSZ: Cell<Option<(u64, Instant)>> = const { Cell::new(None) };
    /// The trip behind the engine-specific error the current thread is about
    /// to return — see [`record_trip`].
    static LAST_TRIP: Cell<Option<BudgetTrip>> = const { Cell::new(None) };
}

/// [`check`] *and* the memory ceilings — the checkpoint a heavy exact-
/// arithmetic loop should call.
///
/// Returns [`BudgetTrip::Budget`] for the wall/step/cancel cases (so an
/// existing `check` call site can be upgraded without changing what it
/// reports) and [`BudgetTrip::Memory`] / [`BudgetTrip::AddressSpace`] for the
/// memory ones.
pub fn check_all() -> Result<(), BudgetTrip> {
    check()?;
    check_memory()
}

/// The memory half of [`check_all`], without the wall/step/cancel checks.
pub fn check_memory() -> Result<(), BudgetTrip> {
    let framed = STACK.with(|s| {
        let stack = s.borrow();
        stack.last().and_then(|f| {
            f.max_bytes.map(|limit| {
                (
                    limit,
                    memory::gmp_live_bytes().saturating_sub(f.bytes_at_entry),
                )
            })
        })
    });
    if let Some((limit, used)) = framed {
        if used > limit {
            return Err(BudgetTrip::Memory { limit, used });
        }
    }
    if memory::address_space_limit().is_some() {
        let probe = PROBE_TICK.with(|t| {
            let n = t.get().wrapping_add(1);
            t.set(n);
            n % PROBE_INTERVAL == 1
        });
        if probe {
            if let (Some(limit), Some(used)) =
                (memory::address_space_limit(), memory::address_space_used())
            {
                let now = Instant::now();
                let prev = LAST_VSZ.with(|c| c.replace(Some((used, now))));
                let reserve = address_space_reserve(limit, used, prev, now);
                if used.saturating_add(reserve) >= limit {
                    return Err(BudgetTrip::AddressSpace {
                        limit,
                        used,
                        reserve,
                    });
                }
            }
        }
    }
    Ok(())
}

/// Pre-flight check for a caller that is about to allocate `bytes` in one go.
///
/// Unlike [`check_memory`] this always pays for the address-space probe: a
/// site that can name its allocation size up front is exactly the site where
/// one step can cross the whole reserve, so it is worth a syscall to refuse
/// *before* the allocation rather than after the next checkpoint.
pub fn check_alloc(bytes: u64) -> Result<(), BudgetTrip> {
    check()?;
    check_alloc_memory(bytes)
}

/// The memory half of [`check_alloc`]: the active budget's `max_bytes` and
/// the process's `RLIMIT_AS`, without the wall/step/cancel checks.
fn check_alloc_memory(bytes: u64) -> Result<(), BudgetTrip> {
    let framed = STACK.with(|s| {
        let stack = s.borrow();
        stack.last().and_then(|f| {
            f.max_bytes.map(|limit| {
                (
                    limit,
                    memory::gmp_live_bytes().saturating_sub(f.bytes_at_entry),
                )
            })
        })
    });
    if let Some((limit, used)) = framed {
        if used.saturating_add(bytes) > limit {
            return Err(BudgetTrip::Memory {
                limit,
                used: used.saturating_add(bytes),
            });
        }
    }
    if let (Some(limit), Some(used)) = (memory::address_space_limit(), memory::address_space_used())
    {
        let projected = used.saturating_add(bytes);
        let reserve = memory::reserve_bytes(limit);
        if projected.saturating_add(reserve) >= limit {
            return Err(BudgetTrip::AddressSpace {
                limit,
                used: projected,
                reserve,
            });
        }
    }
    Ok(())
}

/// The largest integer GMP can represent, in bits.
///
/// An `mpz_t` records its length in limbs in a C `int`, so a value past
/// `INT_MAX` 64-bit limbs cannot exist. GMP does not fail when asked for one:
/// `mpz_pow_ui` and friends detect the overflow and raise `SIGFPE`
/// (`__gmp_overflow_in_mpz`), and FLINT's `fmpz` inherits the same ceiling.
pub const GMP_MAX_BITS: u64 = (i32::MAX as u64) * 64;

/// The largest single integer [`preflight_bignum_bits`] lets a computation
/// produce: 2³¹ bits (256 MiB, about 646 million decimal digits).
///
/// A sanity ceiling well inside [`GMP_MAX_BITS`], not a memory budget. It
/// keeps an innocent-looking call (`σ_k(2)` with `k = 2³²`, a number-field
/// power with a twelve-digit exponent) from spending minutes and gigabytes
/// before it is refused, and it keeps sizes within what `rug` reports as a
/// `u32` bit count.
pub const MAX_INTEGER_BITS: u64 = 1 << 31;

/// Pre-flight check for an allocation of `bytes` that a caller can size before
/// handing the work to GMP or FLINT, neither of which can fail gracefully.
///
/// Refuses with [`BudgetTrip::Oversized`] when `bytes` exceeds the machine's
/// physical memory ([`memory::physical_memory`]) — such a request can only end
/// in `abort()` — and otherwise defers to [`check_alloc`], which applies the
/// active budget's `max_bytes` and the process's `RLIMIT_AS`.
///
/// Only memory is consulted — not the wall clock, the step count or
/// cancellation — so a constructor that calls this refuses exactly the sizes
/// that cannot fit and nothing else. Nor does it apply a work cap of its own:
/// a request that fits is allowed however long it takes.
pub fn preflight_bytes(bytes: u64) -> Result<(), BudgetTrip> {
    if let Some(ceiling) = memory::physical_memory() {
        if bytes > ceiling {
            return Err(BudgetTrip::Oversized {
                requested: bytes,
                ceiling,
            });
        }
    }
    check_alloc_memory(bytes)
}

/// [`preflight_bytes`] for a computation whose largest integer has about
/// `bits` bits and which holds about `copies` integers of that size at once.
///
/// Also refuses, with [`BudgetTrip::Oversized`], an integer past
/// [`MAX_INTEGER_BITS`] (and so past [`GMP_MAX_BITS`], which GMP cannot
/// represent at any memory size).
pub fn preflight_bignum_bits(bits: u64, copies: u64) -> Result<(), BudgetTrip> {
    if bits > MAX_INTEGER_BITS {
        return Err(BudgetTrip::Oversized {
            requested: bits.div_ceil(8),
            ceiling: MAX_INTEGER_BITS / 8,
        });
    }
    preflight_bytes(bits.div_ceil(8).saturating_mul(copies.max(1)))
}

// ---------------------------------------------------------------------------
// Out-of-band trip reporting
// ---------------------------------------------------------------------------

/// Record the trip that is about to be reported as an engine-specific error.
///
/// Engines whose error enums are public and exhaustive (`HolonomicError`,
/// `Telescoping2dError`, `SosError`, …) cannot grow a `Budget` variant without
/// a major semver break, so they return their own "gave up" variant and leave
/// the real cause here for the bindings to pick up with [`take_trip`] — the
/// pattern [`crate::calculus::limits::last_budget_trip`] established for
/// wall-clock trips inside `LimitError::DepthExceeded`.
///
/// Call [`clear_trip`] at the outermost entry of such an engine so a stale
/// trip from an earlier call can never be attributed to this one.
pub fn record_trip(trip: BudgetTrip) {
    LAST_TRIP.with(|c| c.set(Some(trip)));
}

/// Take (and clear) the trip recorded by [`record_trip`] on this thread.
pub fn take_trip() -> Option<BudgetTrip> {
    LAST_TRIP.with(|c| c.take())
}

/// Clear any recorded trip. Call at the outermost entry of an engine that
/// reports trips out of band.
pub fn clear_trip() {
    LAST_TRIP.with(|c| c.set(None));
}

// ---------------------------------------------------------------------------
// Error type
// ---------------------------------------------------------------------------

/// A [`Budget`] was exceeded, or cancellation was requested.
///
/// This is a fine, expected answer for a fan-out search loop — not a crash —
/// so it carries a stable code and remediation like every other
/// `alkahest-core` error, rather than surfacing as a panic or an OS kill.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BudgetError {
    /// The active budget's wall-clock limit elapsed.
    WallClock { limit: Duration, elapsed: Duration },
    /// A work limit was exceeded: either the active budget's step counter
    /// passed `max_steps`, or a single step asked [`check_growth`] to spend
    /// more units at once than the budget (or
    /// [`DEFAULT_MAX_GROWTH_UNITS`]) allows. Both are the same currency, so
    /// they carry the same code and the same remedy.
    Steps { limit: u64, taken: u64 },
    /// [`request_cancel`] was called and not yet cleared.
    Cancelled,
}

impl fmt::Display for BudgetError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            BudgetError::WallClock { limit, elapsed } => write!(
                f,
                "budget exceeded: wall-clock limit {limit:?} elapsed ({elapsed:?} elapsed)"
            ),
            BudgetError::Steps { limit, taken } => write!(
                f,
                "budget exceeded: step limit {limit} reached ({taken} steps taken)"
            ),
            BudgetError::Cancelled => write!(f, "budget: operation was cancelled"),
        }
    }
}

impl std::error::Error for BudgetError {}

impl AlkahestError for BudgetError {
    fn code(&self) -> &'static str {
        match self {
            BudgetError::WallClock { .. } => "E-BUDGET-001",
            BudgetError::Steps { .. } => "E-BUDGET-002",
            BudgetError::Cancelled => "E-BUDGET-003",
        }
    }

    fn remediation(&self) -> Option<&'static str> {
        match self {
            BudgetError::WallClock { .. } => Some(
                "raise Budget(wall_ms=...), or accept a heuristic/numeric result for this \
                 candidate instead of an exact one",
            ),
            BudgetError::Steps { .. } => Some(
                "raise Budget(max_steps=...), or accept a partial/heuristic result for this \
                 candidate instead of an exact one",
            ),
            BudgetError::Cancelled => Some(
                "call alkahest.clear_cancel() (Python) or budget::clear_cancel() (Rust) before \
                 starting the next candidate",
            ),
        }
    }
}

/// Why a [`check_all`] checkpoint stopped a call.
///
/// # Why this is not a `BudgetError` variant
///
/// [`BudgetError`] is a public *exhaustive* enum: adding `Memory` to it is a
/// major semver break, and so is marking it `#[non_exhaustive]` to allow one
/// later. This enum is new, so it can be `#[non_exhaustive]` from birth and
/// grow without breaking anyone; [`From<BudgetError>`] keeps the two in one
/// `?`-chain.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum BudgetTrip {
    /// A wall-clock, step, or cancellation trip — see [`BudgetError`].
    Budget(BudgetError),
    /// The active budget's `max_bytes` ceiling was reached
    /// ([`enter_with_memory`], `Budget(max_bytes=...)` in Python).
    Memory {
        /// The ceiling, in bytes.
        limit: u64,
        /// GMP bytes held live by the guarded block when it tripped.
        used: u64,
    },
    /// The process is within [`memory::reserve_bytes`] of its `RLIMIT_AS`, so
    /// the next large allocation would `abort()` inside GMP or the Rust
    /// allocator rather than fail.
    ///
    /// Unlike [`BudgetTrip::Memory`] this fires with **no budget active**: the
    /// operator who set `ulimit -v` (or a container memory limit) already said
    /// how much the process may have, and refusing inside that number is
    /// strictly better than dying at it.
    AddressSpace {
        /// The process's soft `RLIMIT_AS`, in bytes.
        limit: u64,
        /// Address space mapped by the process when it tripped, in bytes.
        used: u64,
        /// Headroom the guard was keeping below `limit` — the flat reserve of
        /// [`memory::reserve_bytes`], widened by the growth observed between
        /// the last two probes.
        reserve: u64,
    },
    /// One result is larger than the machine can hold — more than its
    /// physical memory, or one integer past [`MAX_INTEGER_BITS`] (GMP itself
    /// stops at [`GMP_MAX_BITS`]) — so the allocation would `abort()` (or GMP
    /// would raise `SIGFPE`), or run for minutes first. Fires with no budget
    /// active; see [`preflight_bytes`] and [`preflight_bignum_bits`].
    Oversized {
        /// Estimated size of the refused allocation, in bytes.
        requested: u64,
        /// The ceiling it exceeds, in bytes.
        ceiling: u64,
    },
}

impl From<BudgetError> for BudgetTrip {
    fn from(e: BudgetError) -> Self {
        BudgetTrip::Budget(e)
    }
}

impl BudgetTrip {
    /// The [`BudgetError`] behind a wall/step/cancel trip, or `None` for a
    /// memory trip.
    pub fn budget_error(&self) -> Option<BudgetError> {
        match self {
            BudgetTrip::Budget(e) => Some(*e),
            _ => None,
        }
    }
}

impl fmt::Display for BudgetTrip {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            BudgetTrip::Budget(e) => e.fmt(f),
            BudgetTrip::Memory { limit, used } => write!(
                f,
                "budget exceeded: exact-arithmetic memory limit {limit} bytes reached \
                 ({used} bytes held live)"
            ),
            BudgetTrip::AddressSpace {
                limit,
                used,
                reserve,
            } => write!(
                f,
                "budget exceeded: refusing before the process address-space limit \
                 ({used} of {limit} bytes mapped, reserve {reserve} bytes) — the allocation \
                 that would follow cannot fail safely, GMP and the Rust allocator both abort"
            ),
            BudgetTrip::Oversized { requested, ceiling } => write!(
                f,
                "refusing a result of about {requested} bytes: it is larger than the \
                 {ceiling}-byte ceiling (physical memory, or the largest single integer \
                 computed here), and the allocation cannot fail safely — GMP and FLINT abort"
            ),
        }
    }
}

impl std::error::Error for BudgetTrip {}

impl AlkahestError for BudgetTrip {
    fn code(&self) -> &'static str {
        match self {
            BudgetTrip::Budget(e) => e.code(),
            BudgetTrip::Memory { .. } => "E-BUDGET-004",
            BudgetTrip::AddressSpace { .. } => "E-BUDGET-005",
            BudgetTrip::Oversized { .. } => "E-BUDGET-006",
        }
    }

    fn remediation(&self) -> Option<&'static str> {
        match self {
            BudgetTrip::Budget(e) => e.remediation(),
            BudgetTrip::Memory { .. } => Some(
                "raise Budget(max_bytes=...), or ask for a smaller problem (fewer unknowns, \
                 a lower order/degree) — the exact coefficients, not the shape of the system, \
                 are what grew",
            ),
            BudgetTrip::AddressSpace { .. } => Some(
                "raise the process address-space limit (ulimit -v, or the container/cgroup \
                 memory limit), or ask for a smaller problem — this refusal replaces the \
                 uncatchable abort that would otherwise follow",
            ),
            BudgetTrip::Oversized { .. } => Some(
                "ask for a smaller result (a smaller exponent, degree or shape) — no limit \
                 setting can make this one fit; this refusal replaces the uncatchable abort \
                 that would otherwise follow",
            ),
        }
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{Mutex, MutexGuard};

    /// The growth term tightens the guard only for growth observed across one
    /// recent probe interval; a stale sample from an earlier, unrelated call
    /// must not be read as growth (see [`GROWTH_WINDOW`]).
    #[test]
    fn address_space_growth_counts_only_across_a_recent_probe() {
        const GB: u64 = 1 << 30;
        let limit = 16 * GB;
        let flat = memory::reserve_bytes(limit);
        let now = Instant::now();
        let recent = now - Duration::from_millis(5);
        let stale = now - (GROWTH_WINDOW + Duration::from_secs(60));

        assert_eq!(address_space_reserve(limit, 15 * GB, None, now), flat);
        // A fast-growing computation: 1 GB in one interval reserves 4 GB.
        assert_eq!(
            address_space_reserve(limit, 15 * GB, Some((14 * GB, recent)), now),
            4 * GB
        );
        // The same difference against a sample from long ago is history, not
        // growth: this is the 16 GB full-suite case, which used to refuse at
        // 15.6 GB mapped and pass again one probe later.
        assert_eq!(
            address_space_reserve(limit, 15 * GB, Some((7 * GB, stale)), now),
            flat
        );
        // Shrinkage is never growth.
        assert_eq!(
            address_space_reserve(limit, 7 * GB, Some((15 * GB, recent)), now),
            flat
        );
    }

    /// Serializes the tests below, which share process-wide state (the memory
    /// ones read the global GMP byte counter).
    ///
    /// No test in this crate's unit-test binary may call [`request_cancel`].
    /// `CANCELLED` is a process-wide `AtomicBool` by design (see the module
    /// docs), and *every* test in the binary reaches [`check`] — `simplify`
    /// calls it once per pass and stops rewriting when it trips. A lock here
    /// serializes only this module: while one of these tests held the flag
    /// set, a quaternion test running in parallel got back an unsimplified
    /// `k·k` and failed `k² = -1`. The cancellation tests therefore live in
    /// `tests/budget_cancel.rs`, which is a separate process.
    static TEST_SERIAL: Mutex<()> = Mutex::new(());

    fn serial() -> MutexGuard<'static, ()> {
        TEST_SERIAL.lock().unwrap_or_else(|e| e.into_inner())
    }

    #[test]
    fn no_budget_active_never_trips() {
        let _serial = serial();
        assert!(!is_active());
        assert_eq!(seed(), None);
        for _ in 0..1000 {
            assert!(check().is_ok());
        }
    }

    #[test]
    fn step_budget_trips_after_limit() {
        let _serial = serial();
        let _guard = enter(Budget::new().with_max_steps(3));
        assert!(check().is_ok());
        assert!(check().is_ok());
        assert!(check().is_ok());
        let err = check().unwrap_err();
        assert_eq!(err.code(), "E-BUDGET-002");
        assert_eq!(err, BudgetError::Steps { limit: 3, taken: 4 });
    }

    #[test]
    fn wall_budget_trips_after_elapsed() {
        let _serial = serial();
        let _guard = enter(Budget::new().with_wall(Duration::from_millis(10)));
        assert!(check().is_ok());
        std::thread::sleep(Duration::from_millis(25));
        let err = check().unwrap_err();
        assert_eq!(err.code(), "E-BUDGET-001");
        assert!(matches!(err, BudgetError::WallClock { .. }));
    }

    #[test]
    fn seed_round_trips_through_active_budget() {
        let _serial = serial();
        assert_eq!(seed(), None);
        {
            let _guard = enter(Budget::new().with_seed(7));
            assert_eq!(seed(), Some(7));
        }
        // Popped: seed is no longer visible.
        assert_eq!(seed(), None);
    }

    #[test]
    fn nested_budgets_shadow_not_merge() {
        let _serial = serial();
        let _outer = enter(Budget::new().with_seed(1).with_max_steps(1000));
        assert_eq!(seed(), Some(1));
        {
            // Inner budget has no seed set — it does NOT inherit the outer
            // seed, matching alkahest.context(...)'s non-merging semantics.
            let _inner = enter(Budget::new().with_max_steps(2));
            assert_eq!(seed(), None);
            assert!(check().is_ok());
            assert!(check().is_ok());
            assert_eq!(check().unwrap_err().code(), "E-BUDGET-002");
        }
        // Back to the outer frame: its own step counter is untouched by the
        // inner frame's checks, and its seed is visible again.
        assert_eq!(seed(), Some(1));
        assert!(check().is_ok());
    }

    #[test]
    fn guard_pops_on_early_return_via_question_mark() {
        let _serial = serial();
        fn inner() -> Result<(), BudgetError> {
            let _guard = enter(Budget::new().with_max_steps(1));
            check()?;
            check()?; // trips — guard must still pop on this early return.
            unreachable!();
        }
        assert!(inner().is_err());
        assert!(!is_active());
    }

    /// The whole point of the size half: a step that asks for more work than
    /// any budget allows is refused *before* it allocates, with no budget
    /// entered at all. There is no later checkpoint to refuse at — the
    /// allocation it prevents aborts the process.
    #[test]
    fn growth_is_refused_without_an_active_budget() {
        let _serial = serial();
        assert!(!is_active());
        // A modest request is fine and does not count as a step.
        assert!(check_growth(1_000).is_ok());
        let err = check_growth(DEFAULT_MAX_GROWTH_UNITS + 1).unwrap_err();
        assert_eq!(err.code(), "E-BUDGET-002");
        assert_eq!(
            err,
            BudgetError::Steps {
                limit: DEFAULT_MAX_GROWTH_UNITS,
                taken: DEFAULT_MAX_GROWTH_UNITS + 1,
            }
        );
    }

    /// `max_steps` is the tunable for both halves, which is what makes
    /// `E-BUDGET-002`'s "raise Budget(max_steps=...)" remediation true of a
    /// growth refusal and not just of a step-count one.
    #[test]
    fn max_steps_replaces_the_default_growth_ceiling() {
        let _serial = serial();
        {
            let _guard = enter(Budget::new().with_max_steps(10));
            assert!(check_growth(10).is_ok());
            assert_eq!(
                check_growth(11).unwrap_err(),
                BudgetError::Steps {
                    limit: 10,
                    taken: 11
                }
            );
        }
        {
            // Raised well past the default: the same request now passes.
            let _guard = enter(Budget::new().with_max_steps(DEFAULT_MAX_GROWTH_UNITS * 4));
            assert!(check_growth(DEFAULT_MAX_GROWTH_UNITS + 1).is_ok());
        }
    }

    /// A growth checkpoint is also an ordinary one, so a call site needs only
    /// one of them: the wall clock still trips through it (and so does
    /// cancellation — see `tests/budget_cancel.rs`).
    #[test]
    fn growth_checkpoint_still_reports_the_wall_clock() {
        let _serial = serial();
        let _guard = enter(Budget::new().with_wall(Duration::from_millis(10)));
        assert!(check_growth(1).is_ok());
        std::thread::sleep(Duration::from_millis(25));
        assert_eq!(check_growth(1).unwrap_err().code(), "E-BUDGET-001");
    }

    /// The comparison is per request and holds no state, so a refusal cannot
    /// latch onto unrelated later work on the same thread.
    #[test]
    fn a_growth_refusal_does_not_poison_later_calls() {
        let _serial = serial();
        assert!(check_growth(DEFAULT_MAX_GROWTH_UNITS + 1).is_err());
        assert!(check_growth(1).is_ok());
        assert!(check().is_ok());
    }

    /// ~8 MB of GMP limbs, well clear of the 1 MiB ceiling the memory tests
    /// set and of any noise from tests running in parallel (the GMP live-byte
    /// total is process-wide by necessity — see `memory::gmp_live_bytes`).
    fn big_gmp_integer() -> rug::Integer {
        let mut z = rug::Integer::from(1);
        z <<= 64_000_000;
        z
    }

    #[test]
    fn memory_budget_trips_when_gmp_memory_passes_max_bytes() {
        let _serial = serial();
        let _guard = enter_with_memory(Budget::new(), Some(1 << 20));
        assert!(
            check_all().is_ok(),
            "an empty frame must not trip before anything is allocated"
        );
        let z = big_gmp_integer();
        let err = check_all().unwrap_err();
        assert_eq!(err.code(), "E-BUDGET-004");
        assert!(
            matches!(err, BudgetTrip::Memory { limit, used } if limit == 1 << 20 && used > limit),
            "{err:?}"
        );
        drop(z);
    }

    #[test]
    fn a_generous_memory_budget_does_not_trip() {
        let _serial = serial();
        // 1 TiB: nothing this process can allocate reaches it, so this pins
        // that the ceiling is a ceiling and not an unconditional refusal.
        let _guard = enter_with_memory(Budget::new(), Some(1 << 40));
        let z = big_gmp_integer();
        assert!(check_all().is_ok());
        drop(z);
    }

    #[test]
    fn check_alloc_refuses_before_the_allocation_is_made() {
        let _serial = serial();
        let _guard = enter_with_memory(Budget::new(), Some(1 << 20));
        // The refusal is entirely pre-flight — this frame has allocated
        // nothing — which is the whole point: GMP has no failure path to take
        // once it has been asked for the memory.
        let err = check_alloc(4 << 20).unwrap_err();
        assert_eq!(err.code(), "E-BUDGET-004");
        // The reported `used` is the *projected* total, i.e. it accounts for
        // the allocation that has not happened. Not asserted exactly: GMP's
        // live-byte total is process-wide (its allocation hooks are), so a
        // test running in parallel can contribute to this frame's delta.
        assert!(
            matches!(err, BudgetTrip::Memory { used, .. } if used >= 4 << 20),
            "{err:?}"
        );
    }

    #[test]
    fn max_bytes_is_visible_and_scoped_to_its_frame() {
        let _serial = serial();
        assert_eq!(max_bytes(), None);
        {
            let _guard = enter_with_memory(Budget::new(), Some(4096));
            assert_eq!(max_bytes(), Some(4096));
            // A plain `enter` shadows it, matching the non-merging nesting of
            // every other budget field.
            let _inner = enter(Budget::new());
            assert_eq!(max_bytes(), None);
        }
        assert_eq!(max_bytes(), None);
    }

    #[test]
    fn trips_round_trip_through_the_out_of_band_carrier() {
        let _serial = serial();
        clear_trip();
        assert_eq!(take_trip(), None);
        record_trip(BudgetTrip::Memory { limit: 1, used: 2 });
        let taken = take_trip().expect("recorded");
        assert_eq!(taken.code(), "E-BUDGET-004");
        // Taking clears, so a later unrelated error cannot inherit it.
        assert_eq!(take_trip(), None);
    }

    #[test]
    fn budget_trip_codes_have_remediation() {
        for trip in [
            BudgetTrip::Budget(BudgetError::Cancelled),
            BudgetTrip::Memory { limit: 1, used: 2 },
            BudgetTrip::AddressSpace {
                limit: 3,
                used: 2,
                reserve: 1,
            },
        ] {
            assert!(trip.code().starts_with("E-BUDGET-"));
            assert!(trip.remediation().is_some());
            assert!(!trip.to_string().is_empty());
        }
        assert_eq!(
            BudgetTrip::from(BudgetError::Cancelled).budget_error(),
            Some(BudgetError::Cancelled)
        );
    }

    #[test]
    fn error_codes_have_remediation() {
        for err in [
            BudgetError::WallClock {
                limit: Duration::from_secs(1),
                elapsed: Duration::from_secs(2),
            },
            BudgetError::Steps { limit: 1, taken: 2 },
            BudgetError::Cancelled,
        ] {
            assert!(err.code().starts_with("E-BUDGET-"));
            assert!(err.remediation().is_some());
            // Display must not panic and should mention something useful.
            assert!(!err.to_string().is_empty());
        }
    }
}