dellingr 0.4.0

//! This module provides the `State` struct, which handles the primary
//! components of the VM.

mod anchor;
mod eval;
mod eval_control;
mod eval_index;
mod eval_store;
mod frame;
mod lua_val;
mod metamethod;
mod object;
mod rng;
#[cfg(feature = "snapshot")]
mod save_state;
mod stack;
mod table;
mod table_ops;

pub use anchor::Anchor;
pub use lua_val::LuaType;
pub use lua_val::RustFunc;
#[cfg(feature = "snapshot")]
pub use save_state::{LoadError, SaveDiagnostics, SaveError, SaveState};

use indexmap::IndexMap;
#[cfg(feature = "snapshot")]
use std::collections::BTreeMap;
use std::collections::BTreeSet;
use std::sync::Arc;

use super::Instr;
use super::Result;
use super::compiler;
use super::compiler::Bytecode;
use super::compiler::RuntimeCaches;
use super::cost_meter::CostMeter;
use super::error::Error;
use super::error::ErrorKind;
use super::error::StackFrame;
use super::error::TypeError;
use super::host::{DefaultCallbacks, HostCallbacks};
use super::instr::{ArgCount, Builtin, RetCount};
use super::patterns::LuaPattern;

use anchor::Registry;

pub(super) use lua_val::Val;
pub(super) use object::ObjectPtr;
use object::{GcHeap, Markable, StringPtr, UpvaluePool, UpvalueRef};
use rng::VmRng;
use table::{Table, TableShape};

/// Maximum size in bytes of any Lua string.
pub const MAX_STRING_BYTES: usize = 16 * 1024 * 1024;

const GMATCH_PATTERN_CACHE_ENTRIES: usize = 256;
const GMATCH_PATTERN_CACHE_BYTES: usize = MAX_STRING_BYTES;

struct GmatchPatternEntry {
    pattern: std::result::Result<LuaPattern, crate::patterns::PatternError>,
    byte_len: usize,
}

pub(crate) fn check_string_size(size: usize) -> Result<()> {
    if size > MAX_STRING_BYTES {
        return Err(Error::without_location(ErrorKind::StringSizeExceeded {
            size,
            limit: MAX_STRING_BYTES,
        }));
    }
    Ok(())
}

pub(crate) fn checked_string_growth(current: usize, additional: usize) -> Result<usize> {
    // Saturating rather than wrapping: an overflowing total is reported as
    // usize::MAX, which is above the cap and so rejected, instead of wrapping
    // to a small value that would pass the check.
    let size = current.saturating_add(additional);
    check_string_size(size)?;
    Ok(size)
}

/// Values temporarily held outside the visible VM stack.
pub(super) struct TransientRoots {
    values: Vec<Val>,
    suspended_envs: Vec<SuspendedEnvironment>,
}

impl TransientRoots {
    #[cfg(feature = "snapshot")]
    fn is_empty(&self) -> bool {
        self.values.is_empty() && self.suspended_envs.is_empty()
    }
}

struct SuspendedEnvironment {
    globals: IndexMap<String, Val>,
    builtins: [Val; Builtin::COUNT],
}

impl Markable for SuspendedEnvironment {
    fn mark_reachable(&self, heap: &GcHeap, worklist: &mut Vec<ObjectPtr>) {
        self.globals.mark_reachable(heap, worklist);
        self.builtins.mark_reachable(heap, worklist);
    }
}

impl Markable for TransientRoots {
    fn mark_reachable(&self, heap: &GcHeap, worklist: &mut Vec<ObjectPtr>) {
        self.values.mark_reachable(heap, worklist);
        self.suspended_envs.mark_reachable(heap, worklist);
    }
}

/// Marks all GC roots. Called before garbage collection to identify reachable objects.
///
/// This function is the single source of truth for what constitutes a GC root.
/// All allocation functions that may trigger GC must call this with the same set of roots.
///
#[hotpath::measure]
pub(super) fn mark_gc_roots(state: &State, worklist: &mut Vec<ObjectPtr>) {
    // Mark all roots - closed upvalues are now marked transitively when
    // marking LuaFn closures that reference them
    state.stack.mark_reachable(&state.heap, worklist);
    state.globals.mark_reachable(&state.heap, worklist);
    state.builtins.mark_reachable(&state.heap, worklist);
    for (identity, entry) in &state.bytecode_caches {
        debug_assert_eq!(*identity, Arc::as_ptr(&entry.bytecode) as usize);
        entry.runtime.literals.mark_reachable(&state.heap, worklist);
    }
    state.transient_roots.mark_reachable(&state.heap, worklist);
    state.registry.mark_reachable(&state.heap, worklist);
    if let Some(cache) = &state.table_library_fallback {
        cache.names.mark_reachable(&state.heap, worklist);
    }
    #[cfg(feature = "snapshot")]
    {
        // The canonical environment and its pristine baseline must survive even
        // when user code currently reaches none of it: save/load diffs live
        // against this snapshot.
        // Must go through `GcHeap::mark`, which sets the object's color before
        // queueing it. Pushing onto the worklist directly leaves the object
        // Unmarked, so its children are traced but sweep still frees it - and a
        // later save then dereferences a dangling canonical environment table.
        for ptr in state.env_tokens.keys() {
            state.heap.mark(*ptr, worklist);
        }
        for baseline in state.env_baselines.values() {
            for (key, value) in &baseline.entries {
                key.mark_reachable(&state.heap, worklist);
                value.mark_reachable(&state.heap, worklist);
            }
            if let Some(metatable) = &baseline.metatable {
                metatable.mark_reachable(&state.heap, worklist);
            }
        }
    }
    // Note: open upvalues point to stack (already marked), closed upvalues
    // are marked transitively through the closures that reference them
}

/// Cached proof that the installed table library has no extension keys.
///
/// `table` deliberately remains unrooted: rebinding invalidates this cache,
/// and ObjectPtr's generational identity makes a collected slot safe to compare
/// before it is ever dereferenced.
pub(super) struct TableLibraryFallbackCache {
    // Private, not pub(super): the fields are reached only from `vm` and its
    // descendant modules (eval_index), and `TableShape` itself is vm-private.
    table: ObjectPtr,
    names: [Val; 7],
    shape: TableShape,
}

