cljrs_runtime/tiered/

jit_state.rs

//! JIT invocation counters and native function pointer cache.
//!
//! Bridges the Tier-1 IR interpreter and the background JIT compiler
//! ([`JitBackend`]). Each JIT-eligible arity has one [`JitEntry`] keyed by
//! `ir_arity_id`. The flow is:
//!
//! 1. [`JitState::record_call`] bumps the invocation counter; crossing the
//!    threshold enqueues a compile on the runtime's backend.
//! 2. The background worker compiles the function and calls
//!    [`JitState::store_native_fn`] on the runtime that asked.
//! 3. [`JitState::get_native_fn`] returns the pointer; `call_cljrs_fn` calls it.
//! 4. [`dispatch_jit_call`] transmutes the raw pointer to the correct arity
//!    and invokes the native code.
//!
//! ## What is per-runtime and what is not
//!
//! Every *table* here belongs to one runtime, reached through
//! [`GlobalEnv::jit`](crate::env::env::GlobalEnv::jit): counters, argument
//! profiles, published pointers, OSR entries, the specialization ban list,
//! and the bootstrap watermark.  Two runtimes in one process never promote,
//! deoptimize, or invalidate each other's code.
//!
//! Three things are deliberately process-wide, because they describe the
//! process rather than a runtime:
//!
//! - **Thresholds** — configuration, set once from CLI flags or the
//!   environment before any runtime is built.
//! - **Active native frames** ([`push_jit_frame`], [`live_epochs`]) — a
//!   *thread* can have frames from several runtimes on its stack, and code
//!   reclamation asks "is any thread executing this module?".
//! - **[`dispatch_jit_call`]** — a pure transmute-and-call.

use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock, RwLock, Weak};

use cljrs_ir::IrFunction;
use cljrs_value::Value;

use crate::tiered::backend::JitBackend;
use crate::tiered::tiers::Tiers;

// ── Thresholds ────────────────────────────────────────────────────────────────

pub const DEFAULT_JIT_THRESHOLD: u32 = 1_000;

/// Tier-0 invocation count at which a function is lowered to IR in the
/// background (Phase 10.7).  Deliberately far below [`DEFAULT_JIT_THRESHOLD`]:
/// tree-walking is slow, so warm functions should reach the IR interpreter
/// quickly, while one-shot top-level code never pays for lowering.
pub const DEFAULT_IR_THRESHOLD: u32 = 50;

/// Per-process override set by the CLI or programmatic callers.
/// 0 means "read from CLJRS_IR_THRESHOLD env var or use the default".
/// `u32::MAX` disables background lowering entirely.
static IR_THRESHOLD_OVERRIDE: AtomicU32 = AtomicU32::new(0);

pub fn set_ir_threshold(t: u32) {
    IR_THRESHOLD_OVERRIDE.store(t, Ordering::Relaxed);
}

pub fn ir_threshold() -> u32 {
    let v = IR_THRESHOLD_OVERRIDE.load(Ordering::Relaxed);
    if v != 0 {
        return v;
    }
    match std::env::var("CLJRS_IR_THRESHOLD")
        .ok()
        .and_then(|s| s.parse::<u32>().ok())
    {
        // Like `--ir-threshold 0`: disable background lowering entirely.
        Some(0) => u32::MAX,
        Some(t) => t,
        None => DEFAULT_IR_THRESHOLD,
    }
}

/// Per-process override set by the CLI or programmatic callers.
/// 0 means "read from CLJRS_JIT_THRESHOLD env var or use the default".
static JIT_THRESHOLD_OVERRIDE: AtomicU32 = AtomicU32::new(0);

pub fn set_jit_threshold(t: u32) {
    JIT_THRESHOLD_OVERRIDE.store(t, Ordering::Relaxed);
}

pub fn jit_threshold() -> u32 {
    let v = JIT_THRESHOLD_OVERRIDE.load(Ordering::Relaxed);
    if v != 0 {
        return v;
    }
    std::env::var("CLJRS_JIT_THRESHOLD")
        .ok()
        .and_then(|s| s.parse::<u32>().ok())
        .unwrap_or(DEFAULT_JIT_THRESHOLD)
}

/// Per-process override; 0 means "env var or default".
static OSR_THRESHOLD_OVERRIDE: AtomicU32 = AtomicU32::new(0);

pub fn set_osr_threshold(t: u32) {
    OSR_THRESHOLD_OVERRIDE.store(t, Ordering::Relaxed);
}

/// Back-edge count at which a loop header is considered hot.  Defaults to the
/// invocation threshold ([`jit_threshold`]); override with
/// `CLJRS_OSR_THRESHOLD` or [`set_osr_threshold`].
pub fn osr_threshold() -> u32 {
    let v = OSR_THRESHOLD_OVERRIDE.load(Ordering::Relaxed);
    if v != 0 {
        return v;
    }
    std::env::var("CLJRS_OSR_THRESHOLD")
        .ok()
        .and_then(|s| s.parse::<u32>().ok())
        .unwrap_or_else(jit_threshold)
}

/// Entry-guard failures tolerated before a specialization is discarded.
pub fn deopt_limit() -> u32 {
    std::env::var("CLJRS_JIT_DEOPT_LIMIT")
        .ok()
        .and_then(|s| s.parse::<u32>().ok())
        .unwrap_or(10)
}

// ── Per-arity state ───────────────────────────────────────────────────────────

