harn_vm/value/

env.rs

use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::{Arc, Weak};

use crate::chunk::CompiledFunctionRef;

use super::{VmError, VmMutex, VmValue};

/// A compiled closure value.
#[derive(Debug, Clone)]
pub struct VmClosure {
    pub func: CompiledFunctionRef,
    pub env: VmEnv,
    /// Source directory for this closure's originating module.
    /// When set, `render()` and other source-relative builtins resolve
    /// paths relative to this directory instead of the entry pipeline.
    pub source_dir: Option<PathBuf>,
    /// Module-local named functions that should resolve before builtin fallback.
    /// This lets selectively imported functions keep private sibling helpers
    /// without exporting them into the caller's environment.
    pub module_functions: Option<WeakModuleFunctionRegistry>,
    /// Shared, mutable module-level env: holds top-level `var` / `let`
    /// bindings declared at the module root (caches, counters, lazily
    /// initialized registries). All closures created from the same
    /// module import point at the same shared mutable env, so a
    /// mutation inside one function is visible to every other function
    /// in that module on subsequent calls. `closure.env` still holds
    /// the per-closure lexical snapshot (captured function args from
    /// enclosing scopes, etc.) and is unchanged by this — `module_state`
    /// is a separate lookup layer consulted after the local env and
    /// before globals. Created in `import_declarations` after the
    /// module's init chunk runs, so the initial values from `var x = ...`
    /// land in it.
    pub module_state: Option<WeakModuleState>,
    /// Strong owners of this closure's module scope, pinned only when the
    /// closure is stored in a process/thread-local registry that outlives the
    /// VM that created it (reminder providers, session/lifecycle hooks). See
    /// [`RetainedModuleScope`] and [`VmClosure::retained_for_host_registry`].
    /// `None` for the overwhelmingly common short-lived closure, whose module
    /// scope stays alive through the live VM's `module_cache`.
    pub retained_module_scope: Option<Arc<RetainedModuleScope>>,
}

pub type ModuleFunctionRegistry = Arc<VmMutex<BTreeMap<String, Arc<VmClosure>>>>;
pub type WeakModuleFunctionRegistry = Weak<VmMutex<BTreeMap<String, Arc<VmClosure>>>>;
pub type ModuleState = Arc<VmMutex<VmEnv>>;
pub type WeakModuleState = Weak<VmMutex<VmEnv>>;

/// Strong owners of a closure's module function table and module-level state.
///
/// A [`VmClosure`] resolves sibling module `pub fn`s through its module's
/// function registry, which it references only via a [`Weak`]
/// ([`VmClosure::module_functions`] / [`module_state`](VmClosure::module_state)).
/// The sole strong owner of that registry is normally the registering VM's
/// `module_cache`. When a closure is registered into a process/thread-local
/// registry (reminder providers, session/lifecycle hooks) it outlives that VM;
/// once the VM tears down, the `Weak` dangles and a sibling-fn call inside the
/// invoked closure falls through name resolution to host-bridge dispatch. This
/// pins strong owners so the `Weak` stays upgradeable for the closure's whole
/// retained lifetime.
///
/// The fields are intentionally unread — their sole purpose is to keep the
/// referenced `Arc`s alive.
#[derive(Debug)]
pub struct RetainedModuleScope {
    _functions: Option<ModuleFunctionRegistry>,
    _state: Option<ModuleState>,
}

impl VmClosure {
    pub(crate) fn module_functions(&self) -> Option<ModuleFunctionRegistry> {
        self.module_functions
            .as_ref()
            .and_then(WeakModuleFunctionRegistry::upgrade)
    }

    pub(crate) fn module_state(&self) -> Option<ModuleState> {
        self.module_state
            .as_ref()
            .and_then(WeakModuleState::upgrade)
    }

