narju/eval/

mod.rs

//! The evaluator: staging state ([`EvalCtx`]), the lift operator, and
//! the entry points [`evalms`] / [`evalmsg`].
//!
//! Evaluation proper is the CK machine in [`machine`]; the `trans`
//! translation from source-as-data to expressions is in [`staging`].
//! This module holds what the machine threads through every step — the
//! state base.scm keeps in global mutable variables (`fresh`, the open
//! statement block, the `stFun` memoization registry, the `run` level)
//! plus narju's additions: the cell heap and the I/O handle.
//!
//! The staging model follows the paper's normalization-by-evaluation
//! reading: evaluating under `lift` produces *values of kind code*,
//! and every reflected statement lands in the innermost open block,
//! which `reifyv`/`reifyc` drain into ANF Let-chains when their scope
//! closes.

pub mod machine;
pub mod staging;

use std::rc::Rc;

use crate::{
    core::{code, rc_exp, Env, Exp, RcExp, Val},
    error::NarjuError,
    io::{terminal::TerminalIo, NarjuIo},
};

/// One stFun registry key: the `(env, body)` of a closure whose
/// η-expansion has been memoised at some de Bruijn level. Keys hold `Rc`
/// clones so the allocations are pinned alive for the entry's lifetime —
/// earlier versions keyed on `Rc::as_ptr` integers and suffered stale
/// hits after the originals were freed and their addresses reused (the
/// F2-self-compilation OOM/wrong-value bug). Matching is *structural*
/// (`val_equal`/`env_equal` below), following base.scm:83's `equal?` on
/// stFun entries — so structurally identical closures reaching separate
/// lift sites share one residual lambda.
#[derive(Clone)]
struct FunKey {
    /// The closure's captured environment.
    env: Env,
    /// The closure's body — shared with the originating `Lam` node.
    body: RcExp,
}

/// Full-structural value equality — the reference's `equal?` (base.scm:83
/// matches stFun entries with `(equal? env ...)` / `(equal? e ...)`).
///
/// Deliberately distinct from `Val::PartialEq` (which backs the `eq?`
/// primitive and compares closures by body only): base.scm closures are
/// plain `(clo env body)` lists, so `equal?` recurses into the captured
/// env as well. Using the body-only compare here would conflate closures
/// that differ only in captured state — an incorrect memoization hit.
///
/// Exp comparison delegates to the derived `PartialEq`; its one Val
/// embedding (`Exp::Proc`) routes through `Val::PartialEq`'s body-only
/// Clo arm, but Proc cannot occur in source `Lam` bodies (it's only
/// constructed by staged `log`), so FunKey bodies don't hit that corner
/// in practice.
fn val_equal(a: &Val, b: &Val) -> bool {
    match (a, b) {
        (Val::Cst(x), Val::Cst(y)) => x == y,
        (Val::Sym(x), Val::Sym(y)) => x == y,
        (Val::Nil, Val::Nil) => true,
        (Val::Cell(x), Val::Cell(y)) => x == y,
        (Val::Tup(a1, a2), Val::Tup(b1, b2)) => {
            (Rc::ptr_eq(a1, b1) || val_equal(a1, b1))
                && (Rc::ptr_eq(a2, b2) || val_equal(a2, b2))
        }
        (Val::Code(x), Val::Code(y)) => Rc::ptr_eq(x, y) || x == y,
        (Val::Clo(ea, ba), Val::Clo(eb, bb)) => {
            (Rc::ptr_eq(ba, bb) || ba == bb) && env_equal(ea, eb)
        }
        _ => false,
    }
}

/// Structural equality on environments, with an `Rc::ptr_eq` fast path.
fn env_equal(a: &Env, b: &Env) -> bool {
    Rc::ptr_eq(a, b)
        || (a.len() == b.len() && a.iter().zip(b.iter()).all(|(x, y)| val_equal(x, y)))
}

/// The stFun analogue: an assoc list scanned linearly with structural
/// equality, exactly the reference's `List[(Int,Env,Exp)]` + `findFun`
/// shape. Kept in an `Rc` so `ScopeSnapshot` save/restore is an O(1)
/// refcount bump; `register_fun` copies-on-write via `Rc::make_mut`.
type FunRegistry = Rc<Vec<(FunKey, usize)>>;

