cljrs_runtime/env/

dynamics.rs

//! Thread-local dynamic variable binding stack.
//!
//! `binding` forms push a frame onto `BINDING_STACK` for the duration of their
//! body; the RAII `BindingGuard` pops it on drop (handles both normal return
//! and panics).

use std::cell::RefCell;
use std::collections::HashMap;

use cljrs_gc::GcPtr;
use cljrs_gc::Trace as _;
use cljrs_value::{Value, Var};

/// Opaque key for a Var in the binding stack (pointer identity).
/// Stable because the GC is non-moving.
pub type VarKey = usize;

pub fn var_key_of(var: &GcPtr<Var>) -> VarKey {
    var.get() as *const Var as usize
}

struct BindingFrame {
    values: HashMap<VarKey, Value>,
    #[cfg(feature = "no-gc")]
    leases: HashMap<VarKey, cljrs_value::publish::ValueLease>,
}
impl BindingFrame {
    fn new(values: HashMap<VarKey, Value>) -> Self {
        #[cfg(feature = "no-gc")]
        let (values, leases) = {
            let mut output = HashMap::new();
            let mut leases = HashMap::new();
            for (key, value) in values {
                let (value, lease) =
                    cljrs_value::publish::prepare_for_owner(value, &Default::default());
                output.insert(key, value);
                leases.insert(key, lease);
            }
            (output, leases)
        };
        Self {
            values,
            #[cfg(feature = "no-gc")]
            leases,
        }
    }
}

thread_local! {
    static BINDING_STACK: RefCell<Vec<BindingFrame>> =
        const { RefCell::new(Vec::new()) };
}
cljrs_gc::mutator_slot!("cljrs-runtime/BINDING_STACK", refcell BINDING_STACK: Vec<BindingFrame>, trace trace_saved);

// ── RAII guard ────────────────────────────────────────────────────────────────

/// Pops the innermost binding frame when dropped.
pub struct BindingGuard;

impl Drop for BindingGuard {
    fn drop(&mut self) {
        pop_frame();
    }
}

// ── Stack manipulation ────────────────────────────────────────────────────────

/// Push a new dynamic binding frame; return a guard that pops it on drop.
pub fn push_frame(bindings: HashMap<VarKey, Value>) -> BindingGuard {
    let frame = BindingFrame::new(bindings);
    BINDING_STACK.with(|s| s.borrow_mut().push(frame));
    BindingGuard
}

/// Pop the frame `guard` owns and hand it back, so a caller that yields between
/// uses (the async evaluator) can keep a `binding` frame off the thread-local
/// stack while other tasks on the same thread run, and push it again later.
/// Must be called with the innermost frame's guard.
pub fn take_frame(guard: BindingGuard) -> HashMap<VarKey, Value> {
    std::mem::forget(guard);
    BINDING_STACK.with(|s| {
        s.borrow_mut()
            .pop()
            .map(|frame| {
                #[cfg(feature = "no-gc")]
                for lease in frame.leases.values() {
                    lease.protect_read();
                }
                frame.values
            })
            .unwrap_or_default()
    })
}

fn pop_frame() {
    BINDING_STACK.with(|s| {
        s.borrow_mut().pop();
    });
}

// ── Lookup ────────────────────────────────────────────────────────────────────

/// Check the thread-local stack first (innermost frame wins); fall back to the
/// root binding stored in the `Var` itself.
pub fn deref_var(var: &GcPtr<Var>) -> Option<Value> {
    let key = var_key_of(var);
    let tl = BINDING_STACK.with(|s| {
        s.borrow().iter().rev().find_map(|frame| {
            let value = frame.values.get(&key)?.clone();
            #[cfg(feature = "no-gc")]
            if let Some(lease) = frame.leases.get(&key) {
                lease.protect_read();
            }
            Some(value)
        })
    });
    tl.or_else(|| var.get().deref())
}

/// True if `var` has any thread-local binding on this thread.
pub fn is_thread_bound(var: &GcPtr<Var>) -> bool {
    let key = var_key_of(var);
    BINDING_STACK.with(|s| {
        s.borrow()
            .iter()
            .any(|frame| frame.values.contains_key(&key))
    })
}

/// Set the innermost thread-local binding for `var`.
/// Returns `false` if no thread-local binding exists (caller should fall back
/// to setting the root).
pub fn set_thread_local(var: &GcPtr<Var>, val: Value) -> bool {
    let key = var_key_of(var);
    #[cfg(feature = "no-gc")]
    {
        let target = BINDING_STACK.with(|stack| {
            stack
                .borrow()
                .iter()
                .rposition(|frame| frame.values.contains_key(&key))
        });
        let Some(index) = target else {
            return false;
        };
        let (value, lease) = cljrs_value::publish::prepare_for_owner(val, &Default::default());
        let old = BINDING_STACK.with(|stack| {
            let mut stack = stack.borrow_mut();
            stack[index].values.insert(key, value);
            stack[index].leases.insert(key, lease)
        });
        drop(old);
        true
    }
    #[cfg(not(feature = "no-gc"))]
    BINDING_STACK.with(|s| {
        for frame in s.borrow_mut().iter_mut().rev() {
            if let std::collections::hash_map::Entry::Occupied(mut e) = frame.values.entry(key) {
                e.insert(val);
                return true;
            }
        }
        false
    })
}

// ── Binding conveyance ────────────────────────────────────────────────────────

/// Snapshot the current thread's entire binding stack (for conveyance into a
/// child thread, e.g. `future`).
pub fn capture_current() -> Vec<HashMap<VarKey, Value>> {
    BINDING_STACK.with(|s| {
        s.borrow()
            .iter()
            .map(|frame| {
                #[cfg(feature = "no-gc")]
                for lease in frame.leases.values() {
                    lease.protect_read();
                }
                frame.values.clone()
            })
            .collect()
    })
}

/// Install a previously captured snapshot on the current (new) thread.
pub fn install_frames(frames: Vec<HashMap<VarKey, Value>>) {
    let frames = frames.into_iter().map(BindingFrame::new).collect();
    BINDING_STACK.with(|s| *s.borrow_mut() = frames);
}

// ── GC root tracing ───────────────────────────────────────────────────────────

/// Trace all values in the current thread's binding stack as GC roots.
/// Call this during the GC root enumeration phase.
pub fn trace_current(visitor: &mut cljrs_gc::MarkVisitor) {
    BINDING_STACK.with(|s| {
        for frame in s.borrow().iter() {
            for val in frame.values.values() {
                val.trace(visitor);
            }
        }
    });
}

fn trace_saved(frames: &Vec<BindingFrame>, visitor: &mut cljrs_gc::MarkVisitor) {
    for frame in frames {
        for value in frame.values.values() {
            value.trace(visitor);
        }
    }
}