pub struct JitEntry {
    pub invocation_count: AtomicU32,
    /// A background IR-lowering request has been enqueued for this arity
    /// (Phase 10.7).  Set once when the warm threshold is crossed and the
    /// request is accepted; keeps Tier-0 dispatch from re-enqueueing.
    pub lower_queued: AtomicBool,
    pub compile_queued: AtomicBool,
    /// Finalized native function pointer, or null if not yet compiled.
    /// Calling convention: SystemV/C ABI, N ptr-sized params, one ptr-sized return.
    pub native_fn_ptr: AtomicPtr<()>,
    /// Reclamation epoch of the published native code (0 = none).  Assigned by
    /// the JIT worker when it registers the compiled module; used by code
    /// unloading to identify which `JITModule` backs this pointer and to track
    /// whether a frame executing this code is live at a safepoint.
    pub epoch: AtomicU64,
    /// Observed argument-type bitmasks, one byte per IR parameter (Phase
    /// 10.6).  OR-accumulated by [`JitState::record_call`] until the compile
    /// is queued; the JIT worker reads it to decide per-parameter
    /// specializations ([`JitState::arg_type_profile`]).
    pub arg_profile: Mutex<Vec<u8>>,
    /// Entry-guard failures of the published specialized code.  Crossing
    /// [`deopt_limit`] discards the specialization (see
    /// [`JitState::record_deopt`]).
    pub deopt_count: AtomicU32,
}

impl JitEntry {
    fn new() -> Self {
        Self {
            invocation_count: AtomicU32::new(0),
            lower_queued: AtomicBool::new(false),
            compile_queued: AtomicBool::new(false),
            native_fn_ptr: AtomicPtr::new(std::ptr::null_mut()),
            epoch: AtomicU64::new(0),
            arg_profile: Mutex::new(Vec::new()),
            deopt_count: AtomicU32::new(0),
        }
    }
}

// SAFETY: all fields are atomics, inherently Send+Sync.
unsafe impl Send for JitEntry {}
unsafe impl Sync for JitEntry {}

// ── Argument type profiles (Phase 10.6) ──────────────────────────────────────

/// Profile bitmask bits: the observed type classes of one argument position.
pub const PROFILE_LONG: u8 = 1;
pub const PROFILE_DOUBLE: u8 = 2;
pub const PROFILE_OTHER: u8 = 0x80;

/// Classify a value for the argument-type profile.
#[inline]
fn profile_tag(v: &Value) -> u8 {
    match v {
        Value::Long(_) => PROFILE_LONG,
        Value::Double(_) => PROFILE_DOUBLE,
        _ => PROFILE_OTHER,
    }
}

// ── OSR (on-stack replacement) state — Phase 10.4 ────────────────────────────
//
// A single hot call containing a `loop*`/`recur` never returns to re-dispatch,
// so the invocation counter above can never promote it.  The IR interpreter
// instead counts loop back-edges per execution; when a header crosses
// [`osr_threshold`] it requests compilation of an OSR-entry variant (built by
// `cljrs_ir::osr::build_osr_function` on the JIT worker).  Once the worker
// publishes the compiled entry here, the interpreter transfers its register
// file into the native frame at the next loop-header entry.

/// A published, compiled OSR entry for one `(arity_id, loop header)` pair.
#[derive(Clone)]
pub struct OsrSlot {
    /// Native OSR-entry code: C ABI, `live_ins.len()` `*const Value` params,
    /// one `*const Value` return.
    pub fn_ptr: *const (),
    /// Reclamation epoch of the backing module (see [`push_jit_frame`]).
    pub epoch: u64,
    /// Interpreter registers to pass, in parameter order
    /// (`cljrs_ir::osr::OsrFunction::live_ins`).
    pub live_ins: Arc<[cljrs_ir::VarId]>,
}

// SAFETY: `fn_ptr` is executable code owned by the JIT code cache; it carries
// no thread affinity.  All other fields are plain data.
unsafe impl Send for OsrSlot {}
unsafe impl Sync for OsrSlot {}

enum OsrState {
    /// Compilation requested, worker has not finished yet.
    Queued,
    /// Native entry published.
    Ready(OsrSlot),
    /// Compilation declined or failed — stop polling, stay at Tier 1.
    Failed,
}

/// Result of polling for a compiled OSR entry.
pub enum OsrPoll {
    NotRequested,
    Pending,
    Ready(OsrSlot),
    Failed,
}

// ── One runtime's JIT tables ─────────────────────────────────────────────────

/// The Tier-2 state of one runtime.
///
/// Owned by that runtime's [`Tiers`]; see the module docs for the split
/// between this and the handful of genuinely process-wide items below.
pub struct JitState {
    entries: RwLock<HashMap<u64, Arc<JitEntry>>>,
    osr: RwLock<HashMap<(u64, u32), OsrState>>,
    /// Arities whose specialization repeatedly deoptimized; the JIT worker
    /// compiles these generically (all parameters boxed).
    spec_banned: RwLock<HashSet<u64>>,
    /// Every arity id below this was defined before the compiler became ready
    /// (the `clojure.core` bootstrap); see [`JitState::is_bootstrap_arity`].
    bootstrap_watermark: AtomicU64,
    /// The compiler attached to this runtime, if any.
    backend: OnceLock<Arc<dyn JitBackend>>,
    /// Handle to the enclosing tier state, so a compile request can name the
    /// runtime the worker must publish into.
    tiers: Weak<Tiers>,
}

impl JitState {
    pub(crate) fn new(tiers: Weak<Tiers>) -> Self {
        Self {
            entries: RwLock::new(HashMap::new()),
            osr: RwLock::new(HashMap::new()),
            spec_banned: RwLock::new(HashSet::new()),
            bootstrap_watermark: AtomicU64::new(0),
            backend: OnceLock::new(),
            tiers,
        }
    }

    // ── Backend ──────────────────────────────────────────────────────────────

    /// Attach a JIT compiler to this runtime.  Idempotent: the first backend
    /// installed wins, later attempts are ignored.
    pub fn install_backend(&self, backend: Arc<dyn JitBackend>) {
        let _ = self.backend.set(backend);
    }

