handball/src/eval.rs

use std::assert_matches::assert_matches;
use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

use crate::value::Value;
use crate::{ast, prims};

#[derive(/*debug2::Debug,*/ Debug, Clone)]
crate struct Lambda {
    crate func: Rc<ast::Func>,
    crate captures: Rc<RefCell<Env>>,
}

pub(crate) fn eval(t: &ast::Tree, env: Rc<RefCell<Env>>) -> Result<Value, RTError> {
    Ok(match t {
        ast::Tree::Leaf(l) => match l {
            ast::Literal::Sym(s) => match env.borrow().lookup(s) {
                Some(v) => v,
                None => return err(format!("Undefined variable `{}`", s)),
            },
            ast::Literal::Num(v) => Value::Num(*v),
            ast::Literal::Bool(b) => Value::Bool(*b),
        },
        ast::Tree::Func(l) => Value::Lambda(Lambda {
            func: Rc::clone(l),
            captures: Rc::clone(&env),
        }),
        ast::Tree::Branch(args) => {
            let Some(fun) = args.get(0) else { return err("No func given".to_owned()) };

            let fun = eval(fun, Rc::clone(&env))?.as_func()?;

            let args = args
                .iter()
                .skip(1)
                .map(|a| eval(a, Rc::clone(&env)))
                .collect::<Result<Vec<_>, _>>()?;

            match fun {
                Callable::Func(f) => f(&args)?,
                Callable::Lambda(l) => {
                    if l.func.args.len() == args.len() {
                        // let mut env = env.child();
                        let mut env = Env::child(Rc::clone(&l.captures));

                        for (x, y) in l.func.args.iter().zip(args) {
                            env.define(x.clone(), y)
                        }
                        let env = Rc::new(RefCell::new(env));

                        let [main @ .., tail] = &l.func.body[..] else { unreachable!("Body has 1+ element by parser") };
                        for i in main {
                            eval(i, Rc::clone(&env))?;
                        }

                        eval(tail, env)?
                    } else {
                        return err(format!(
                            "Need {} args, got {}",
                            l.func.args.len(),
                            args.len()
                        ));
                    }
                }
            }
        }
        ast::Tree::Define(name, to) => {
            let val = eval(to, Rc::clone(&env))?;
            env.borrow_mut().define(name.to_owned(), val);
            Value::Trap
        }
        ast::Tree::Set(name, to) => {
            let val = eval(to, Rc::clone(&env))?;
            env.borrow_mut().set(name, val)?;
            Value::Trap
        }
        ast::Tree::If(box [cond, tcase, fcase]) => {
            let b = eval(cond, Rc::clone(&env))?.as_bool()?;
            let body = if b { tcase } else { fcase };
            eval(body, env)?
        }
    })
}

#[derive(Debug)]
crate struct RTError(crate String);

pub(crate) fn err<T>(s: String) -> Result<T, RTError> {
    Err(RTError(s))
}

pub(crate) enum Callable {
    Func(prims::NativeFunc),
    Lambda(Lambda),
}

#[derive(Clone /*, debug2::Debug*/)]
pub(crate) struct Env {
    vars: HashMap<String, Value>,
    enclosing: Option<Rc<RefCell<Env>>>,
}

impl std::fmt::Debug for Env {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        struct Vars<'a>(&'a HashMap<String, Value>);

        impl std::fmt::Debug for Vars<'_> {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                f.debug_set().entries(self.0.keys()).finish()
            }
        }

        f.debug_struct("Env")
            .field("vars", &Vars(&self.vars))
            .field("enclosing", &self.enclosing)
            .finish()
    }
}

impl Env {
    pub(crate) fn child(this: Rc<RefCell<Self>>) -> Env {
        Env {
            vars: HashMap::new(),
            enclosing: Some(this),
        }
    }

    pub(crate) fn lookup(&self, s: &str) -> Option<Value> {
        if let Some(v) = self.vars.get(s) {
            Some(v.clone())
        } else if let Some(parent) = &self.enclosing {
            parent.borrow().lookup(s)
        } else {
            None
        }
    }

    pub(crate) fn define(&mut self, name: String, val: Value) {
        assert_ne!(val, Value::Trap); // TODO: Better error

        // TODO: Error on previous def
        self.vars.insert(name, val);
    }

    crate fn set(&mut self, name: &str, val: Value) -> Result<(), RTError> {
        if let Some(loc) = self.vars.get_mut(name) {
            *loc = val;
            Ok(())
        } else if let Some(parent) = &self.enclosing {
            parent.borrow_mut().set(name, val)
        } else {
            err(format!("Tried to `set!` un `define`d var `{}`", name))
        }
    }
}

pub(crate) fn default_env() -> Env {
    let mut vars = HashMap::new();

    for (name, fun) in prims::prims() {
        assert_matches!(vars.insert((*name).to_owned(), Value::Func(*fun)), None)
    }

    Env {
        vars,
        enclosing: None,
    }
}