/// Staging and effect state threaded through evaluation.
///
/// One `EvalCtx` is one evaluator instance: it owns the cell heap and
/// the I/O handle, and carries the staging bookkeeping that base.scm
/// keeps in global mutable variables. The staging fields (`fresh`,
/// `block`, `fun`) are saved and restored around each reify scope;
/// `cells`, `level`, and `io` deliberately are not — cell writes and
/// output are effects, and `level` has its own save/restore in the
/// `run`-now path.
pub struct EvalCtx {
    /// Fresh-name counter: the index the next reflected statement's
    /// residual variable will get (the reference's `fresh`).
    pub fresh: usize,
    /// Statements of the innermost open reify scope, in order. Drained
    /// into an ANF Let-chain when the scope closes.
    pub block: Vec<Exp>,
    /// The stFun registry; see [`FunRegistry`].
    fun: FunRegistry,
    /// Current `run` nesting depth: incremented for the
    /// compile-then-execute phase of a run-now and restored after.
    /// The raw REPL displays it in the prompt.
    pub level: usize,
    /// Where `read`, `print`, and unstaged `log` perform their I/O.
    pub io: Box<dyn NarjuIo>,
    /// First-class cell store: `Val::Cell(i)` indexes here. Cells are
    /// per-instance and in-process by design: each evaluator owns its
    /// heap, so instances can be snapshotted or run side by side without
    /// sharing mutable state. Deliberately NOT part of
    /// `ScopeSnapshot`: cell writes are effects and survive scope
    /// save/restore — `catch` does not roll them back.
    pub cells: Vec<Val>,
}

impl EvalCtx {
    /// A fresh evaluator over the given I/O implementation. Use
    /// [`crate::io::headless::HeadlessIo`] for tests and embedding;
    /// `Default` gives a terminal-backed context and panics without a
    /// TTY.
    pub fn new(io: Box<dyn NarjuIo>) -> Self {
        Self {
            fresh: 0,
            block: Vec::new(),
            fun: Rc::new(Vec::new()),
            level: 0,
            io,
            cells: Vec::new(),
        }
    }

    /// Allocate a fresh cell holding `v`; returns the `Val::Cell` handle.
    pub fn cell_new(&mut self, v: Val) -> Val {
        self.cells.push(v);
        Val::Cell(self.cells.len() - 1)
    }

    /// Number of memoised closures in stFun — used by :ctx command.
    pub fn fun_len(&self) -> usize {
        self.fun.len()
    }

    /// Allocate the next residual variable index.
    pub fn fresh_var(&mut self) -> usize {
        let v = self.fresh;
        self.fresh += 1;
        v
    }

    /// Run `f` with the staging fields saved and restored around it —
    /// the shared skeleton under [`EvalCtx::reifyv`].
    fn run_scope<F, T>(&mut self, f: F) -> Result<T, NarjuError>
    where
        F: FnOnce(&mut Self) -> Result<T, NarjuError>,
    {
        let saved_fresh = self.fresh;
        let saved_block = std::mem::take(&mut self.block);
        // O(1) Rc clone — just a reference count bump, not a Vec copy.
        let saved_fun = Rc::clone(&self.fun);
        let result = f(self);
        self.fresh = saved_fresh;
        self.block = saved_block;
        self.fun = saved_fun;
        result
    }

    /// Emit a statement into the current block and return the residual
    /// variable that names its result — the paper's `reflect`. This is
    /// where ANF comes from: every staged operation becomes one `Let`
    /// binding when the enclosing scope drains its block.
    pub fn reflect(&mut self, e: Exp) -> Exp {
        self.block.push(e);
        Exp::Var(self.fresh_var())
    }

    /// [`EvalCtx::reflect`], wrapped as a code value.
    pub fn reflectc(&mut self, e: Exp) -> Val {
        code(self.reflect(e))
    }

    /// Wrap a value-producing thunk in a reify-v scope: save the current
    /// staging state, run the thunk with an empty block, then either
    /// return the thunk's val (block empty) or wrap it in a Let-chain
    /// (block non-empty, val must be Code).
    ///
    /// This is the only public reify-style method retained after the
    /// CK rewrite. The Machine's `Cont::ReifyVExit` implements the same
    /// logic iteratively for in-Machine reifyv scopes; this method
    /// exists for the *top-level* entry `evalmsg`, which wraps one
    /// Machine run in a single reifyv scope.
    pub fn reifyv<F>(&mut self, f: F) -> Result<Val, NarjuError>
    where
        F: FnOnce(&mut Self) -> Result<Val, NarjuError>,
    {
        self.run_scope(|ctx| {
            ctx.block = Vec::new();
            let res = f(ctx)?;
            let stmts = std::mem::take(&mut ctx.block);
            if stmts.is_empty() {
                Ok(res)
            } else {
                let last = match res {
                    Val::Code(e) => Rc::unwrap_or_clone(e),
                    v => {
                        return Err(NarjuError::Stage(format!(
                            "reifyv: non-empty block but non-Code result: {v:?}"
                        )))
                    }
                };
                let e = stmts
                    .into_iter()
                    .rev()
                    .fold(last, |acc, stmt| Exp::Let(rc_exp(stmt), rc_exp(acc)));
                Ok(code(e))
            }
        })
    }

