lambda-throw-cat 0.1.0

//! Tree-walking interpreter with heap-threading, prototype-walking field
//! access, and `throw`/`try`-`catch` non-local control flow.
//!
//! Spike 4 adds two new evaluation cases:
//!
//! - [`Expr::Throw`]: evaluate the inner expression, then emit
//!   [`Error::Thrown`] carrying the resulting value plus the current heap
//!   and fuel.  Existing `?`-propagation through every other case unwinds
//!   the stack until either an enclosing [`Expr::TryCatch`] catches the
//!   signal or it reaches the top boundary as an uncaught exception.
//! - [`Expr::TryCatch`]: evaluate the body; on [`Error::Thrown`], bind the
//!   thrown value to the catch parameter and evaluate the handler with the
//!   heap and fuel that were live at the moment of the throw.  Other
//!   [`Error`] variants propagate unchanged.
//!
//! Substitution traverses the new variants without introducing new binders.
//!
//! [`Error::Thrown`]: crate::error::Error::Thrown

use std::collections::BTreeMap;

use crate::env::Env;
use crate::error::{Error, ThrownPayload};
use crate::heap::{Address, Heap};
use crate::syntax::{Expr, VarName};
use crate::value::Value;

/// A step budget for evaluation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Fuel {
    limit: u64,
    remaining: u64,
}

impl Fuel {
    /// A fresh budget of `limit` steps.
    #[must_use]
    pub fn new(limit: u64) -> Self {
        Self {
            limit,
            remaining: limit,
        }
    }

    /// The original budget.
    #[must_use]
    pub fn limit(&self) -> u64 {
        self.limit
    }

    /// Steps still available.
    #[must_use]
    pub fn remaining(&self) -> u64 {
        self.remaining
    }

    fn spend(self) -> Result<Self, Error> {
        match self.remaining {
            0 => Err(Error::FuelExhausted { limit: self.limit }),
            n => Ok(Self {
                limit: self.limit,
                remaining: n - 1,
            }),
        }
    }
}

/// Evaluate `expr` against `env` with the given heap and step budget.
///
/// # Errors
///
/// See [`Error`] for the full list of failure modes.
///
/// # Examples
///
/// ```
/// # fn main() -> Result<(), lambda_throw_cat::error::Error> {
/// use lambda_throw_cat::env::Env;
/// use lambda_throw_cat::eval::{eval, Fuel};
/// use lambda_throw_cat::heap::Heap;
/// use lambda_throw_cat::syntax::Expr;
///
/// let id = Expr::lam("x", Expr::var("x"));
/// let (value, _heap, _fuel) = eval(&id, &Env::empty(), Heap::empty(), Fuel::new(100))?;
/// assert_eq!(format!("{value}"), "\\x. x");
/// # Ok(())
/// # }
/// ```
pub fn eval(expr: &Expr, env: &Env, heap: Heap, fuel: Fuel) -> Result<(Value, Heap, Fuel), Error> {
    let fuel = fuel.spend()?;
    step(expr, env, heap, fuel)
}

fn step(expr: &Expr, env: &Env, heap: Heap, fuel: Fuel) -> Result<(Value, Heap, Fuel), Error> {
    match expr {
        Expr::Var(name) => eval_var(name, env, heap, fuel),
        Expr::Lam { param, body } => Ok((
            Value::closure(param.clone(), (**body).clone(), env.clone()),
            heap,
            fuel,
        )),
        Expr::App { func, arg } => eval_app(func, arg, env, heap, fuel),
        Expr::Let { name, value, body } => eval_let(name, value, body, env, heap, fuel),
        Expr::Fix { name, body } => eval_fix(expr, name, body, env, heap, fuel),
        Expr::Ref { inner } => eval_ref(inner, env, heap, fuel),
        Expr::Deref { inner } => eval_deref(inner, env, heap, fuel),
        Expr::Assign { target, value } => eval_assign(target, value, env, heap, fuel),
        Expr::Seq { first, second } => eval_seq(first, second, env, heap, fuel),
        Expr::Object { entries, prototype } => {
            eval_object(entries, prototype.as_deref(), env, heap, fuel)
        }
        Expr::Field { object, name } => eval_field(object, name, env, heap, fuel),
        Expr::Throw { inner } => eval_throw(inner, env, heap, fuel),
        Expr::TryCatch {
            body,
            catch_param,
            handler,
        } => eval_try_catch(body, catch_param, handler, env, heap, fuel),
    }
}