/// Information about an active function call, used for stack traces.
#[derive(Clone)]
pub(super) struct CallInfo {
    /// The bytecode being executed.
    pub(super) bytecode: Arc<Bytecode>,
    /// Current instruction pointer.
    pub(super) ip: usize,
}

/// State-local runtime data for one immutable bytecode chunk.
#[derive(Debug)]
pub(super) struct BytecodeRuntime {
    pub(super) literals: Box<[Val]>,
    pub(super) caches: RuntimeCaches,
}

struct BytecodeCacheEntry {
    /// Pins the pointer-derived map identity against allocator ABA reuse.
    bytecode: Arc<Bytecode>,
    runtime: Arc<BytecodeRuntime>,
}

/// The main interface into the Lua VM.
pub struct State {
    /// The global environment. Uses IndexMap for deterministic iteration order
    /// (GC marking, restrict_globals). May be changed to an actual Table in the future.
    pub(super) globals: IndexMap<String, Val>,
    /// Fast array for well-known builtin globals (print, pairs, type, etc.).
    /// Indexed by Builtin enum. Avoids IndexMap lookup for common globals.
    pub(super) builtins: [Val; Builtin::COUNT],
    /// Bumped when the whole global environment is swapped.
    pub(super) globals_version: u64,
    /// Pristine table-library identity, member names, and shape for rejecting
    /// unknown plain-table field misses without probing the library.
    pub(super) table_library_fallback: Option<TableLibraryFallbackCache>,
    /// The main stack which stores values.
    pub(super) stack: Vec<Val>,
    /// The bottom index of the current frame in the stack.
    pub(super) stack_bottom: usize,
    /// The heap which holds any garbage-collected Objects.
    pub(super) heap: GcHeap,
    /// State-local runtime data, keyed by the pinned Bytecode allocation.
    bytecode_caches: IndexMap<usize, BytecodeCacheEntry>,
    /// Identities being turned into closure shells, retained during a GC sweep.
    pending_bytecode_caches: Vec<usize>,
    /// Lua closure objects removed from the visible stack while they execute.
    pub(super) transient_roots: TransientRoots,
    /// Pool for upvalue storage. Avoids per-upvalue heap allocations.
    pub(super) upvalue_pool: UpvaluePool,
    /// Open upvalues currently pointing to stack slots.
    /// Each entry is (stack_index, upvalue_ref). Kept sorted by stack_index ascending
    /// so we can efficiently close them when a function returns.
    pub(super) open_upvalues: Vec<(usize, UpvalueRef)>,
    /// Stack of call base positions for dynamic argument counting.
    /// Pushed by MarkCallBase, popped by Call(255, ...).
    /// Supports nested function calls where each level needs its own base.
    pub(super) vararg_call_bases: Vec<usize>,
    /// Stack of constructor positions for dynamic SetList(0) operations.
    pub(super) table_constructor_bases: Vec<usize>,
    /// Cost budget remaining. When this reaches 0 or below, operations with cost > 0
    /// will fail. The action that pushes you over budget completes before stopping.
    /// Uses i64 to allow going negative (the final action that exceeds budget completes).
    pub(super) cost_remaining: i64,
    /// The original cost budget (for error reporting).
    pub(super) cost_budget: i64,
    /// Whether the host explicitly configured a cost budget.
    pub(super) cost_budget_configured: bool,
    /// Total cost consumed (for reporting).
    pub(super) cost_used: u64,
    /// Compiled patterns retained for active `string.gmatch` iterators.
    ///
    /// Keys are interned string pointers, so cache lookup does not compare
    /// pattern bytes. Insertion order makes eviction deterministic.
    /// It is an implementation cache only: iterator state remains entirely in
    /// Lua values and a snapshot simply rebuilds entries on later iteration.
    gmatch_patterns: IndexMap<StringPtr, GmatchPatternEntry>,
    gmatch_pattern_bytes: usize,
    #[cfg(debug_assertions)]
    gmatch_pattern_compilations: u64,
    /// Current metamethod call depth (for __index/__newindex chains).
    /// Prevents stack overflow from circular metamethod references.
    pub(super) metamethod_depth: u32,
    /// Current function call depth. Prevents stack overflow from deep recursion.
    pub(super) call_depth: u32,
    /// Call stack for generating stack traces on errors.
    /// Each entry represents an active Lua function call.
    pub(super) call_stack: Vec<CallInfo>,
    /// Host callbacks for print output, error handling, etc.
    pub(super) callbacks: Box<dyn HostCallbacks + Send>,
    /// Current source name (for callback context).
    /// Updated when loading a new chunk.
    pub(super) current_source: Option<String>,
    /// User-defined data that RustFuncs can access.
    /// Use `set_user_data<T>()` and `user_data<T>()` to store/retrieve.
    user_data: Option<Box<dyn std::any::Any + Send>>,
    /// Seeded RNG for deterministic math.random(). Defaults to seed 0.
    /// Use `set_rng_seed()` to set a specific seed for replay.
    pub(super) rng: VmRng,
    /// Registry of values retained from Rust via `Anchor` handles. Acts as
    /// an additional GC root set; participates in `mark_gc_roots`. Carries
    /// the State's process-unique `state_id` so cross-State misuse of an
    /// `Anchor` is caught.
    pub(super) registry: Registry,
    /// Lazily assigned deterministic identities used by `string.format("%p")`.
    /// These values deliberately are not GC roots; generational keys prevent a
    /// later allocation from aliasing an identity belonging to a freed value.
    pub(super) format_pointer_ids: Vec<(Val, u64)>,
    pub(super) next_format_pointer_id: u64,
    /// Stable ids for Rust functions that are allowed to survive save/load.
    #[cfg(feature = "snapshot")]
    pub(super) rust_fns_by_id: BTreeMap<String, RustFunc>,
    /// Reverse address lookup for save-time Rust function naming.
    #[cfg(feature = "snapshot")]
    pub(super) rust_fn_ids_by_addr: BTreeMap<usize, String>,
    /// Canonical environment objects (library tables, `_G`, and `_G`'s
    /// metatable) captured once at construction, each mapped to its save token.
    /// Keyed on the objects `open_libs` built, so user shadowing of a builtin
    /// slot (`math = {}`) cannot make the save classifier mistake a user table
    /// for an environment object. See `vm/save_state.rs`.
    #[cfg(feature = "snapshot")]
    pub(super) env_tokens: BTreeMap<ObjectPtr, String>,
    /// Pristine ordered contents captured alongside `env_tokens`. These are
    /// compared at save time so mutations of rebuilt library tables persist.
    #[cfg(feature = "snapshot")]
    pub(super) env_baselines: BTreeMap<ObjectPtr, EnvBaseline>,
}