    /// Look up a closure in the stFun registry, returning the level of
    /// its already-emitted η-expansion if present. A hit means a
    /// recursive (or repeated) lift of the same closure resolves to a
    /// back-reference instead of expanding forever.
    fn find_fun(&self, env: &Env, body: &RcExp) -> Option<usize> {
        // Body first: an `Rc::ptr_eq` hit on the same Lam node is the
        // common case and cheap; env comparison only runs when it passes.
        self.fun.iter().find_map(|(k, level)| {
            let body_eq = Rc::ptr_eq(&k.body, body) || *k.body == **body;
            if body_eq && env_equal(&k.env, env) {
                Some(*level)
            } else {
                None
            }
        })
    }

    /// Record a closure's η-expansion level in the stFun registry
    /// (copy-on-write, so snapshots taken before this stay intact).
    fn register_fun(&mut self, level: usize, env: &Env, body: &RcExp) {
        Rc::make_mut(&mut self.fun).push((
            FunKey {
                env: Rc::clone(env),
                body: Rc::clone(body),
            },
            level,
        ));
    }
}

impl Default for EvalCtx {
    fn default() -> Self {
        let io = TerminalIo::new().unwrap_or_else(|e| panic!("failed to init terminal: {e}"));
        Self::new(Box::new(io))
    }
}

/// Public evaluator entry point — drives the CK machine.
///
/// Consumes the input `Exp` so the top-level `Rc` wrap is a `Rc::new`,
/// not a deep `e.clone()` then wrap. Callers (in `repl::TopLevel` and
/// tests) already own their `Exp` after `parse::lower`, so they have no
/// reason to keep a borrow alive past this call.
///
/// The Machine handles every `Exp` arm natively via the kont stack;
/// there is no recursive baseline. External callers (and `evalmsg`)
/// go through this entry point.
pub fn evalms(env: &Env, e: Exp, ctx: &mut EvalCtx) -> Result<Val, NarjuError> {
    machine::Machine::new(Rc::clone(env), rc_exp(e)).run(ctx)
}

// ── NBE-style lift operator ───────────────────────────────────────────────────

/// Lift a non-closure value to an Exp. Called from machine.rs's
/// `dispatch_unop::Lift` for all variants except `Val::Clo` — the Clo
/// case is handled directly by the CK machine via `Cont::ForceCode`
/// driving the η-expansion through the kont stack.
///
/// This function is **non-recursive**: it never calls back into the
/// evaluator. All branches are O(1) or simple constructions over already-
/// `Code` operands.
pub(super) fn lift(v: Val, ctx: &mut EvalCtx) -> Result<Exp, NarjuError> {
    match v {
        Val::Cst(n) => Ok(Exp::Lit(n)),
        Val::Nil => Ok(Exp::Nil),
        Val::Sym(s) => Ok(Exp::Sym(s)),

        Val::Code(e) => Ok(ctx.reflect(Exp::Lift(e))),

        Val::Tup(a, b) => {
            let ca = match &*a {
                Val::Code(e) => Rc::clone(e),
                v => {
                    return Err(NarjuError::Stage(format!(
                        "lift Tup: car must be Code, got {v:?}"
                    )))
                }
            };
            let cb = match &*b {
                Val::Code(e) => Rc::clone(e),
                v => {
                    return Err(NarjuError::Stage(format!(
                        "lift Tup: cdr must be Code, got {v:?}"
                    )))
                }
            };
            Ok(ctx.reflect(Exp::Cons(ca, cb)))
        }

        Val::Clo(_, _) => {
            // The CK machine's `Cont::ForceCode` handles Clo η-expansion
            // iteratively. dispatch_unop::Lift's caller intercepts Clo
            // before reaching this function, so this branch is dead in
            // practice. Kept as a defensive error rather than
            // `unreachable!` because external callers (if any) shouldn't
            // crash the process.
            Err(NarjuError::Stage(
                "lift(Clo) called from non-CK path — \
                 closure η-expansion now runs via Cont::ForceCode"
                    .into(),
            ))
        }

        // Staging rule: pure bindings fold, cells never fold.
        // Folding a cell into residual syntax would freeze its contents
        // at staging time. Staged cell *ops* residualize (see
        // dispatch_unop/dispatch_binop); the cell value itself is
        // unliftable.
        Val::Cell(_) => Err(NarjuError::Stage(
            "cannot lift a cell: cells never fold under staging".into(),
        )),
    }
}
// ── Top-level entry point ─────────────────────────────────────────────────────

