puyo-lang/src/eval.rs

use core::fmt;
use std::{
    cell::RefCell,
    collections::{HashMap, hash_map::Entry},
    io::{self, Write},
    ops, process,
    rc::Rc,
};

use crate::inst::Inst;

type Key<'a> = &'a str;

#[derive(Debug, Clone)]
enum Val<'a> {
    Num(f64),
    Fun {
        insts: &'a [Inst<'a>],
        env: Rc<RefCell<Env<'a>>>,
    },
}

impl Default for Val<'_> {
    fn default() -> Self {
        Self::Num(0.)
    }
}

impl fmt::Display for Val<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Val::Num(n) => write!(f, "{n}"),
            Val::Fun { insts, env } => {
                write!(f, "{} {}", Inst::Block(insts.to_vec()), env.borrow())
            }
        }
    }
}

impl<'a> Val<'a> {
    fn to_num(&self) -> f64 {
        match self {
            Val::Num(n) => *n,
            Val::Fun { insts: _, env: _ } => 0.,
        }
    }

    fn call(&self, stack: &mut Stack<'a>) {
        match self {
            Val::Num(n) => stack.push(Val::Num(*n)),
            Val::Fun { insts, env } => {
                eval(insts, stack, &env.clone());
            }
        }
    }
}

type EnvMap<'a> = HashMap<Key<'a>, Val<'a>>;

#[derive(Debug, Clone)]
enum Env<'a> {
    Local(EnvMap<'a>),
    Closure {
        local: EnvMap<'a>,
        above: Rc<RefCell<Env<'a>>>,
    },
}

impl fmt::Display for Env<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Env::Local(map) => {
                for key in map.keys() {
                    write!(f, "{key} ")?;
                }
                Ok(())
            }
            Env::Closure { local, above } => {
                for key in local.keys() {
                    write!(f, "{key} ")?;
                }
                write!(f, "| ")?;
                write!(f, "{}", above.borrow())
            }
        }
    }
}

impl Default for Env<'_> {
    fn default() -> Self {
        Env::Local(HashMap::new())
    }
}

impl<'a> Env<'a> {
    fn set(&mut self, key: Key<'a>, val: Val<'a>) {
        match self {
            Env::Local(map) => {
                map.insert(key, val);
            }
            Env::Closure { local, above } => match local.entry(key) {
                Entry::Occupied(mut e) => {
                    e.insert(val);
                }
                Entry::Vacant(e) => {
                    if !above.borrow_mut().set_if_exists(key, val.clone()) {
                        e.insert(val);
                    }
                }
            },
        }
    }

    /// Returns whether the key was set
    fn set_if_exists(&mut self, key: Key<'a>, val: Val<'a>) -> bool {
        match self {
            Env::Local(map) => {
                if let Entry::Occupied(mut e) = map.entry(key) {
                    e.insert(val);
                    true
                } else {
                    false
                }
            }
            Env::Closure { local, above } => {
                if let Entry::Occupied(mut e) = local.entry(key) {
                    e.insert(val);
                    true
                } else {
                    above.borrow_mut().set_if_exists(key, val)
                }
            }
        }
    }

    fn get(&self, key: Key<'a>) -> Option<Val<'a>> {
        match self {
            Env::Local(map) => map.get(key).cloned(),
            Env::Closure { local, above } => local
                .get(key)
                .cloned()
                .or_else(|| above.borrow().get(key).clone()),
        }
    }
}

#[derive(Debug)]
struct Stack<'a>(Vec<Val<'a>>);

impl<'a> Stack<'a> {
    fn new() -> Self {
        Self(Vec::new())
    }

    fn pop(&mut self) -> Option<Val<'a>> {
        self.0.pop()
    }

    fn pop_num(&mut self) -> f64 {
        self.pop_default().to_num()
    }

    fn push(&mut self, val: Val<'a>) {
        self.0.push(val);
    }

    fn push_num(&mut self, n: f64) {
        self.push(Val::Num(n));
    }

    fn pop_default(&mut self) -> Val<'a> {
        self.pop().unwrap_or_default()
    }

    fn run_dyadic(&mut self, f: impl Fn(f64, f64) -> f64) {
        let b = self.pop_default();
        let a = self.pop_default();
        self.push_num(f(a.to_num(), b.to_num()));
    }

    fn run_monadic(&mut self, f: impl Fn(f64) -> f64) {
        let n = self.pop_num();
        self.push_num(f(n));
    }
}

fn eval<'a>(insts: &'a [Inst<'a>], stack: &mut Stack<'a>, env: &Rc<RefCell<Env<'a>>>) {
    for i in insts {
        match i {
            Inst::Set(var) => env.borrow_mut().set(var, stack.pop_default()),
            Inst::Var(var) => match *var {
                "+" => stack.run_dyadic(ops::Add::add),
                "-" => stack.run_dyadic(ops::Sub::sub),
                "*" => stack.run_dyadic(ops::Mul::mul),
                "/" => stack.run_dyadic(ops::Div::div),
                "%" | "rem" => stack.run_dyadic(ops::Rem::rem),
                "sin" => stack.run_monadic(f64::sin),
                "cos" => stack.run_monadic(f64::cos),
                "tan" => stack.run_monadic(f64::tan),
                "abs" => stack.run_monadic(f64::abs),
                "ceil" => stack.run_monadic(f64::ceil),
                "floor" => stack.run_monadic(f64::floor),
                "round" => stack.run_monadic(f64::round),
                "expt" => stack.run_dyadic(f64::powf),
                "=" => stack.run_dyadic(|a, b| ((a - b).abs() < 1e-10).into()),
                "<" => stack.run_dyadic(|a, b| (a < b).into()),
                "dup" => {
                    let top = stack.pop_default();
                    stack.push(top.clone());
                    stack.push(top);
                }
                "swap" => {
                    let a = stack.pop_default();
                    let b = stack.pop_default();
                    stack.push(a);
                    stack.push(b);
                }
                "print" => {
                    print!("{}", stack.pop_default());
                }
                "println" => {
                    println!("{}", stack.pop_default());
                }
                "put" => {
                    #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
                    let byte = stack.pop_num() as u8;
                    io::stdout().write_all(&[byte]).unwrap();
                }
                "explode" => {
                    eprintln!("I ESPLODED");
                    for (i, val) in stack.0.iter().rev().enumerate() {
                        eprintln!("{}. {val}", i + 1);
                    }
                    io::stdout().flush().unwrap();
                    io::stderr().flush().unwrap();
                    process::exit(3);
                }
                "ite" => {
                    let r#else = stack.pop_default();
                    let then = stack.pop_default();
                    let test = stack.pop_default();
                    if test.to_num() == 0.0 {
                        r#else.call(stack);
                    } else {
                        then.call(stack);
                    }
                }
                _ => stack.push(
                    env.borrow()
                        .get(var)
                        .unwrap_or_else(|| panic!("unknown variable: {var}")),
                ),
            },
            Inst::Num(n) => stack.push(Val::Num(*n)),
            Inst::Call => {
                let f = stack.pop_default();
                f.call(stack);
            }
            Inst::Block(insts) => stack.push(Val::Fun {
                insts,
                env: Rc::new(RefCell::new(Env::Closure {
                    local: HashMap::new(),
                    above: env.clone(),
                })),
            }),
        }
    }
}

pub fn run(insts: &[Inst<'_>]) {
    let mut stack = Stack::new();
    eval(insts, &mut stack, &Rc::new(RefCell::new(Env::default())));
}