#[cfg(feature = "snapshot")]
pub(super) struct EnvBaseline {
    pub(super) entries: Vec<(Val, Val)>,
    pub(super) metatable: Option<Val>,
}

/// Maximum call depth to prevent stack overflow from deep recursion.
/// Lua's default is 200, we use 1000 for a bit more headroom.
const MAX_CALL_DEPTH: u32 = 1000;

/// Maximum values on the shared Lua/Rust value stack, not a total host-memory quota.
const MAX_STACK_SIZE: usize = 1_000_000;

// Important note on how the stack is tracked:
// A State uses a single stack for all local variables, temporary values,
// function arguments, and function return values. Both Lua frames and Rust
// frames use this stack. `self.stack_bottom` refers to the first value in the
// stack which belongs to the current frame. Note that Rust functions access
// the stack using 1-based indexing, but Lua code uses 0-based indexing.

// State marking is done through mark_gc_roots() which has direct heap access

impl State {
    const GC_INITIAL_THRESHOLD: usize = 20;

    /// Creates a new, independent state with default callbacks (stdout).
    pub fn new() -> Self {
        Self::with_callbacks(Box::new(DefaultCallbacks))
    }

    /// Creates a new state with custom host callbacks.
    ///
    /// # Example
    ///
    /// ```ignore
    /// struct MyCallbacks { output: Vec<String> }
    /// impl HostCallbacks for MyCallbacks {
    ///     fn on_print(&mut self, _source: Option<&str>, _line: u32, message: &str) {
    ///         self.output.push(message.to_string());
    ///     }
    /// }
    ///
    /// let mut state = State::with_callbacks(Box::new(MyCallbacks { output: vec![] }));
    /// ```
    pub fn with_callbacks(callbacks: Box<dyn HostCallbacks + Send>) -> Self {
        let mut me = Self::empty_with_callbacks(callbacks);
        me.open_libs()
            .expect("standard library initialization starts with an empty value stack");
        me
    }

    /// Record the canonical environment objects right after `open_libs`, before
    /// any user code can shadow a builtin slot. Both save (object -> token) and
    /// load (token -> object) resolve through this snapshot.
    #[cfg(feature = "snapshot")]
    pub(crate) fn capture_env_tokens(&mut self) {
        let mut map = BTreeMap::new();
        for slot in [Builtin::Math, Builtin::String, Builtin::Table, Builtin::G] {
            if let Val::Obj(ptr) = self.builtins[slot as usize] {
                map.insert(ptr, slot.name().to_string());
                if slot == Builtin::G
                    && let Some(table) = self.heap.as_table_ref(ptr)
                    && let Some(mt) = table.get_metatable()
                {
                    map.insert(mt, "_G.metatable".to_string());
                }
            }
        }
        let baselines = map
            .keys()
            .map(|ptr| {
                let table = self
                    .heap
                    .as_table_ref(*ptr)
                    .expect("canonical environment object is a table");
                (
                    *ptr,
                    EnvBaseline {
                        entries: table.entries(),
                        metatable: table.get_metatable().map(Val::Obj),
                    },
                )
            })
            .collect();
        self.env_tokens = map;
        self.env_baselines = baselines;
    }

    /// Creates a new state without opening any of the standard libs.
    /// The global namespace of this state is entirely empty. This corresponds
    /// to the `lua_newstate' function in the C API.
    pub fn empty() -> Self {
        Self::empty_with_callbacks(Box::new(DefaultCallbacks))
    }

    /// Creates an empty state with custom callbacks.
    pub(crate) fn empty_with_callbacks(callbacks: Box<dyn HostCallbacks + Send>) -> Self {
        let state_id = anchor::next_state_id();
        Self {
            globals: IndexMap::new(),
            builtins: std::array::from_fn(|_| Val::Nil),
            globals_version: 0,
            table_library_fallback: None,
            stack: Vec::with_capacity(256), // Pre-size for typical function depth * locals
            stack_bottom: 0,
            heap: GcHeap::with_threshold(Self::GC_INITIAL_THRESHOLD),
            bytecode_caches: IndexMap::new(),
            pending_bytecode_caches: Vec::new(),
            transient_roots: TransientRoots {
                values: Vec::with_capacity(64),
                suspended_envs: Vec::new(),
            },
            upvalue_pool: UpvaluePool::new(),
            open_upvalues: Vec::new(),
            vararg_call_bases: Vec::new(),
            table_constructor_bases: Vec::new(),
            cost_remaining: i64::MAX,
            cost_budget: i64::MAX,
            cost_budget_configured: false,
            cost_used: 0,
            gmatch_patterns: IndexMap::new(),
            gmatch_pattern_bytes: 0,
            #[cfg(debug_assertions)]
            gmatch_pattern_compilations: 0,
            metamethod_depth: 0,
            call_depth: 0,
            call_stack: Vec::with_capacity(64), // Pre-size for call stack
            callbacks,
            current_source: None,
            user_data: None,
            rng: VmRng::seed_from_u64(0),
            registry: Registry::new(state_id),
            format_pointer_ids: Vec::new(),
            next_format_pointer_id: 1,
            #[cfg(feature = "snapshot")]
            rust_fns_by_id: BTreeMap::new(),
            #[cfg(feature = "snapshot")]
            rust_fn_ids_by_addr: BTreeMap::new(),
            #[cfg(feature = "snapshot")]
            env_tokens: BTreeMap::new(),
            #[cfg(feature = "snapshot")]
            env_baselines: BTreeMap::new(),
        }
    }