    /// This runtime's JIT compiler, or `None` when no JIT is linked or
    /// installed (dispatch then stops at Tier 1).
    pub fn backend(&self) -> Option<&Arc<dyn JitBackend>> {
        self.backend.get()
    }

    // ── Entry table ──────────────────────────────────────────────────────────

    fn entry(&self, arity_id: u64) -> Arc<JitEntry> {
        {
            let guard = self.entries.read().unwrap();
            if let Some(e) = guard.get(&arity_id) {
                return e.clone();
            }
        }
        self.entries
            .write()
            .unwrap()
            .entry(arity_id)
            .or_insert_with(|| Arc::new(JitEntry::new()))
            .clone()
    }

    /// Return the compiled native function pointer and its reclamation epoch
    /// for `arity_id`, if native code is currently published.
    ///
    /// The caller **must** keep the returned `epoch` live (via
    /// [`push_jit_frame`]) for the entire native call, so code unloading at a
    /// stop-the-world safepoint does not free the backing module while it
    /// executes.
    pub fn get_native_fn(&self, arity_id: u64) -> Option<(*const (), u64)> {
        let guard = self.entries.read().unwrap();
        let entry = guard.get(&arity_id)?;
        let ptr = entry.native_fn_ptr.load(Ordering::Acquire);
        if ptr.is_null() {
            None
        } else {
            let epoch = entry.epoch.load(Ordering::Acquire);
            Some((ptr as *const (), epoch))
        }
    }

    /// Publish a compiled native function pointer and its reclamation epoch
    /// for `arity_id`.  Called by the JIT worker thread after successful
    /// compilation.
    ///
    /// Stores the epoch before the pointer (release ordering) so that any
    /// reader that observes a non-null pointer also observes the matching
    /// epoch.
    pub fn store_native_fn(&self, arity_id: u64, ptr: *const (), epoch: u64) {
        let entry = self.entry(arity_id);
        entry.epoch.store(epoch, Ordering::Release);
        entry.native_fn_ptr.store(ptr as *mut (), Ordering::Release);
    }

    /// Clear the published native pointer for `arity_id` and return the epoch
    /// that was backing it, if any.
    ///
    /// Called when a var holding this function is redefined: future dispatches
    /// fall back to the interpreter immediately (the pointer is nulled), and
    /// the returned epoch is handed to the code cache so the now-superseded
    /// module is reclaimed at the next safepoint once no frame is executing
    /// it.  Also drops the per-arity table entry, keeping the table bounded
    /// across a long REPL session of redefinitions.
    pub fn take_native_epoch(&self, arity_id: u64) -> Option<u64> {
        let entry = self.entries.write().unwrap().remove(&arity_id)?;
        let ptr = entry
            .native_fn_ptr
            .swap(std::ptr::null_mut(), Ordering::AcqRel);
        if ptr.is_null() {
            None
        } else {
            Some(entry.epoch.load(Ordering::Acquire))
        }
    }

    /// Null any published native code for `arity_id` (whole-function and OSR
    /// entries) and hand the backing epochs to the code cache for reclamation.
    ///
    /// Used by cross-defn invalidation and by the var-rebind hook; a no-op
    /// when nothing was compiled or no JIT is installed.
    pub fn stale_native_code(&self, arity_id: u64) {
        let mut epochs = Vec::new();
        if let Some(epoch) = self.take_native_epoch(arity_id) {
            epochs.push(epoch);
        }
        epochs.extend(self.take_osr_epochs(arity_id));
        if let Some(backend) = self.backend() {
            for epoch in epochs {
                backend.mark_stale(epoch);
            }
        }
    }

    /// Snapshot the accumulated argument-type profile for `arity_id` (one
    /// bitmask byte per IR parameter), if any calls were profiled.
    pub fn arg_type_profile(&self, arity_id: u64) -> Option<Vec<u8>> {
        let entry = self.entries.read().unwrap().get(&arity_id)?.clone();
        let prof = entry.arg_profile.lock().unwrap();
        if prof.is_empty() {
            None
        } else {
            Some(prof.clone())
        }
    }

    /// Record a call to `arity_id`.
    ///
    /// Bumps the invocation counter and folds the call's argument types into
    /// the arity's type profile (`profile_args` are the positional call
    /// arguments matching the IR parameters; for a variadic arity the caller
    /// passes only the fixed prefix, leaving the rest-list parameter
    /// unprofiled — it is padded with [`PROFILE_OTHER`] so it can never be
    /// specialized).  When the counter crosses [`jit_threshold`] for the first
    /// time, submits a compilation request to this runtime's backend.
    ///
    /// Called on every Tier-1 IR dispatch; must be cheap.  Profiling stops
    /// once the compile is queued, so the steady-state cost is one atomic
    /// increment, one compare, and one relaxed load.
    pub fn record_call(&self, arity_id: u64, ir_func: Arc<IrFunction>, profile_args: &[Value]) {
        let entry = self.entry(arity_id);
        let count = entry.invocation_count.fetch_add(1, Ordering::Relaxed) + 1;

        if !entry.compile_queued.load(Ordering::Relaxed) {
            let n_params = ir_func.params.len();
            let mut prof = entry.arg_profile.lock().unwrap();
            if prof.len() < n_params {
                prof.resize(n_params, 0);
            }
            for (i, slot) in prof.iter_mut().enumerate().take(n_params) {
                *slot |= profile_args
                    .get(i)
                    .map(profile_tag)
                    .unwrap_or(PROFILE_OTHER);
            }
        }

        if count < jit_threshold() {
            return;
        }
        // Threshold crossed — enqueue exactly once.
        if entry.compile_queued.swap(true, Ordering::AcqRel) {
            return;
        }
        if let Some(backend) = self.backend() {
            tracing::debug!(target: "jit", "enqueue arity_id={} (count={})", arity_id, count);
            backend.enqueue_function(self.tiers.clone(), arity_id, ir_func);
        }
    }