    /// Return a clone of this closure suitable for storage in a process- or
    /// thread-local registry that outlives the VM that created it (reminder
    /// providers, session/lifecycle hooks). The clone pins strong owners of
    /// this closure's module function table and module-level state
    /// ([`RetainedModuleScope`]), so its body still resolves sibling module
    /// `pub fn`s after the registering VM — the only other strong owner, via
    /// `module_cache` — is dropped.
    ///
    /// The owners are pinned on a *clone* (a fresh `Arc<VmClosure>` that is
    /// never itself a member of any function registry), so retaining a closure
    /// that IS a module `pub fn` cannot form an `Arc` cycle with its registry.
    ///
    /// A no-op refcount bump when there is nothing to pin: the closure is
    /// already pinned, or its `Weak`s do not upgrade — e.g. an entry-chunk
    /// closure whose sibling functions live in captured `env` rather than a
    /// module registry, which resolves without this.
    pub(crate) fn retained_for_host_registry(self: &Arc<Self>) -> Arc<Self> {
        if self.retained_module_scope.is_some() {
            return Arc::clone(self);
        }
        let functions = self.module_functions();
        let state = self.module_state();
        if functions.is_none() && state.is_none() {
            return Arc::clone(self);
        }
        let mut pinned = (**self).clone();
        pinned.retained_module_scope = Some(Arc::new(RetainedModuleScope {
            _functions: functions,
            _state: state,
        }));
        Arc::new(pinned)
    }
}

/// VM environment for variable storage.
///
/// `Scope::vars` is wrapped in `Arc` so that `VmEnv::clone()` is cheap
/// (Arc bump per scope) instead of a deep walk of every BTreeMap. The
/// VM saves and restores `env` snapshots on every function call, and
/// the call hot path dominates orchestration-heavy workloads. With
/// `Arc<BTreeMap<..>>`, the per-scope clone collapses to a refcount
/// bump, and `Arc::make_mut` only does a deep copy when the scope is
/// still shared with a saved snapshot — which is exactly the case where
/// the caller would have needed an isolated copy anyway. Reads still go
/// through the `BTreeMap` directly via `Deref`.
#[derive(Debug, Clone)]
pub struct VmEnv {
    pub(crate) scopes: Vec<Scope>,
}

/// A shared, mutable cell backing a captured binding.
///
/// A local that a nested closure captures is stored behind a `Cell` instead of
/// inline. Cloning a [`Scope`] (which happens on every call and every closure
/// mint) refcount-bumps this `Arc`, so the defining frame and every closure
/// that captured the binding all point at the *same* cell — a write through any
/// of them is observed by all of them. This is what makes closure capture
/// **by reference** (JS/Python/Swift semantics) while keeping distinct
/// variables independent (`let b = a` still copies the value out of `a`'s
/// cell into `b`'s binding). See `docs/design/closure-reference-capture.md`.
pub(crate) type BindingCell = Arc<VmMutex<VmValue>>;

/// One name's binding in a [`Scope`].
///
/// `Value` is the ordinary, unshared binding — a read clones the value out and
/// a write replaces it (copy-on-assignment), exactly as before. `Cell` is a
/// binding captured by a nested closure: the value lives behind a shared
/// [`BindingCell`] so reads clone the inner value out (value semantics for
/// reads is preserved) and writes go *through* the cell rather than replacing
/// the map entry — which also sidesteps the scope-map copy-on-write, so shared
/// mutation survives the per-call env clone.
#[derive(Debug, Clone)]
pub(crate) enum Binding {
    Value { value: VmValue, mutable: bool },
    Cell { cell: BindingCell, mutable: bool },
}

impl Binding {
    #[inline]
    pub(crate) fn mutable(&self) -> bool {
        match self {
            Binding::Value { mutable, .. } | Binding::Cell { mutable, .. } => *mutable,
        }
    }

    /// The current value of this binding, cloned out. Reads never expose the
    /// cell itself — value semantics for reads is identical for both variants.
    #[inline]
    pub(crate) fn read(&self) -> VmValue {
        match self {
            Binding::Value { value, .. } => value.clone(),
            Binding::Cell { cell, .. } => cell.lock().clone(),
        }
    }

    /// Ownership-taking accessor for the iterative teardown paths. A `Value`
    /// yields its inner value directly. A `Cell` yields its inner value only
    /// when this binding holds the *last* reference to the shared cell; a
    /// still-shared cell yields `None` and is left for its own `Arc` drop to
    /// reclaim once the final closure releases it.
    #[inline]
    pub(crate) fn into_teardown_value(self) -> Option<VmValue> {
        match self {
            Binding::Value { value, .. } => Some(value),
            Binding::Cell { cell, .. } => Arc::into_inner(cell).map(VmMutex::into_inner),
        }
    }