    pub(crate) fn reserve_stdlib_capacity(&mut self) {
        self.globals.reserve(Builtin::COUNT + 4);
        self.heap.reserve(8, 96);
    }

    /// Sets the RNG seed for deterministic math.random() behavior.
    pub fn set_rng_seed(&mut self, seed: u64) {
        self.rng = VmRng::seed_from_u64(seed);
    }

    /// Sets the cost budget for this VM.
    /// When the budget is exhausted, operations with cost > 0 will fail.
    /// The action that pushes you over budget always completes before stopping.
    pub fn set_cost_budget(&mut self, budget: i64) {
        self.cost_budget = budget;
        self.cost_remaining = budget;
        self.cost_budget_configured = true;
        self.cost_used = 0;
    }

    /// Returns the total cost consumed since the last budget reset.
    pub fn cost_used(&self) -> u64 {
        self.cost_used
    }

    /// Returns the cost remaining in the budget.
    /// Can be negative if the last action pushed over budget.
    pub fn cost_remaining(&self) -> i64 {
        self.cost_remaining
    }

    /// Test instrumentation for the gmatch compilation cache.
    #[cfg(debug_assertions)]
    #[doc(hidden)]
    pub fn gmatch_pattern_compilations(&self) -> u64 {
        self.gmatch_pattern_compilations
    }

    /// Consume cost from the budget. Returns an error if budget is exhausted
    /// and cost > 0. The action that pushes you over budget completes before
    /// stopping (checked at the START of each operation).
    ///
    /// Use this in RustFuncs to charge for expensive operations.
    #[inline(always)]
    pub fn consume_cost(&mut self, cost: u64) -> Result<()> {
        if cost > 0 && self.cost_remaining <= 0 {
            return Err(self.error(ErrorKind::BudgetExceeded {
                used: self.cost_used,
                budget: self.cost_budget,
            }));
        }
        self.cost_remaining = self.cost_remaining.saturating_sub_unsigned(cost);
        self.cost_used = self.cost_used.saturating_add(cost);
        Ok(())
    }