fn eval_var(
    name: &VarName,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    env.lookup(name)
        .cloned()
        .map(|v| (v, heap, fuel))
        .ok_or_else(|| Error::UnboundVariable { name: name.clone() })
}

fn eval_app(
    func: &Expr,
    arg: &Expr,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    let (func_v, heap, fuel) = eval(func, env, heap, fuel)?;
    let (arg_v, heap, fuel) = eval(arg, env, heap, fuel)?;
    apply(func_v, arg_v, heap, fuel)
}

fn eval_let(
    name: &VarName,
    value: &Expr,
    body: &Expr,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    let (value_v, heap, fuel) = eval(value, env, heap, fuel)?;
    let new_env = env.extend(name.clone(), value_v);
    eval(body, &new_env, heap, fuel)
}

fn eval_fix(
    whole: &Expr,
    name: &VarName,
    body: &Expr,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    let unfolded = subst(body, name, whole);
    eval(&unfolded, env, heap, fuel)
}

fn eval_ref(inner: &Expr, env: &Env, heap: Heap, fuel: Fuel) -> Result<(Value, Heap, Fuel), Error> {
    let (inner_v, heap, fuel) = eval(inner, env, heap, fuel)?;
    let (new_heap, addr) = heap.alloc(inner_v);
    Ok((Value::reference(addr), new_heap, fuel))
}

fn eval_deref(
    inner: &Expr,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    let (ref_v, heap, fuel) = eval(inner, env, heap, fuel)?;
    let addr = expect_ref(&ref_v)?;
    let cell_value = heap
        .load(addr)
        .cloned()
        .ok_or(Error::DanglingReference { address: addr })?;
    Ok((cell_value, heap, fuel))
}

fn eval_assign(
    target: &Expr,
    value: &Expr,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    match target {
        Expr::Field { object, name } => eval_field_assign(object, name, value, env, heap, fuel),
        Expr::Var(_)
        | Expr::Lam { .. }
        | Expr::App { .. }
        | Expr::Let { .. }
        | Expr::Fix { .. }
        | Expr::Ref { .. }
        | Expr::Deref { .. }
        | Expr::Assign { .. }
        | Expr::Seq { .. }
        | Expr::Object { .. }
        | Expr::Throw { .. }
        | Expr::TryCatch { .. } => eval_cell_assign(target, value, env, heap, fuel),
    }
}

fn eval_cell_assign(
    target: &Expr,
    value: &Expr,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    let (target_v, heap, fuel) = eval(target, env, heap, fuel)?;
    let addr = expect_ref(&target_v)?;
    let (value_v, heap, fuel) = eval(value, env, heap, fuel)?;
    let new_heap = heap
        .store(addr, value_v.clone())
        .ok_or(Error::DanglingReference { address: addr })?;
    Ok((value_v, new_heap, fuel))
}

fn eval_field_assign(
    object: &Expr,
    name: &VarName,
    value: &Expr,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    let (obj_v, heap, fuel) = eval(object, env, heap, fuel)?;
    let obj_addr = expect_ref(&obj_v)?;
    let (value_v, heap, fuel) = eval(value, env, heap, fuel)?;
    let cell = heap
        .load(obj_addr)
        .cloned()
        .ok_or(Error::DanglingReference { address: obj_addr })?;
    let new_obj = update_object_property(&cell, name, &value_v)?;
    let new_heap = heap
        .store(obj_addr, new_obj)
        .ok_or(Error::DanglingReference { address: obj_addr })?;
    Ok((value_v, new_heap, fuel))
}

fn eval_seq(
    first: &Expr,
    second: &Expr,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    let (_discarded, heap, fuel) = eval(first, env, heap, fuel)?;
    eval(second, env, heap, fuel)
}

fn eval_object(
    entries: &[(VarName, Expr)],
    prototype_expr: Option<&Expr>,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    let (prototype, heap, fuel) = eval_prototype(prototype_expr, env, heap, fuel)?;
    let (properties, heap, fuel) = eval_entries(entries, env, BTreeMap::new(), heap, fuel)?;
    let object = Value::object(properties, prototype);
    let (new_heap, addr) = heap.alloc(object);
    Ok((Value::reference(addr), new_heap, fuel))
}