/// Paper: `evalmsg(env, e) = reifyv(lambda () evalms(env, e))`.
///
/// Consumes `e` and threads it into the inner closure via `move`. The
/// previous `&Exp` signature forced a deep clone at the `evalms` entry;
/// the new shape lets `e` flow straight through to the Machine.
pub fn evalmsg(env: &Env, e: Exp, ctx: &mut EvalCtx) -> Result<Val, NarjuError> {
    ctx.reifyv(move |ctx| evalms(env, e, ctx))
}

// ── Prompt formatting ─────────────────────────────────────────────────────────

/// Convert a Val to a coloured REPL prompt string, matching the
/// aesthetic of the raw λ↑↓ REPL's `make_prompt`.
///
/// - `Val::Sym(s)`  → `"<ROSE>s<DIM>›<RESET> "`
/// - `Val::Tup`     → walks the cons-list, joining elements:
///   `(cons 'purple (cons 0 nil))` → `"<ROSE>purple<DIM>[<BLOOM>0<DIM>]›<RESET> "`
///   First element is the name, subsequent elements are bracketed.
/// - anything else  → bare `"<ROSE>purple<DIM>›<RESET> "`
///
/// ANSI escapes are emitted unconditionally — `HeadlessIo` never displays
/// the prompt (it only returns `None` from `read_line`), and `TerminalIo`
/// uses rustyline which handles ANSI escapes transparently.
fn val_to_prompt(v: &Val) -> String {
    use crate::colours::{BLOOM, DIM, RESET, ROSE};

    let render = |name: &str, levels: &[String]| -> String {
        if levels.is_empty() {
            format!("{ROSE}{name}{DIM}›{RESET} ")
        } else {
            let levels = levels.join(",");
            format!("{ROSE}{name}{DIM}[{RESET}{BLOOM}{levels}{DIM}]›{RESET} ")
        }
    };

    match v {
        Val::Sym(s) => render(s, &[]),
        Val::Tup(_, _) => {
            let mut parts = Vec::new();
            let mut cur = v;
            loop {
                match cur {
                    Val::Tup(h, t) => {
                        parts.push(match h.as_ref() {
                            Val::Sym(s) => s.to_string(),
                            Val::Cst(n) => n.to_string(),
                            _ => "?".to_string(),
                        });
                        cur = t;
                    }
                    Val::Nil => break,
                    _ => break,
                }
            }
            if parts.is_empty() {
                return render("purple", &[]);
            }
            let name = parts[0].clone();
            let levels: Vec<String> = parts[1..].to_vec();
            render(&name, &levels)
        }
        _ => render("purple", &[]),
    }
}

/// Render a value the way the REPL prints results. Identical to the
/// `fmt::Display` impl on [`Val`]; kept as a named function because
/// call sites that format for `catch` payloads and prompts read better
/// with one.
pub fn display(v: &Val) -> String {
    match v {
        Val::Cst(n) => n.to_string(),
        Val::Nil => "nil".into(),
        Val::Sym(s) => format!("'{s}"),
        Val::Clo(_, _) => "#<closure>".into(),
        Val::Code(e) => format!("#<code {} nodes>", crate::core::count_nodes(e)),
        Val::Tup(a, b) => display_list(a, b),
        Val::Cell(i) => format!("#<cell {i}>"),
    }
}

/// Walk a cons-list for display, given head and tail as `RcVal`.
fn display_list(head: &crate::core::RcVal, tail: &crate::core::RcVal) -> String {
    let mut items = vec![display(head.as_ref())];
    let mut cur = tail.as_ref();
    loop {
        match cur {
            Val::Tup(h, t) => {
                items.push(display(h.as_ref()));
                cur = t.as_ref();
            }
            Val::Nil => return format!("({})", items.join(" ")),
            v => return format!("({} . {})", items.join(" "), display(v)),
        }
    }
}