    /// Returns a meter for runtime work that does not need VM state.
    pub(crate) fn cost_meter(&mut self) -> CostMeter<'_> {
        if self.cost_budget_configured {
            CostMeter::finite_budget(&mut self.cost_remaining, &mut self.cost_used)
        } else {
            CostMeter::count_only(&mut self.cost_used)
        }
    }

    /// Compile and retain a `string.gmatch` pattern on its first iteration.
    ///
    /// Calling this from the iterator rather than its constructor preserves
    /// gmatch's deferred pattern-validation behaviour.
    pub(crate) fn memoize_gmatch_pattern(&mut self, idx: isize) -> Result<()> {
        let idx = self.convert_idx(idx)?;
        let val = &self.stack[idx];
        let ptr = val.as_string_ptr().ok_or_else(|| {
            Error::without_location(ErrorKind::ArgError(crate::error::ArgError {
                arg_number: idx as isize + 1,
                func_name: None,
                expected: Some(LuaType::String),
                received: Some(val.typ(&self.heap)),
            }))
        })?;
        if let Some(cached) = self.gmatch_patterns.get(&ptr) {
            return cached
                .pattern
                .as_ref()
                .map(|_| ())
                .map_err(|err| self.error(ErrorKind::RuntimeError(err.to_string())));
        }
        let pattern_len = val
            .as_string(&self.heap)
            .expect("string pointer came from a string value")
            .len();
        self.consume_cost(pattern_len.max(1) as u64)?;
        #[cfg(debug_assertions)]
        {
            self.gmatch_pattern_compilations = self.gmatch_pattern_compilations.saturating_add(1);
        }
        let compiled = LuaPattern::from_bytes_try(
            self.stack[idx]
                .as_string(&self.heap)
                .expect("string pointer came from a string value"),
        );
        let error = compiled.as_ref().err().cloned();
        while self.gmatch_patterns.len() >= GMATCH_PATTERN_CACHE_ENTRIES
            || self.gmatch_pattern_bytes.saturating_add(pattern_len) > GMATCH_PATTERN_CACHE_BYTES
        {
            let (_, evicted) = self
                .gmatch_patterns
                .shift_remove_index(0)
                .expect("a nonempty gmatch cache has an oldest entry");
            self.gmatch_pattern_bytes = self.gmatch_pattern_bytes.saturating_sub(evicted.byte_len);
        }
        self.gmatch_pattern_bytes = self.gmatch_pattern_bytes.saturating_add(pattern_len);
        self.gmatch_patterns.insert(
            ptr,
            GmatchPatternEntry {
                pattern: compiled,
                byte_len: pattern_len,
            },
        );
        error.map_or(Ok(()), |err| {
            Err(self.error(ErrorKind::RuntimeError(err.to_string())))
        })
    }

    /// Returns a gmatch subject borrow, its cached matcher, and a cost meter.
    ///
    /// These values come from disjoint `State` fields. Keeping the split here
    /// lets the iterator match directly against the stack-resident subject.
    pub(crate) fn gmatch_subject_matcher_and_cost_meter(
        &mut self,
        subject_idx: isize,
        pattern_idx: isize,
    ) -> Result<(&[u8], &mut LuaPattern, CostMeter<'_>)> {
        let subject_idx = self.convert_idx(subject_idx)?;
        let pattern_idx = self.convert_idx(pattern_idx)?;
        let State {
            stack,
            heap,
            cost_budget_configured,
            cost_remaining,
            cost_used,
            gmatch_patterns,
            ..
        } = self;
        let subject_val = &stack[subject_idx];
        let subject = subject_val.as_string(heap).ok_or_else(|| {
            Error::without_location(ErrorKind::ArgError(crate::error::ArgError {
                arg_number: subject_idx as isize + 1,
                func_name: None,
                expected: Some(LuaType::String),
                received: Some(subject_val.typ(heap)),
            }))
        })?;
        let pattern_val = &stack[pattern_idx];
        let pattern_ptr = pattern_val.as_string_ptr().ok_or_else(|| {
            Error::without_location(ErrorKind::ArgError(crate::error::ArgError {
                arg_number: pattern_idx as isize + 1,
                func_name: None,
                expected: Some(LuaType::String),
                received: Some(pattern_val.typ(heap)),
            }))
        })?;
        // Both of these are caller-ordering invariants rather than script-
        // reachable states: the iterator memoizes before matching, and a cached
        // compilation failure is surfaced before we get here. They are still
        // reported rather than asserted, so a future caller that gets the order
        // wrong sees an error instead of a panic in a host callback.
        let entry = gmatch_patterns.get_mut(&pattern_ptr).ok_or_else(|| {
            Error::without_location(ErrorKind::InternalError(
                "gmatch pattern was not memoized before matching".into(),
            ))
        })?;
        let matcher = match entry.pattern.as_mut() {
            Ok(matcher) => matcher,
            Err(_) => {
                return Err(Error::without_location(ErrorKind::InternalError(
                    "a failed gmatch compilation must be reported before matching".into(),
                )));
            }
        };
        let meter = if *cost_budget_configured {
            CostMeter::finite_budget(cost_remaining, cost_used)
        } else {
            CostMeter::count_only(cost_used)
        };
        Ok((subject, matcher, meter))
    }

    pub(crate) fn budget_exceeded_error(&self) -> Error {
        self.error(ErrorKind::BudgetExceeded {
            used: self.cost_used,
            budget: self.cost_budget,
        })
    }

    // ========================================================================
    // User data
    // ========================================================================

    /// Store arbitrary user data that RustFuncs can access.
    ///
    /// Useful for passing context to Rust callbacks, like a command collector.
    /// `T` must be `Send` because `State` is `Send`: any data the embedder
    /// hands to the VM must be safe to move across threads with the State.
    ///
    /// # Example
    /// ```ignore
    /// let collector = Arc::new(Mutex::new(CommandCollector::default()));
    /// state.set_user_data(collector.clone());
    ///
    /// state.push_rust_fn(|state| {
    ///     let collector = state.user_data::<Arc<Mutex<CommandCollector>>>().unwrap();
    ///     collector.lock().unwrap().turn = Some(0.5);
    ///     Ok(0)
    /// })?;
    /// ```
    pub fn set_user_data<T: Send + 'static>(&mut self, data: T) {
        self.user_data = Some(Box::new(data));
    }

    /// Get a reference to the stored user data.
    /// Returns None if no data is stored or if the type doesn't match.
    pub fn user_data<T: Send + 'static>(&self) -> Option<&T> {
        self.user_data.as_ref()?.downcast_ref()
    }

    /// Get a mutable reference to the stored user data.
    /// Returns None if no data is stored or if the type doesn't match.
    pub fn user_data_mut<T: Send + 'static>(&mut self) -> Option<&mut T> {
        self.user_data.as_mut()?.downcast_mut()
    }

    /// Clear the stored user data.
    pub fn clear_user_data(&mut self) {
        self.user_data = None;
    }

    // ========================================================================
    // Callbacks
    // ========================================================================

    /// Get a mutable reference to the host callbacks.
    ///
    /// Use this to retrieve collected print output or other callback state.
    pub fn callbacks_mut(&mut self) -> &mut dyn HostCallbacks {
        self.callbacks.as_mut()
    }

    /// Replace the host callbacks with new ones, returning the old callbacks.
    pub fn replace_callbacks(
        &mut self,
        callbacks: Box<dyn HostCallbacks + Send>,
    ) -> Box<dyn HostCallbacks + Send> {
        std::mem::replace(&mut self.callbacks, callbacks)
    }

    // ========================================================================
    // Memory tracking
    // ========================================================================

    /// Returns the number of GC-managed objects (tables and closures).
    pub fn object_count(&self) -> usize {
        self.heap.object_count()
    }

    /// Returns the number of interned strings.
    pub fn string_count(&self) -> usize {
        self.heap.string_count()
    }

    /// Returns the total number of heap allocations (objects + strings).
    pub fn heap_size(&self) -> usize {
        self.heap.allocation_count()
    }

    // ========================================================================
    // Host-controlled GC
    // ========================================================================

    /// Returns true if `heap_size()` has reached the GC threshold.
    /// The threshold counts objects plus distinct interned strings and is
    /// checked before allocation.
    pub fn gc_should_run(&self) -> bool {
        self.heap.is_full()
    }

    /// Returns the current GC threshold in total heap allocations (objects plus
    /// distinct interned strings). Each collection recomputes it from survivors.
    pub fn gc_threshold(&self) -> usize {
        self.heap.threshold()
    }

    /// Sets the GC threshold in total heap allocations (objects plus distinct
    /// interned strings). Collection triggers before allocation when
    /// `heap_size() >= threshold`; each collection recomputes the threshold
    /// from survivors.
    /// Set to `usize::MAX` to disable automatic GC: that value is a sentinel
    /// that collections preserve instead of recomputing.
    pub fn gc_set_threshold(&mut self, threshold: usize) {
        self.heap.set_threshold(threshold);
    }

    /// Disables automatic GC by setting threshold to usize::MAX.
    /// After calling this, GC only runs when you explicitly call `gc_collect()`;
    /// explicit collections keep automatic GC disabled. Re-enable it with
    /// `gc_set_threshold`.
    pub fn gc_disable_auto(&mut self) {
        self.heap.set_threshold(usize::MAX);
    }

    /// Forces a full garbage collection cycle.
    /// This marks all reachable objects and frees unreachable ones.
    #[hotpath::measure]
    pub fn gc_collect(&mut self) {
        self.gmatch_patterns.clear();
        self.gmatch_pattern_bytes = 0;
        // Mark all roots
        let mut worklist = self.heap.take_mark_worklist();
        mark_gc_roots(self, &mut worklist);
        self.heap
            .drain_mark_worklist(&mut worklist, &self.upvalue_pool);
        self.heap.restore_mark_worklist(worklist);
        self.sweep_bytecode_caches();
        // Sweep unmarked objects
        self.heap.collect();
    }

    /// Drop state-local runtime entries which no reachable closure, active
    /// frame, or pending closure shell can use. Removing an entry after roots
    /// were marked intentionally leaves its literals alive for one extra GC.
    fn sweep_bytecode_caches(&mut self) {
        let mut live = BTreeSet::new();
        for bytecode in self.heap.reachable_lua_bytecodes() {
            Self::retain_bytecode_tree(&mut live, &bytecode);
        }
        for call in &self.call_stack {
            Self::retain_bytecode_tree(&mut live, &call.bytecode);
        }
        live.extend(self.pending_bytecode_caches.iter().copied());
        self.bytecode_caches
            .retain(|identity, _| live.contains(identity));
    }

    fn retain_bytecode_tree(live: &mut BTreeSet<usize>, bytecode: &Arc<Bytecode>) {
        if !live.insert(Arc::as_ptr(bytecode) as usize) {
            return;
        }
        for nested in &bytecode.nested {
            Self::retain_bytecode_tree(live, nested);
        }
    }

    /// Resolve the State-local runtime bundle. Prevalidation happens before
    /// mutating the string pool; each completed intern is transiently rooted so
    /// an allocation-triggered GC cannot collect an earlier literal.
    pub(super) fn resolve_bytecode_runtime(
        &mut self,
        bytecode: &Arc<Bytecode>,
    ) -> Result<Arc<BytecodeRuntime>> {
        let identity = Arc::as_ptr(bytecode) as usize;
        if let Some(entry) = self.bytecode_caches.get(&identity) {
            return Ok(Arc::clone(&entry.runtime));
        }
        for literal in &bytecode.string_literals {
            check_string_size(literal.len())?;
        }
        let watermark = self.transient_roots.values.len();
        let result = (|| {
            let mut literals = Vec::with_capacity(bytecode.string_literals.len());
            for literal in &bytecode.string_literals {
                let value = self.alloc_string(literal)?;
                self.transient_roots.values.push(value);
                literals.push(value);
            }
            let runtime = Arc::new(BytecodeRuntime {
                literals: literals.into_boxed_slice(),
                caches: RuntimeCaches::new(bytecode),
            });
            self.bytecode_caches.insert(
                identity,
                BytecodeCacheEntry {
                    bytecode: Arc::clone(bytecode),
                    runtime: Arc::clone(&runtime),
                },
            );
            Ok(runtime)
        })();
        self.transient_roots.values.truncate(watermark);
        result
    }

    #[cfg(feature = "snapshot")]
    pub(super) fn resolve_bytecode_runtime_no_gc(
        &mut self,
        bytecode: &Arc<Bytecode>,
    ) -> Arc<BytecodeRuntime> {
        let identity = Arc::as_ptr(bytecode) as usize;
        if let Some(entry) = self.bytecode_caches.get(&identity) {
            return Arc::clone(&entry.runtime);
        }
        let literals = bytecode
            .string_literals
            .iter()
            .map(|literal| Val::Str(self.heap.alloc_string(literal)))
            .collect::<Vec<_>>()
            .into_boxed_slice();
        let runtime = Arc::new(BytecodeRuntime {
            literals,
            caches: RuntimeCaches::new(bytecode),
        });
        self.bytecode_caches.insert(
            identity,
            BytecodeCacheEntry {
                bytecode: Arc::clone(bytecode),
                runtime: Arc::clone(&runtime),
            },
        );
        runtime
    }

    // ========================================================================
    // Host callbacks
    // ========================================================================

    /// Called by the built-in `print()` function.
    /// Routes output through host callbacks with source context.
    pub(crate) fn host_print(&mut self, message: &str) {
        // Get current line from call stack (if available)
        let line = self
            .call_stack
            .last()
            .and_then(|info| {
                info.bytecode
                    .line_info
                    .get(info.ip.saturating_sub(1))
                    .copied()
            })
            .unwrap_or(0);

        let source = self.current_source.as_deref();
        self.callbacks.on_print(source, line, message);
    }

    /// Called when an error occurs. Notifies host callbacks.
    pub(crate) fn host_error(&mut self, error: &Error) {
        let source = self.current_source.as_deref();
        self.callbacks.on_error(source, error);
    }

    /// Returns the current source name (if set).
    pub fn current_source(&self) -> Option<&str> {
        self.current_source.as_deref()
    }

    /// Pushes onto the stack the value of the global `name`.
    #[hotpath::measure]
    pub fn get_global(&mut self, name: &str) -> Result<()> {
        // Check builtins first for common names
        let val = if let Some(slot) = Builtin::from_name(name) {
            self.builtins[slot as usize]
        } else {
            self.globals.get(name).copied().unwrap_or_default()
        };
        self.push_val(val)
    }

    /// Instr::pop()s a value from the stack and sets it as the new value of global
    /// `name`.
    #[hotpath::measure]
    pub fn set_global(&mut self, name: &str) {
        let val = self.pop_val();
        self.set_global_value(name, val);
    }

    #[cfg(not(feature = "snapshot"))]
    pub(crate) fn set_global_rust_fn(&mut self, name: &str, func: RustFunc) {
        self.set_global_value(name, Val::RustFn(func));
    }

    /// Register a Rust function under a stable id and install it as a global.
    ///
    /// Save/load uses `id` to resolve reachable [`RustFunc`] values across
    /// processes; ids should be stable strings owned by the embedder.
    #[cfg(feature = "snapshot")]
    pub fn set_global_named_rust_fn(
        &mut self,
        name: &str,
        id: &str,
        func: RustFunc,
    ) -> std::result::Result<(), SaveError> {
        self.register_rust_fn(id, func)?;
        self.set_global_value(name, Val::RustFn(func));
        Ok(())
    }

    #[cfg(feature = "snapshot")]
    pub(crate) fn set_global_stdlib_rust_fn(&mut self, name: &str, id: &str, func: RustFunc) {
        self.set_global_named_rust_fn(name, id, func)
            .expect("stdlib Rust function registration cannot fail");
    }

    /// Register a Rust function id without installing the function anywhere.
    ///
    /// This is useful when the host pushes or stores the function manually but
    /// still wants it to survive save/load if Lua makes it reachable.
    #[cfg(feature = "snapshot")]
    pub fn register_rust_fn(
        &mut self,
        id: &str,
        func: RustFunc,
    ) -> std::result::Result<(), save_state::SaveError> {
        let addr = func as usize;
        // The only genuine error is an id collision: one id bound to two
        // different function addresses (an embedder bug). Re-registering the
        // same id->fn pair is idempotent.
        if let Some(existing_func) = self.rust_fns_by_id.get(id) {
            if *existing_func as usize != addr {
                return Err(save_state::SaveError::DuplicateFunctionRegistration {
                    id: id.to_string(),
                });
            }
            return Ok(());
        }
        self.rust_fns_by_id.insert(id.to_string(), func);
        // Several ids may legitimately share one address: intentional aliasing,
        // or identical-code folding collapsing distinct fns under release LTO.
        // Both are behaviorally safe (same code), so we never error here; we
        // keep the lexicographically-smallest id as the reverse name so saves
        // stay deterministic regardless of registration order.
        match self.rust_fn_ids_by_addr.get(&addr) {
            Some(existing_id) if existing_id.as_str() <= id => {}
            _ => {
                self.rust_fn_ids_by_addr.insert(addr, id.to_string());
            }
        }
        Ok(())
    }

    pub(super) fn set_global_value(&mut self, name: &str, val: Val) {
        self.set_global_value_owned(name.to_string(), val);
    }

    pub(super) fn set_global_value_owned(&mut self, name: String, val: Val) {
        // Update builtins array if this is a well-known name. Rebinding a
        // builtin slot poisons inline caches that hold a direct ObjectPtr
        // to the previous library table (string-method IC, method-lookup
        // IC reaching the lib via __index), so bump globals_version to
        // force re-resolution. User globals don't poison those ICs - the
        // ICs only key on builtin slots - so the bump stays narrow.
        if let Some(slot) = Builtin::from_name(&name) {
            self.builtins[slot as usize] = val;
            self.globals_version = self.globals_version.wrapping_add(1);
            if slot == Builtin::Table {
                self.invalidate_table_library_fallback_rebind(val);
            }
        }
        self.globals.insert(name, val);
    }

    pub(super) fn capture_table_library_fallback(&mut self) {
        let Some(table) = self.builtins[Builtin::Table as usize].as_object_ptr() else {
            self.table_library_fallback = None;
            return;
        };
        let Some(library) = self.heap.as_table_ref(table) else {
            self.table_library_fallback = None;
            return;
        };
        if library.get_metatable().is_some() {
            self.table_library_fallback = None;
            return;
        }
        let Ok(names) = library.live_string_keys().try_into() else {
            self.table_library_fallback = None;
            return;
        };
        self.table_library_fallback = Some(TableLibraryFallbackCache {
            table,
            names,
            shape: library.fallback_shape(),
        });
    }

    pub(super) fn invalidate_table_library_fallback_rebind(&mut self, val: Val) {
        if self
            .table_library_fallback
            .as_ref()
            .is_some_and(|cache| val.as_object_ptr() != Some(cache.table))
        {
            self.table_library_fallback = None;
        }
    }

    /// Only the snapshot load path needs an unconditional drop: an ordered
    /// `"table"` environment replay rebuilds the canonical object in place,
    /// which could coincidentally restore pristine-looking shape fields.
    #[cfg(feature = "snapshot")]
    pub(super) fn drop_table_library_fallback(&mut self) {
        self.table_library_fallback = None;
    }

    /// Execute a function with a restricted global environment.
    /// Only globals in the whitelist are accessible during execution.
    /// The original environment is restored after the function completes (or errors).
    pub fn with_restricted_env<F, R>(&mut self, whitelist: &[&str], f: F) -> R
    where
        F: FnOnce(&mut Self) -> R,
    {
        // Build restricted environment
        let mut restricted_globals = IndexMap::new();
        let mut restricted_builtins: [Val; Builtin::COUNT] = std::array::from_fn(|_| Val::Nil);

        for name in whitelist {
            // Copy from builtins if it's a well-known name
            if let Some(slot) = Builtin::from_name(name) {
                restricted_builtins[slot as usize] = self.builtins[slot as usize];
            }
            // Also copy from globals
            if let Some(val) = self.globals.get(*name) {
                restricted_globals.insert((*name).to_string(), *val);
            }
        }

        // Swap to restricted environment
        let saved_globals = std::mem::replace(&mut self.globals, restricted_globals);
        let saved_builtins = std::mem::replace(&mut self.builtins, restricted_builtins);
        self.transient_roots
            .suspended_envs
            .push(SuspendedEnvironment {
                globals: saved_globals,
                builtins: saved_builtins,
            });
        self.globals_version = self.globals_version.wrapping_add(1);

        // Execute the function under an unwind guard so a panic in `f` still
        // restores the original environment before propagating (L11). Without
        // this, an embedder that catches the panic and reuses the State would
        // find it stuck in the restricted environment. (Under panic=abort
        // catch_unwind never returns, but reuse is moot there.)
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| f(self)));

        // Restore original environment (on both normal return and unwind)
        let saved = self
            .transient_roots
            .suspended_envs
            .pop()
            .expect("suspended environment missing during restoration");
        self.globals = saved.globals;
        self.builtins = saved.builtins;
        self.globals_version = self.globals_version.wrapping_add(1);

        match result {
            Ok(r) => r,
            Err(payload) => std::panic::resume_unwind(payload),
        }
    }

    pub(super) fn with_rooted_value<F, R>(&mut self, value: Val, f: F) -> R
    where
        F: FnOnce(&mut Self) -> R,
    {
        self.with_rooted_values(&[value], f)
    }

    pub(super) fn with_rooted_values<F, R>(&mut self, values: &[Val], f: F) -> R
    where
        F: FnOnce(&mut Self) -> R,
    {
        let watermark = self.transient_roots.values.len();
        self.transient_roots.values.extend_from_slice(values);
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| f(self)));
        self.transient_roots.values.truncate(watermark);
        match result {
            Ok(result) => result,
            Err(payload) => std::panic::resume_unwind(payload),
        }
    }

    /// Allocates a string on the heap.
    #[hotpath::measure]
    pub(super) fn alloc_string(&mut self, bytes: impl AsRef<[u8]>) -> Result<Val> {
        check_string_size(bytes.as_ref().len())?;
        // Check if GC is needed before allocating
        if self.heap.is_full() {
            self.gc_collect();
        }
        let ptr = self.heap.alloc_string(bytes.as_ref());
        Ok(Val::Str(ptr))
    }

    /// Construct an [`Error`] of the given kind with no source position.
    ///
    /// The position is filled in later, on the way out of the Lua frame the
    /// error surfaced in (see `locate_in_frame`), because that is the first
    /// point with access to the frame's line information. Errors that never
    /// unwind through a Lua frame - those raised by a host `RustFunc` called
    /// directly - keep a zero position.
    pub fn error(&self, kind: ErrorKind) -> Error {
        Error::new(kind, 0, 0)
    }

    pub(super) fn type_error(&self, e: TypeError) -> Error {
        self.error(ErrorKind::TypeError(e))
    }

    /// Build a stack trace from the current call stack and the active frame.
    /// The frame represents the innermost (current) function where the error occurred.
    #[allow(private_interfaces)]
    pub(super) fn build_stack_trace(&self, current_frame: &frame::Frame) -> Vec<StackFrame> {
        let mut trace = Vec::with_capacity(self.call_stack.len() + 1);

        // First entry: the current frame where the error occurred
        trace.push(current_frame.to_stack_frame());

        // Add entries from call_stack (most recent first).
        // Skip the last entry (current function) since we already have it from current_frame.
        // The remaining entries are the callers, with ip pointing to their call sites.
        for call_info in self.call_stack.iter().rev().skip(1) {
            // Get line number from the ip (call site)
            let line = if call_info.ip > 0 {
                call_info
                    .bytecode
                    .line_info
                    .get(call_info.ip - 1)
                    .copied()
                    .unwrap_or(0)
            } else {
                call_info.bytecode.line_info.first().copied().unwrap_or(0)
            };
            trace.push(StackFrame {
                function_name: call_info.bytecode.name.clone(),
                source: call_info.bytecode.source.clone(),
                line,
            });
        }

        trace
    }
}