    /// Record a Tier-0 (tree-walk) call to `arity_id`.
    ///
    /// Returns `true` exactly when the warm threshold is crossed and no
    /// lowering request has been enqueued yet — the caller should snapshot the
    /// function and enqueue, then call [`Self::mark_lower_queued`] on success.
    /// Uses `>=` so a failed enqueue (full queue) retries on the next call.
    pub fn record_interp_call(&self, arity_id: u64) -> bool {
        let threshold = ir_threshold();
        if threshold == u32::MAX {
            return false;
        }
        let entry = self.entry(arity_id);
        let count = entry.invocation_count.fetch_add(1, Ordering::Relaxed) + 1;
        count >= threshold && !entry.lower_queued.load(Ordering::Relaxed)
    }

    /// Whether a JIT compile has been queued (or published) for `arity_id`.
    /// Used by the cold-IR sweep: an arity with an in-flight compile still
    /// needs its IR as the deoptimization fallback.
    pub fn compile_queued(&self, arity_id: u64) -> bool {
        self.entries
            .read()
            .unwrap()
            .get(&arity_id)
            .is_some_and(|e| e.compile_queued.load(Ordering::Relaxed))
    }

    /// Whether the cold-IR sweep must keep `arity_id`'s IR: native code is
    /// published (the IR is its deoptimization fallback) or a compile is in
    /// flight (the worker will need it).
    pub fn pins_ir(&self, arity_id: u64) -> bool {
        self.get_native_fn(arity_id).is_some() || self.compile_queued(arity_id)
    }

    /// Whether a background lowering request is already queued for `arity_id`.
    /// Fast skip for the Tier-0 dispatch path.
    pub fn lower_queued(&self, arity_id: u64) -> bool {
        self.entries
            .read()
            .unwrap()
            .get(&arity_id)
            .is_some_and(|e| e.lower_queued.load(Ordering::Relaxed))
    }

    /// Mark that a background lowering request was accepted for `arity_id`.
    pub fn mark_lower_queued(&self, arity_id: u64) {
        self.entry(arity_id)
            .lower_queued
            .store(true, Ordering::Relaxed);
    }

    /// Clear the lowering-queued flag for `arity_id`, re-arming the dispatch
    /// seam's enqueue.  Used by the lowering worker when it abandons an arity
    /// after exhausting its rebind-retry budget.
    pub fn clear_lower_queued(&self, arity_id: u64) {
        if let Some(entry) = self.entries.read().unwrap().get(&arity_id) {
            entry.lower_queued.store(false, Ordering::Relaxed);
        }
    }

    /// Called by the lowering worker when it publishes IR for `arity_id`.
    ///
    /// Restarts the invocation counter so the JIT threshold counts pure Tier-1
    /// calls (and the argument-type profile gets a full window).  Deliberately
    /// does *not* drop the `JitEntry`: `lower_queued` must stay set, or the
    /// Tier-0 path could re-enqueue between publish and the next IR dispatch.
    pub fn on_ir_published(&self, arity_id: u64) {
        self.entry(arity_id)
            .invocation_count
            .store(0, Ordering::Relaxed);
    }

    /// Drop the `JitEntry` for `arity_id` iff it has no published native code
    /// and no queued compile.  Returns whether it was dropped.
    ///
    /// Used by the cold-IR TTL sweep: dropping the entry clears the counters
    /// and `lower_queued`, so an evicted function can re-warm from zero.
    pub fn evict_entry_if_cold(&self, arity_id: u64) -> bool {
        let mut guard = self.entries.write().unwrap();
        let Some(entry) = guard.get(&arity_id) else {
            return false;
        };
        if !entry.native_fn_ptr.load(Ordering::Acquire).is_null()
            || entry.compile_queued.load(Ordering::Relaxed)
        {
            return false;
        }
        guard.remove(&arity_id);
        true
    }

    // ── Bootstrap watermark (Phase 10.7) ─────────────────────────────────────
    //
    // Arity ids are minted from a monotonic counter, so a single snapshot
    // taken when the compiler becomes ready separates bootstrap-era
    // definitions (the clojure.core bootstrap) from everything defined
    // afterwards (user code).  Background lowering excludes bootstrap
    // arities: they were never lowered under eager lowering either (the
    // compiler was not ready when they were defined), and some bootstrap
    // patterns are known to miscompile under the JIT (see TODO.md Phase 10.7
    // notes).  User code gets the warm tier; the bootstrap stays at
    // tree-walk, exactly as before.

    /// Record the bootstrap/user boundary: every arity id strictly below `w`
    /// was defined before this runtime's compiler became ready.  Called by the
    /// runtime builder with `crate::interp::arity::next_arity_id()`.
    pub fn set_bootstrap_watermark(&self, w: u64) {
        self.bootstrap_watermark.store(w, Ordering::Relaxed);
    }

    /// Whether `arity_id` belongs to a bootstrap-era definition (excluded from
    /// background lowering).
    pub fn is_bootstrap_arity(&self, arity_id: u64) -> bool {
        arity_id < self.bootstrap_watermark.load(Ordering::Relaxed)
    }

    // ── Deoptimization (Phase 10.6) ──────────────────────────────────────────
    //
    // A specialized compilation guards its parameter types at entry; on a
    // guard failure the native code returns a unique sentinel pointer owned by
    // the compiler.  The dispatch seam detects the sentinel, re-executes the
    // call at Tier 1 (sound: guards precede all side effects), and counts the
    // failure.  Crossing the deopt limit discards the specialized code and
    // bans the arity from further specialization, so the next compile is
    // generic.