    /// Whether this binding *uniquely* owns a deeply-nested container that the
    /// default recursive drop could overflow the native stack on. A `Value`
    /// checks its container directly. A `Cell` only qualifies when unshared
    /// (`strong_count == 1`) — a cell still held by a live closure must not be
    /// torn down from here — and is peeked with `try_lock` so a drop never
    /// blocks.
    #[inline]
    fn owns_recursive_container(&self) -> bool {
        match self {
            Binding::Value { value, .. } => super::recursion::is_recursive_container(value),
            Binding::Cell { cell, .. } => {
                Arc::strong_count(cell) == 1
                    && cell
                        .try_lock()
                        .map(|v| super::recursion::is_recursive_container(&v))
                        .unwrap_or(false)
            }
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct Scope {
    pub(crate) vars: Arc<BTreeMap<String, Binding>>,
}

/// Process-wide shared empty binding map.
///
/// Every block entry pushes a fresh [`Scope`], but inside a function body its
/// bindings compile to local slots (`DefLocalSlot`) rather than env writes, so
/// the pushed scope is overwhelmingly *empty* — a hot loop whose body is a
/// block would otherwise `Arc::new(BTreeMap::new())`-allocate (and free) one
/// map per iteration. Sharing a single immutable empty map makes
/// [`Scope::empty`] a refcount bump instead; the first real `define`/`assign`
/// copies-on-write away from this shared map via `Arc::make_mut` (the insert
/// paths already do), so a scope that never binds anything never allocates.
static EMPTY_SCOPE_VARS: std::sync::LazyLock<Arc<BTreeMap<String, Binding>>> =
    std::sync::LazyLock::new(|| Arc::new(BTreeMap::new()));

impl Scope {
    #[inline]
    fn empty() -> Self {
        Self {
            vars: Arc::clone(&EMPTY_SCOPE_VARS),
        }
    }
}

impl Drop for Scope {
    fn drop(&mut self) {
        // Deeply nested script values (e.g. `x = [x]` built in a loop, which
        // adds no VM call frames and so never trips `max_vm_frames`) live in
        // scope bindings. Their default recursive drop would overflow the
        // native stack and abort the whole process — an uncatchable failure.
        // When this scope holds the last reference to its bindings and any
        // value is a nested container, tear the bindings down iteratively
        // instead. `Arc::get_mut` succeeds only for a uniquely-owned scope, so
        // shared snapshots fall through to the cheap default drop and the real
        // teardown happens later at the last owner (also a `Scope`).
        //
        // A still-shared `Cell` may outlive this scope (a live closure holds
        // it), so its `Arc` refcount — not this map's — governs when its inner
        // value drops. `into_teardown_value` therefore only reclaims a cell we
        // uniquely own; shared cells fall through to their own `Arc` drop.
        if let Some(map) = Arc::get_mut(&mut self.vars) {
            if map.values().any(Binding::owns_recursive_container) {
                let bindings = std::mem::take(map);
                super::recursion::dismantle_values(
                    bindings
                        .into_values()
                        .filter_map(Binding::into_teardown_value),
                );
            }
        }
    }
}

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

impl VmEnv {
    pub fn new() -> Self {
        Self {
            scopes: vec![Scope::empty()],
        }
    }

    pub fn push_scope(&mut self) {
        self.scopes.push(Scope::empty());
    }

    /// Clone the scope stack for a fresh call frame, reserving room for the
    /// one empty scope every invocation pushes for the callee's body.
    ///
    /// `Vec::clone` allocates at exactly `len` capacity, so the `push_scope`
    /// that immediately follows on the call hot path would otherwise force a
    /// reallocation and copy of the whole scope stack. Reserving the extra
    /// slot up front folds those two allocations into one. When a caller does
    /// not end up pushing (no path currently does, but it stays correct if one
    /// is added), the only cost is a single unused `Scope` slot of capacity.
    pub(crate) fn cloned_for_call(&self) -> VmEnv {
        let mut scopes = Vec::with_capacity(self.scopes.len() + 1);
        scopes.extend(self.scopes.iter().cloned());
        VmEnv { scopes }
    }

    pub fn pop_scope(&mut self) {
        if self.scopes.len() > 1 {
            self.scopes.pop();
        }
    }

    pub fn scope_depth(&self) -> usize {
        self.scopes.len()
    }

    pub fn truncate_scopes(&mut self, target_depth: usize) {
        let min_depth = target_depth.max(1);
        while self.scopes.len() > min_depth {
            self.scopes.pop();
        }
    }

    pub fn get(&self, name: &str) -> Option<VmValue> {
        for scope in self.scopes.iter().rev() {
            if let Some(binding) = scope.vars.get(name) {
                return Some(binding.read());
            }
        }
        None
    }

    pub(crate) fn contains(&self, name: &str) -> bool {
        self.scopes
            .iter()
            .rev()
            .any(|scope| scope.vars.contains_key(name))
    }

    pub fn define(&mut self, name: &str, value: VmValue, mutable: bool) -> Result<(), VmError> {
        self.define_binding(name, Binding::Value { value, mutable })
    }

    /// Define `name` as a **captured** binding: a fresh shared cell holding
    /// `value`. Emitted for a local that a nested closure captures. A closure
    /// minted after this point clones the enclosing env (refcount-bumping the
    /// cell), so its reads and writes of `name` flow through the same cell as
    /// the defining frame. Called once per activation, so each activation gets
    /// a distinct cell (per-iteration loop captures stay independent).
    pub(crate) fn define_cell(
        &mut self,
        name: &str,
        value: VmValue,
        mutable: bool,
    ) -> Result<(), VmError> {
        self.define_binding(
            name,
            Binding::Cell {
                cell: Arc::new(VmMutex::new(value)),
                mutable,
            },
        )
    }

    fn define_binding(&mut self, name: &str, binding: Binding) -> Result<(), VmError> {
        if let Some(scope) = self.scopes.last_mut() {
            if let Some(existing) = scope.vars.get(name) {
                if !existing.mutable() && !binding.mutable() {
                    return Err(VmError::Runtime(format!(
                        "Cannot redeclare immutable variable '{name}' in the same scope (use 'let' for mutable bindings)"
                    )));
                }
            }
            if let Some(Binding::Value { value, .. }) =
                Arc::make_mut(&mut scope.vars).insert(name.to_string(), binding)
            {
                super::recursion::dismantle(value);
            }
        }
        Ok(())
    }

    pub fn all_variables(&self) -> crate::value::DictMap {
        let mut vars = crate::value::DictMap::new();
        for scope in &self.scopes {
            for (name, binding) in scope.vars.iter() {
                vars.insert(crate::value::intern_key(name), binding.read());
            }
        }
        vars
    }

    pub fn assign(&mut self, name: &str, value: VmValue) -> Result<(), VmError> {
        for scope in self.scopes.iter_mut().rev() {
            let Some(existing) = scope.vars.get(name) else {
                continue;
            };
            if !existing.mutable() {
                return Err(VmError::ImmutableAssignment(name.to_string()));
            }
            match existing {
                // Write *through* the shared cell: the entry is not replaced,
                // so the scope-map copy-on-write is sidestepped and every
                // holder of this cell (the defining frame, sibling closures)
                // observes the update.
                Binding::Cell { cell, .. } => {
                    let previous = std::mem::replace(&mut *cell.lock(), value);
                    super::recursion::dismantle(previous);
                }
                Binding::Value { .. } => {
                    // Iterative teardown so overwriting a deeply nested binding
                    // cannot overflow the stack on drop (scalars are a no-op).
                    // The prior binding here is always a `Value` (a name is
                    // either always boxed or never — see the compiler's capture
                    // pre-pass), so only that arm needs draining.
                    if let Some(Binding::Value { value, .. }) = Arc::make_mut(&mut scope.vars)
                        .insert(
                            name.to_string(),
                            Binding::Value {
                                value,
                                mutable: true,
                            },
                        )
                    {
                        super::recursion::dismantle(value);
                    }
                }
            }
            return Ok(());
        }
        Err(VmError::UndefinedVariable(name.to_string()))
    }

    /// Debugger-only variant of `assign` that rebinds the name even if
    /// the existing binding was declared with `let`. Pipeline authors
    /// overwhelmingly use `let`, so a strict mutability check would
    /// make the DAP `setVariable` request useless for "what-if"
    /// iteration — which is the whole point of the feature. Preserves
    /// the original mutability flag so the VM's runtime behavior is
    /// unchanged after the debugger overrides.
    pub fn assign_debug(&mut self, name: &str, value: VmValue) -> Result<(), VmError> {
        for scope in self.scopes.iter_mut().rev() {
            let Some(existing) = scope.vars.get(name) else {
                continue;
            };
            match existing {
                // Preserve the shared-cell identity so a debugger override of a
                // captured binding is still observed by the closures holding it.
                Binding::Cell { cell, .. } => {
                    *cell.lock() = value;
                }
                Binding::Value { mutable, .. } => {
                    let mutable = *mutable;
                    Arc::make_mut(&mut scope.vars)
                        .insert(name.to_string(), Binding::Value { value, mutable });
                }
            }
            return Ok(());
        }
        Err(VmError::UndefinedVariable(name.to_string()))
    }
}

/// Compute Levenshtein edit distance between two strings.
fn levenshtein(a: &str, b: &str) -> usize {
    let a: Vec<char> = a.chars().collect();
    let b: Vec<char> = b.chars().collect();
    let (m, n) = (a.len(), b.len());
    let mut prev = (0..=n).collect::<Vec<_>>();
    let mut curr = vec![0; n + 1];
    for i in 1..=m {
        curr[0] = i;
        for j in 1..=n {
            let cost = usize::from(a[i - 1] != b[j - 1]);
            curr[j] = (prev[j] + 1).min(curr[j - 1] + 1).min(prev[j - 1] + cost);
        }
        std::mem::swap(&mut prev, &mut curr);
    }
    prev[n]
}

/// Find the closest match from a list of candidates using Levenshtein distance.
/// Returns `Some(suggestion)` if a candidate is within `max_dist` edits.
pub fn closest_match<'a>(name: &str, candidates: impl Iterator<Item = &'a str>) -> Option<String> {
    let max_dist = match name.len() {
        0..=2 => 1,
        3..=5 => 2,
        _ => 3,
    };
    candidates
        .filter(|c| *c != name && !c.starts_with("__"))
        .map(|c| (c, levenshtein(name, c)))
        .filter(|(_, d)| *d <= max_dist)
        // Prefer smallest distance, then closest length to original, then alphabetical
        .min_by(|(a, da), (b, db)| {
            da.cmp(db)
                .then_with(|| {
                    let a_diff = (a.len() as isize - name.len() as isize).unsigned_abs();
                    let b_diff = (b.len() as isize - name.len() as isize).unsigned_abs();
                    a_diff.cmp(&b_diff)
                })
                .then_with(|| a.cmp(b))
        })
        .map(|(c, _)| c.to_string())
}

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

    #[test]
    fn empty_scopes_share_one_backing_map() {
        // Pushing block scopes (the per-iteration cost in a loop body) must not
        // allocate: every empty scope shares the process-wide empty map.
        let mut env = VmEnv::new();
        env.push_scope();
        env.push_scope();
        for scope in &env.scopes {
            assert!(Arc::ptr_eq(&scope.vars, &EMPTY_SCOPE_VARS));
        }
    }

    #[test]
    fn define_copies_on_write_without_disturbing_siblings() {
        let mut env = VmEnv::new();
        env.push_scope(); // shares EMPTY
        env.define("x", VmValue::Int(1), true).unwrap();
        // The bound scope copied on write away from the shared empty map...
        let top = env.scopes.last().unwrap();
        assert!(!Arc::ptr_eq(&top.vars, &EMPTY_SCOPE_VARS));
        // ...while the root scope (untouched) still shares it.
        assert!(Arc::ptr_eq(&env.scopes[0].vars, &EMPTY_SCOPE_VARS));
        assert!(matches!(env.get("x"), Some(VmValue::Int(1))));
        // Popping the scope drops the binding entirely.
        env.pop_scope();
        assert!(env.get("x").is_none());
    }
}