impl Default for State {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod runtime_tests {
    use super::*;

    #[test]
    fn closures_of_one_bytecode_share_runtime_bundle() {
        let mut state = State::empty();
        let bytecode = Arc::new(Bytecode {
            string_literals: vec![b"field".to_vec()],
            ..Bytecode::default()
        });
        state
            .push_closure(Arc::clone(&bytecode), Vec::new())
            .expect("first closure allocates");
        state
            .push_closure(Arc::clone(&bytecode), Vec::new())
            .expect("second closure allocates");
        let second = state
            .pop_val()
            .as_lua_function(&state.heap)
            .expect("second value is a closure");
        let first = state
            .pop_val()
            .as_lua_function(&state.heap)
            .expect("first value is a closure");
        assert!(Arc::ptr_eq(&first.runtime, &second.runtime));
        assert_eq!(state.bytecode_caches.len(), 1);
    }

    #[test]
    fn warming_one_factory_closure_populates_the_shared_slots() {
        let mut state = State::new();
        state
            .load_string(
                "function make() return function(t) return t.v end end \
                 r1 = make() r2 = make() warmed = r1({ v = 3 })",
            )
            .expect("factory source compiles");
        state
            .call(ArgCount::Fixed(0), RetCount::Fixed(0))
            .expect("factory source runs");
        state.get_global("r1").expect("r1 exists");
        let r1 = state
            .pop_val()
            .as_lua_function(&state.heap)
            .expect("r1 is a closure");
        state.get_global("r2").expect("r2 exists");
        let r2 = state
            .pop_val()
            .as_lua_function(&state.heap)
            .expect("r2 is a closure");
        assert!(Arc::ptr_eq(&r1.runtime, &r2.runtime));
        // r1 executed once; that warmup must be visible through r2's handle -
        // an actually-populated slot, not just pointer equality.
        assert!(
            r2.runtime
                .caches
                .field_lookup
                .iter()
                .any(|slot| slot.get_field().is_some()),
            "executing one closure must warm the shared field slot"
        );
    }