    /// Whether `result` is the deopt sentinel returned by a failed entry
    /// guard.  Always false without a backend: nothing produced native code.
    #[inline]
    pub fn is_deopt_result(&self, result: *const Value) -> bool {
        self.backend()
            .is_some_and(|b| b.deopt_sentinel() == result as usize)
    }

    /// Take (and clear) the thread's pending exception, if any.
    ///
    /// Called by the JIT-native and OSR dispatch seams immediately after
    /// native code returns.  Returns `None` when no JIT is installed.
    pub fn take_pending_exception(&self) -> Option<Value> {
        self.backend().and_then(|b| b.take_pending_exception())
    }

    /// Whether `arity_id` may be compiled with type specializations.
    pub fn specialization_allowed(&self, arity_id: u64) -> bool {
        !self.spec_banned.read().unwrap().contains(&arity_id)
    }

    /// Record an entry-guard deopt for `arity_id`.
    ///
    /// Once the failure count crosses [`deopt_limit`], the specialized code is
    /// unpublished (dispatch falls back to Tier 1 immediately), its module is
    /// handed to the code cache for reclamation, the arity is banned from
    /// re-specialization, and its invocation counter restarts so the generic
    /// recompile triggers through the ordinary hot path.
    pub fn record_deopt(&self, arity_id: u64) {
        let entry = self.entry(arity_id);
        let failures = entry.deopt_count.fetch_add(1, Ordering::Relaxed) + 1;
        tracing::debug!(
            target: "jit",
            "deopt arity_id={} (failure #{} of {})",
            arity_id,
            failures,
            deopt_limit()
        );
        if failures < deopt_limit() {
            return;
        }
        self.spec_banned.write().unwrap().insert(arity_id);
        // Unpublish + reclaim the specialized code.  `take_native_epoch` also
        // drops the JitEntry, so the arity re-counts from zero and re-enqueues
        // a (now generic) compile when hot again.
        if let Some(epoch) = self.take_native_epoch(arity_id)
            && let Some(backend) = self.backend()
        {
            tracing::debug!(
                target: "jit",
                "specialization discarded arity_id={} epoch={}",
                arity_id,
                epoch
            );
            backend.mark_stale(epoch);
        }
    }

    // ── OSR table ────────────────────────────────────────────────────────────

    /// Poll for a compiled OSR entry for the loop at `header` in `arity_id`.
    pub fn osr_poll(&self, arity_id: u64, header: u32) -> OsrPoll {
        match self.osr.read().unwrap().get(&(arity_id, header)) {
            None => OsrPoll::NotRequested,
            Some(OsrState::Queued) => OsrPoll::Pending,
            Some(OsrState::Ready(slot)) => OsrPoll::Ready(slot.clone()),
            Some(OsrState::Failed) => OsrPoll::Failed,
        }
    }

    /// Request OSR compilation for the loop at `header` in `arity_id`.
    /// Idempotent: only the first request per `(arity_id, header)` enqueues
    /// (and pays for one `IrFunction` clone); with no JIT installed the entry
    /// is marked failed so callers stop polling.
    pub fn osr_request(&self, arity_id: u64, header: u32, ir_func: &IrFunction) {
        let has_backend = self.backend().is_some();
        {
            let mut guard = self.osr.write().unwrap();
            match guard.entry((arity_id, header)) {
                std::collections::hash_map::Entry::Occupied(_) => return,
                std::collections::hash_map::Entry::Vacant(slot) => {
                    slot.insert(if has_backend {
                        OsrState::Queued
                    } else {
                        OsrState::Failed
                    });
                }
            }
        }
        if let Some(backend) = self.backend() {
            tracing::debug!(
                target: "jit",
                "osr enqueue arity_id={} header=bb{}",
                arity_id,
                header
            );
            backend.enqueue_osr(
                self.tiers.clone(),
                arity_id,
                header,
                Arc::new(ir_func.clone()),
            );
        }
    }

    /// Publish a compiled OSR entry.  Called by the JIT worker thread.
    pub fn store_osr_fn(
        &self,
        arity_id: u64,
        header: u32,
        ptr: *const (),
        epoch: u64,
        live_ins: Vec<cljrs_ir::VarId>,
    ) {
        self.osr.write().unwrap().insert(
            (arity_id, header),
            OsrState::Ready(OsrSlot {
                fn_ptr: ptr,
                epoch,
                live_ins: live_ins.into(),
            }),
        );
    }

    /// Record that OSR compilation for `(arity_id, header)` declined or
    /// failed, so interpreters stop polling and the loop stays at Tier 1.
    pub fn mark_osr_failed(&self, arity_id: u64, header: u32) {
        self.osr
            .write()
            .unwrap()
            .insert((arity_id, header), OsrState::Failed);
    }

    /// Drop all OSR entries for `arity_id` (the owning var was rebound),
    /// returning the epochs of published code so the caller can hand them to
    /// the code cache for reclamation once no frame executes them.
    pub fn take_osr_epochs(&self, arity_id: u64) -> Vec<u64> {
        let mut guard = self.osr.write().unwrap();
        let keys: Vec<(u64, u32)> = guard
            .keys()
            .filter(|(a, _)| *a == arity_id)
            .copied()
            .collect();
        let mut epochs = Vec::new();
        for key in keys {
            if let Some(OsrState::Ready(slot)) = guard.remove(&key) {
                epochs.push(slot.epoch);
            }
        }
        epochs
    }