fn eval_prototype(
    prototype_expr: Option<&Expr>,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Option<Address>, Heap, Fuel), Error> {
    prototype_expr
        .iter()
        .copied()
        .try_fold((None, heap, fuel), |(_, heap, fuel), expr| {
            let (value, heap, fuel) = eval(expr, env, heap, fuel)?;
            let addr = extract_object_reference(&value, &heap)?;
            Ok((Some(addr), heap, fuel))
        })
}

fn eval_entries(
    entries: &[(VarName, Expr)],
    env: &Env,
    initial: BTreeMap<VarName, Value>,
    heap: Heap,
    fuel: Fuel,
) -> Result<(BTreeMap<VarName, Value>, Heap, Fuel), Error> {
    entries.iter().try_fold(
        (initial, heap, fuel),
        |(props, heap, fuel), (name, value_expr)| {
            let (value, heap, fuel) = eval(value_expr, env, heap, fuel)?;
            let new_props: BTreeMap<VarName, Value> = props
                .into_iter()
                .filter(|(k, _)| k != name)
                .chain(std::iter::once((name.clone(), value)))
                .collect();
            Ok((new_props, heap, fuel))
        },
    )
}

fn eval_field(
    object: &Expr,
    name: &VarName,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    let (obj_v, heap, fuel) = eval(object, env, heap, fuel)?;
    let addr = expect_ref(&obj_v)?;
    let value = lookup_property(addr, name, &heap)?;
    Ok((value, heap, fuel))
}

fn eval_throw(
    inner: &Expr,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    let (value, heap, fuel) = eval(inner, env, heap, fuel)?;
    Err(Error::Thrown(Box::new(ThrownPayload::new(
        value, heap, fuel,
    ))))
}

fn eval_try_catch(
    body: &Expr,
    catch_param: &VarName,
    handler: &Expr,
    env: &Env,
    heap: Heap,
    fuel: Fuel,
) -> Result<(Value, Heap, Fuel), Error> {
    match eval(body, env, heap, fuel) {
        Ok(triple) => Ok(triple),
        Err(Error::Thrown(payload)) => {
            let (value, heap, fuel) = payload.into_parts();
            let new_env = env.extend(catch_param.clone(), value);
            eval(handler, &new_env, heap, fuel)
        }
        Err(other) => Err(other),
    }
}

fn lookup_property(addr: Address, name: &VarName, heap: &Heap) -> Result<Value, Error> {
    let cell = heap
        .load(addr)
        .ok_or(Error::DanglingReference { address: addr })?;
    let (properties, prototype) = expect_object_fields(cell)?;
    properties
        .get(name)
        .cloned()
        .map_or_else(|| lookup_in_proto(prototype, name, heap), Ok)
}

fn lookup_in_proto(
    prototype: Option<Address>,
    name: &VarName,
    heap: &Heap,
) -> Result<Value, Error> {
    prototype.map_or_else(
        || Err(Error::PropertyNotFound { name: name.clone() }),
        |addr| lookup_property(addr, name, heap),
    )
}

fn apply(func: Value, arg: Value, heap: Heap, fuel: Fuel) -> Result<(Value, Heap, Fuel), Error> {
    match func {
        Value::Closure { param, body, env } => {
            let new_env = env.extend(param, arg);
            eval(&body, &new_env, heap, fuel)
        }
        Value::Ref(address) => Err(Error::NotAFunction {
            found: format!("ref({address})"),
        }),
        Value::Object { .. } => Err(Error::NotAFunction {
            found: format!("{func}"),
        }),
    }
}

fn expect_ref(value: &Value) -> Result<Address, Error> {
    match value {
        Value::Ref(address) => Ok(*address),
        Value::Closure { .. } | Value::Object { .. } => Err(Error::NotAReference {
            found: format!("{value}"),
        }),
    }
}

fn expect_object_fields(
    value: &Value,
) -> Result<(&BTreeMap<VarName, Value>, Option<Address>), Error> {
    match value {
        Value::Object {
            properties,
            prototype,
        } => Ok((properties, *prototype)),
        Value::Closure { .. } | Value::Ref(_) => Err(Error::NotAnObject {
            found: format!("{value}"),
        }),
    }
}