    #[test]
    fn runtime_cache_sweep_is_state_reachability_based() {
        let mut state = State::empty();
        state.gc_disable_auto();
        let bytecode = Arc::new(Bytecode {
            string_literals: vec![b"temporary".to_vec()],
            ..Bytecode::default()
        });
        state
            .push_closure(Arc::clone(&bytecode), Vec::new())
            .expect("closure allocates");
        state.pop_val();
        state.gc_collect();
        assert!(state.bytecode_caches.is_empty());
        // Removed entries were roots for the collection that removed them.
        assert_eq!(state.string_count(), 1);
        state.gc_collect();
        assert_eq!(state.string_count(), 0);
    }

    #[test]
    fn live_parent_retains_nested_runtime_entry() {
        let mut state = State::empty();
        state.gc_disable_auto();
        let child = Arc::new(Bytecode::default());
        let parent = Arc::new(Bytecode {
            nested: vec![Arc::clone(&child)],
            ..Bytecode::default()
        });
        state
            .push_closure(Arc::clone(&child), Vec::new())
            .expect("child closure allocates");
        state.pop_val();
        state
            .push_closure(Arc::clone(&parent), Vec::new())
            .expect("parent closure allocates");
        state.gc_collect();
        assert_eq!(state.bytecode_caches.len(), 2);
        state.pop_val();
        state.gc_collect();
        assert!(state.bytecode_caches.is_empty());
    }
}

#[cfg(test)]
mod tests;