    /// Drop all OSR entries for `arity_id` and hand their published epochs to
    /// the code cache for reclamation.  Used by the cold-IR TTL sweep:
    /// OSR-entry code is only reachable from Tier-1 interpretation of the
    /// evicted IR, so it is equally cold.  A no-op when nothing was compiled
    /// or no JIT is installed.
    pub fn stale_osr_code(&self, arity_id: u64) {
        let epochs = self.take_osr_epochs(arity_id);
        if let Some(backend) = self.backend() {
            for epoch in epochs {
                backend.mark_stale(epoch);
            }
        }
    }

    /// Every epoch this runtime still has published, whole-function and OSR.
    ///
    /// Only the [`Drop`] impl needs this: everywhere else, code is unpublished
    /// one arity at a time as it is superseded.
    fn published_epochs(&self) -> Vec<u64> {
        let mut epochs = Vec::new();
        for entry in self.entries.read().unwrap().values() {
            // An entry keeps its epoch after the pointer is nulled, so the
            // pointer — not the epoch — is what says the code is still live.
            if !entry.native_fn_ptr.load(Ordering::Acquire).is_null() {
                epochs.push(entry.epoch.load(Ordering::Acquire));
            }
        }
        for state in self.osr.read().unwrap().values() {
            if let OsrState::Ready(slot) = state {
                epochs.push(slot.epoch);
            }
        }
        epochs
    }
}

/// Release this runtime's compiled code when the runtime goes away.
///
/// The code cache is process-global (one executable-memory budget, one
/// worker), so nothing else would ever look at these modules again: their only
/// referents were the tables being dropped right here.  Staling them hands
/// them to the cache's reclaim path, which frees each one at the next
/// stop-the-world safepoint at which no thread is executing it — the same
/// live-frame scan that guards redefinition, so a native frame still on some
/// thread's stack is as safe here as it is there.
///
/// Without this, a host that creates and drops runtimes — the case
/// per-runtime tier state exists to support — would leak every module it ever
/// compiled.  The CLI, with one runtime for the life of the process, never
/// reaches it.
impl Drop for JitState {
    fn drop(&mut self) {
        let Some(backend) = self.backend.get() else {
            return;
        };
        for epoch in self.published_epochs() {
            backend.mark_stale(epoch);
        }
    }
}

// ── Active JIT frame tracking (for code unloading) ──────────────────────────────
//
// Code unloading (Phase 10.2) frees a superseded `JITModule` only once no frame
// is executing its code.  We track that precisely: each in-flight native call
// pushes its code epoch onto a per-thread stack for the duration of the call.
//
// This is per *thread*, not per runtime: one thread can have native frames from
// several runtimes on its stack, and the question reclamation asks is "is any
// thread executing this module?".
//
// At a stop-the-world safepoint every mutator thread is parked, so the reclaimer
// can read all threads' stacks to compute the set of epochs that must not be
// freed.  Cross-thread reads are sound because:
//   - Each thread mutates only its own stack, and only while running (never
//     while parked: push/pop bracket the native call and release before any
//     safepoint poll inside it).
//   - At STW all other threads are frozen, so their stacks are stable.
//
// The per-thread stack lives behind a `Mutex` that is essentially uncontended
// on the hot path (only the owning thread locks it, briefly); the reclaimer
// contends for it only at the rare STW safepoint.

struct ThreadFrames {
    stack: Mutex<Vec<u64>>,
}

static FRAME_REGISTRY: RwLock<Vec<Weak<ThreadFrames>>> = RwLock::new(Vec::new());

thread_local! {
    static MY_FRAMES: Arc<ThreadFrames> = {
        let frames = Arc::new(ThreadFrames { stack: Mutex::new(Vec::new()) });
        let mut reg = FRAME_REGISTRY.write().unwrap();
        // Opportunistically drop registrations for threads that have exited.
        reg.retain(|w| w.strong_count() > 0);
        reg.push(Arc::downgrade(&frames));
        frames
    };
}

/// RAII guard that pops one active-frame epoch on drop (including on unwind,
/// e.g. when native code throws back through the dispatch boundary).
pub struct JitFrameGuard {
    epoch: u64,
}

impl Drop for JitFrameGuard {
    fn drop(&mut self) {
        MY_FRAMES.with(|f| {
            let mut stack = f.stack.lock().unwrap();
            // Pop the matching epoch.  Frames are strictly LIFO, so the epoch
            // is almost always the top; search defensively in case of nesting.
            if let Some(pos) = stack.iter().rposition(|&e| e == self.epoch) {
                stack.remove(pos);
            }
        });
    }
}

/// Register that this thread is about to enter native code backed by `epoch`.
///
/// Returns a guard that unregisters the frame on drop.  Must wrap the native
/// call so code unloading never frees a module that is executing.
pub fn push_jit_frame(epoch: u64) -> JitFrameGuard {
    MY_FRAMES.with(|f| f.stack.lock().unwrap().push(epoch));
    JitFrameGuard { epoch }
}

/// The epoch of the innermost native frame on this thread, if any.
///
/// Used by the closure-escape path: a closure value materialized by
/// `rt_make_fn*` captures a raw pointer into the module of the currently
/// executing native code, so that module (this epoch) must be pinned against
/// reclamation.
pub fn current_jit_epoch() -> Option<u64> {
    MY_FRAMES.with(|f| f.stack.lock().unwrap().last().copied())
}

/// Collect the set of epochs with at least one active native frame across all
/// mutator threads.
///
/// **Must be called at a stop-the-world safepoint** (all other mutator threads
/// parked), so each thread's frame stack is stable while it is read.  Used by
/// the JIT code cache to decide which stale modules are safe to free.
pub fn live_epochs() -> HashSet<u64> {
    let mut live = HashSet::new();
    let reg = FRAME_REGISTRY.read().unwrap();
    for weak in reg.iter() {
        if let Some(frames) = weak.upgrade() {
            for &epoch in frames.stack.lock().unwrap().iter() {
                live.insert(epoch);
            }
        }
    }
    live
}