fn update_object_property(cell: &Value, name: &VarName, value: &Value) -> Result<Value, Error> {
    match cell {
        Value::Object {
            properties,
            prototype,
        } => {
            let new_props: BTreeMap<VarName, Value> = properties
                .iter()
                .filter(|(k, _)| *k != name)
                .map(|(k, v)| (k.clone(), v.clone()))
                .chain(std::iter::once((name.clone(), value.clone())))
                .collect();
            Ok(Value::object(new_props, *prototype))
        }
        Value::Closure { .. } | Value::Ref(_) => Err(Error::NotAnObject {
            found: format!("{cell}"),
        }),
    }
}

fn extract_object_reference(value: &Value, heap: &Heap) -> Result<Address, Error> {
    let addr = expect_ref(value)?;
    let cell = heap
        .load(addr)
        .ok_or(Error::DanglingReference { address: addr })?;
    expect_object_fields(cell).map(|_| addr)
}

fn subst(expr: &Expr, target: &VarName, replacement: &Expr) -> Expr {
    match expr {
        Expr::Var(name) => {
            if name == target {
                replacement.clone()
            } else {
                Expr::Var(name.clone())
            }
        }
        Expr::Lam { param, body } => Expr::Lam {
            param: param.clone(),
            body: Box::new(subst_under_binder(body, param, target, replacement)),
        },
        Expr::App { func, arg } => Expr::App {
            func: Box::new(subst(func, target, replacement)),
            arg: Box::new(subst(arg, target, replacement)),
        },
        Expr::Let { name, value, body } => Expr::Let {
            name: name.clone(),
            value: Box::new(subst(value, target, replacement)),
            body: Box::new(subst_under_binder(body, name, target, replacement)),
        },
        Expr::Fix { name, body } => Expr::Fix {
            name: name.clone(),
            body: Box::new(subst_under_binder(body, name, target, replacement)),
        },
        Expr::Ref { inner } => Expr::Ref {
            inner: Box::new(subst(inner, target, replacement)),
        },
        Expr::Deref { inner } => Expr::Deref {
            inner: Box::new(subst(inner, target, replacement)),
        },
        Expr::Assign {
            target: t,
            value: v,
        } => Expr::Assign {
            target: Box::new(subst(t, target, replacement)),
            value: Box::new(subst(v, target, replacement)),
        },
        Expr::Seq { first, second } => Expr::Seq {
            first: Box::new(subst(first, target, replacement)),
            second: Box::new(subst(second, target, replacement)),
        },
        Expr::Object { entries, prototype } => Expr::Object {
            entries: entries
                .iter()
                .map(|(k, v)| (k.clone(), subst(v, target, replacement)))
                .collect(),
            prototype: prototype
                .as_ref()
                .map(|p| Box::new(subst(p, target, replacement))),
        },
        Expr::Field { object, name } => Expr::Field {
            object: Box::new(subst(object, target, replacement)),
            name: name.clone(),
        },
        Expr::Throw { inner } => Expr::Throw {
            inner: Box::new(subst(inner, target, replacement)),
        },
        Expr::TryCatch {
            body,
            catch_param,
            handler,
        } => Expr::TryCatch {
            body: Box::new(subst(body, target, replacement)),
            catch_param: catch_param.clone(),
            handler: Box::new(subst_under_binder(
                handler,
                catch_param,
                target,
                replacement,
            )),
        },
    }
}

fn subst_under_binder(body: &Expr, binder: &VarName, target: &VarName, replacement: &Expr) -> Expr {
    if binder == target {
        body.clone()
    } else {
        subst(body, target, replacement)
    }
}

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

    fn empty_fuel() -> Fuel {
        Fuel::new(10_000)
    }

    #[test]
    fn empty_object_allocates_one_cell() -> Result<(), Error> {
        let e = Expr::object(vec![]);
        let (value, heap, _fuel) = eval(&e, &Env::empty(), Heap::empty(), empty_fuel())?;
        let is_ref = matches!(value, Value::Ref(_));
        let one_cell = heap.len() == 1;
        (is_ref && one_cell)
            .then_some(())
            .ok_or(Error::UnboundVariable {
                name: VarName::from("expected one cell"),
            })
    }
}