// ── JIT function dispatch ─────────────────────────────────────────────────────

/// Transmute `fn_ptr` to the correct arity and invoke native JIT code.
///
/// # Safety
/// - `fn_ptr` must be a valid JIT-compiled function using the default
///   platform C calling convention (SystemV on x86-64 Linux/macOS).
/// - The function must accept exactly `args.len()` parameters of type
///   `*const Value` and return `*const Value`.
/// - The returned pointer is valid until the GC collects the backing object;
///   callers must clone the `Value` before yielding a safepoint.
#[allow(clippy::too_many_arguments)]
pub unsafe fn dispatch_jit_call(fn_ptr: *const (), args: &[*const Value]) -> *const Value {
    unsafe {
        match args.len() {
            0 => {
                let f: unsafe extern "C" fn() -> *const Value = std::mem::transmute(fn_ptr);
                f()
            }
            1 => {
                let f: unsafe extern "C" fn(*const Value) -> *const Value =
                    std::mem::transmute(fn_ptr);
                f(args[0])
            }
            2 => {
                let f: unsafe extern "C" fn(*const Value, *const Value) -> *const Value =
                    std::mem::transmute(fn_ptr);
                f(args[0], args[1])
            }
            3 => {
                let f: unsafe extern "C" fn(
                    *const Value,
                    *const Value,
                    *const Value,
                ) -> *const Value = std::mem::transmute(fn_ptr);
                f(args[0], args[1], args[2])
            }
            4 => {
                let f: unsafe extern "C" fn(
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                ) -> *const Value = std::mem::transmute(fn_ptr);
                f(args[0], args[1], args[2], args[3])
            }
            5 => {
                let f: unsafe extern "C" fn(
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                ) -> *const Value = std::mem::transmute(fn_ptr);
                f(args[0], args[1], args[2], args[3], args[4])
            }
            6 => {
                let f: unsafe extern "C" fn(
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                ) -> *const Value = std::mem::transmute(fn_ptr);
                f(args[0], args[1], args[2], args[3], args[4], args[5])
            }
            7 => {
                let f: unsafe extern "C" fn(
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                ) -> *const Value = std::mem::transmute(fn_ptr);
                f(
                    args[0], args[1], args[2], args[3], args[4], args[5], args[6],
                )
            }
            8 => {
                let f: unsafe extern "C" fn(
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                    *const Value,
                ) -> *const Value = std::mem::transmute(fn_ptr);
                f(
                    args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7],
                )
            }
            n => panic!("JIT dispatch: unsupported arity {n} (max 8 in Phase 10.1)"),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tiered::tiers::Tiers;

    /// A backend that records the epochs handed to it for reclamation, so a
    /// test can assert on what the runtime released and when.
    #[derive(Default)]
    struct RecordingBackend {
        staled: Mutex<Vec<u64>>,
    }

    impl RecordingBackend {
        fn staled(&self) -> Vec<u64> {
            let mut v = self.staled.lock().unwrap().clone();
            v.sort_unstable();
            v
        }
    }

    impl JitBackend for RecordingBackend {
        fn enqueue_function(&self, _: Weak<Tiers>, _: u64, _: Arc<cljrs_ir::IrFunction>) {}
        fn enqueue_osr(&self, _: Weak<Tiers>, _: u64, _: u32, _: Arc<cljrs_ir::IrFunction>) {}
        fn mark_stale(&self, epoch: u64) {
            self.staled.lock().unwrap().push(epoch);
        }
        fn take_pending_exception(&self) -> Option<Value> {
            None
        }
        fn deopt_sentinel(&self) -> usize {
            0
        }
        fn compile_async_arity(&self, _: &Value, _: usize, _: &mut crate::env::env::Env) {}
    }

    fn jit() -> Arc<Tiers> {
        Tiers::new(0xF000_0000)
    }

    #[test]
    fn epoch_round_trips_through_store_get_take() {
        let t = jit();
        let id = 0xF100_0001;
        let ptr = 0x1234usize as *const ();
        t.jit().store_native_fn(id, ptr, 777);
        assert_eq!(t.jit().get_native_fn(id), Some((ptr, 777)));
        // take returns the epoch and removes the entry.
        assert_eq!(t.jit().take_native_epoch(id), Some(777));
        assert_eq!(t.jit().get_native_fn(id), None);
        assert_eq!(t.jit().take_native_epoch(id), None);
    }

    /// Published native code belongs to the runtime that compiled it: a
    /// second runtime asked about the same arity id sees nothing.
    #[test]
    fn native_code_is_per_runtime() {
        let a = jit();
        let b = jit();
        let id = 0xF100_0002;
        a.jit().store_native_fn(id, 0x1234usize as *const (), 778);
        assert!(a.jit().get_native_fn(id).is_some());
        assert!(b.jit().get_native_fn(id).is_none());
        assert!(!b.jit().compile_queued(id));
    }

    /// Dropping a runtime releases the code it compiled.  Nothing else ever
    /// looks at those modules again — the tables that referenced them go away
    /// with the runtime — so without this the process-global code cache would
    /// hold them forever.
    #[test]
    fn dropping_a_runtime_stales_its_published_code() {
        use cljrs_ir::VarId;
        let backend = Arc::new(RecordingBackend::default());
        let t = jit();
        t.jit().install_backend(backend.clone());

        // Two published whole-function epochs, one published OSR entry, one
        // header that failed to compile (nothing to release), and one epoch
        // already superseded by a redefinition before the drop.
        t.jit()
            .store_native_fn(0xF300_0001, 0x1000usize as *const (), 1001);
        t.jit()
            .store_native_fn(0xF300_0002, 0x2000usize as *const (), 1002);
        t.jit().store_osr_fn(
            0xF300_0001,
            4,
            0x3000usize as *const (),
            1003,
            vec![VarId(1)],
        );
        t.jit().mark_osr_failed(0xF300_0002, 9);
        t.jit()
            .store_native_fn(0xF300_0003, 0x4000usize as *const (), 1004);
        t.jit().stale_native_code(0xF300_0003);
        assert_eq!(backend.staled(), vec![1004], "redefinition released 1004");

        drop(t);
        assert_eq!(
            backend.staled(),
            vec![1001, 1002, 1003, 1004],
            "the drop must release every epoch still published"
        );
    }

    #[test]
    fn frame_guard_marks_epoch_live_then_clears() {
        let e = 0xBEEF_0001;
        assert!(!live_epochs().contains(&e));
        let guard = push_jit_frame(e);
        assert!(live_epochs().contains(&e));
        drop(guard);
        assert!(!live_epochs().contains(&e));
    }

    #[test]
    fn osr_slot_round_trips_and_rebind_takes_epochs() {
        use cljrs_ir::VarId;
        let t = jit();
        let id = 0xF200_0001;
        // Unpublished header polls as NotRequested.
        assert!(matches!(t.jit().osr_poll(id, 1), OsrPoll::NotRequested));

        let ptr = 0x5678usize as *const ();
        t.jit()
            .store_osr_fn(id, 1, ptr, 901, vec![VarId(3), VarId(4)]);
        match t.jit().osr_poll(id, 1) {
            OsrPoll::Ready(slot) => {
                assert_eq!(slot.fn_ptr, ptr);
                assert_eq!(slot.epoch, 901);
                assert_eq!(&*slot.live_ins, &[VarId(3), VarId(4)]);
            }
            _ => panic!("expected Ready"),
        }

        // A second loop in the same arity that failed to compile.
        t.jit().mark_osr_failed(id, 7);
        assert!(matches!(t.jit().osr_poll(id, 7), OsrPoll::Failed));

        // Rebinding the var takes only the published epochs and clears all
        // entries for the arity.
        let epochs = t.jit().take_osr_epochs(id);
        assert_eq!(epochs, vec![901]);
        assert!(matches!(t.jit().osr_poll(id, 1), OsrPoll::NotRequested));
        assert!(matches!(t.jit().osr_poll(id, 7), OsrPoll::NotRequested));
    }

    #[test]
    fn osr_request_without_backend_marks_failed() {
        // No backend is installed on a bare `Tiers` (that is
        // `cljrs_compiler::jit::install`'s job), so a request must
        // immediately fail closed rather than leave interpreters polling
        // forever.
        let t = jit();
        let id = 0xF200_0002;
        let ir = IrFunction::new(None, None);
        t.jit().osr_request(id, 2, &ir);
        assert!(matches!(t.jit().osr_poll(id, 2), OsrPoll::Failed));
    }

    #[test]
    fn record_interp_call_warm_lifecycle() {
        // Single test for the whole lifecycle: set_ir_threshold is a
        // process-global override, so splitting this across tests would race.
        set_ir_threshold(3);
        let t = jit();
        let id = 0xF300_0001;

        // Below the threshold: no trigger.
        assert!(!t.jit().record_interp_call(id));
        assert!(!t.jit().record_interp_call(id));
        // Crossing (and staying past) the threshold triggers until a request
        // is accepted — a full queue must be able to retry.
        assert!(t.jit().record_interp_call(id));
        assert!(t.jit().record_interp_call(id));

        // Once queued, the trigger disarms.
        t.jit().mark_lower_queued(id);
        assert!(t.jit().lower_queued(id));
        assert!(!t.jit().record_interp_call(id));

        // Worker abandons the request: re-armed.
        t.jit().clear_lower_queued(id);
        assert!(t.jit().record_interp_call(id));

        // Worker publishes: counter restarts so jit_threshold measures pure
        // Tier-1 calls; the (re-armed) trigger stays quiet below threshold.
        t.jit().on_ir_published(id);
        assert!(!t.jit().record_interp_call(id));

        // u32::MAX disables background lowering entirely.
        set_ir_threshold(u32::MAX);
        for _ in 0..10 {
            assert!(!t.jit().record_interp_call(id));
        }
        set_ir_threshold(0); // restore env/default for other tests
    }

    #[test]
    fn evict_entry_if_cold_respects_native_and_queued() {
        let t = jit();
        let id = 0xF300_0002;
        t.jit().mark_lower_queued(id); // creates the entry
        // Published native code pins the entry.
        t.jit().store_native_fn(id, 0x4242usize as *const (), 555);
        assert!(!t.jit().evict_entry_if_cold(id));
        assert_eq!(t.jit().take_native_epoch(id), Some(555)); // also drops the entry

        // A fresh cold entry is evictable.
        t.jit().mark_lower_queued(id);
        assert!(t.jit().evict_entry_if_cold(id));
        assert!(!t.jit().lower_queued(id));
        assert!(!t.jit().evict_entry_if_cold(id)); // already gone
    }

    #[test]
    fn nested_frames_pop_in_lifo_order() {
        let a = 0xBEEF_1001;
        let b = 0xBEEF_1002;
        let ga = push_jit_frame(a);
        let gb = push_jit_frame(b);
        let live = live_epochs();
        assert!(live.contains(&a) && live.contains(&b));
        drop(gb);
        assert!(live_epochs().contains(&a));
        assert!(!live_epochs().contains(&b));
        drop(ga);
        assert!(!live_epochs().contains(&a));
    }
}