minilua/src/run.rs

use std::{collections::HashMap, fmt::Debug, rc::Rc};

use crate::{
	ast::{Constant, Expression, Location, Statement},
	types::{Value, ValueInner},
};

#[derive(Clone)]
struct Function {
	name: &'static str,
	inner: Rc<dyn Fn(&mut Scope, Vec<Value>) -> Value>,
}

impl Debug for Function {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		write!(f, "<function {}>", self.name)
	}
}

#[derive(Debug, Clone, Default)]
pub struct Scope {
	variables: HashMap<String, ValueInner>,
	functions: HashMap<&'static str, Function>,
}

impl Scope {
	pub fn get(&self, name: impl AsRef<str>) -> Value {
		self.variables.get(name.as_ref()).cloned()
	}

	pub fn set(&mut self, name: String, value: Value) {
		if let Some(value) = value {
			self.variables.insert(name, value);
		} else {
			self.variables.remove(&name);
		}
	}

	pub fn add_fn(
		&mut self,
		name: &'static str,
		f: impl Fn(&mut Scope, Vec<Value>) -> Value + 'static,
	) {
		self.functions.insert(
			name,
			Function {
				name: name,
				inner: Rc::new(f),
			},
		);
	}
}

pub trait Eval {
	fn eval(&self, scope: &mut Scope) -> Value;
}

impl Eval for Constant {
	fn eval(&self, _scope: &mut Scope) -> Value {
		match self {
			Constant::Nil => None,
			Constant::Boolean(value) => Some(ValueInner::Boolean(*value)),
			Constant::Number(value) => Some(ValueInner::Number(*value)),
			Constant::String(value) => Some(ValueInner::String(value.clone())),
		}
	}
}

impl Eval for Location {
	fn eval(&self, scope: &mut Scope) -> Value {
		match self {
			Location::Variable { name } => scope.get(name),
			Location::Field { table, index } => {
				let table = table.eval(scope);
				let Some(ValueInner::Table(table)) = table else {
					panic!("attempt to index non-table value {table:?}")
				};
				let index = index.eval(scope);
				let Some(index) = index else {
					panic!("attempt to index with a nil value")
				};
				table.get(index)
			}
		}
	}
}

impl Eval for Expression {
	fn eval(&self, scope: &mut Scope) -> Value {
		match self {
			Expression::Constant(inner) => inner.eval(scope),
			Expression::Variable(inner) => inner.eval(scope),
			Expression::Call { callee, args } => {
				let Some(function) = scope.functions.get(callee.as_str()).cloned() else {
					panic!("attempt to call non-existent function {callee}")
				};
				let args = args.into_iter().map(|arg| arg.eval(scope)).collect();
				(function.inner)(scope, args)
			}
		}
	}
}

pub trait Exec {
	fn exec(&self, scope: &mut Scope);
}

impl Exec for Statement {
	fn exec(&self, scope: &mut Scope) {
		match self {
			Statement::Assign { target, source } => {
				let value = source.eval(scope);
				match target {
					Location::Variable { name } => {
						scope.set(name.into(), value);
					}
					Location::Field { table, index } => {
						let table = table.eval(scope);
						let Some(ValueInner::Table(table)) = table else {
							panic!("attempt to index non-table value {table:?}")
						};
						let index = index.eval(scope);
						let Some(index) = index else {
							panic!("attempt to index with a nil value")
						};
						table.set(index, value);
					}
				};
			}
			Statement::Call { callee, args } => {
				Expression::Call {
					callee: callee.clone(),
					args: args.clone(),
				}
				.eval(scope);
			}
		}
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::{ast, types};

	#[test]
	fn test_consts() {
		let mut scope = Scope::default();
		assert_eq!(ast::Constant::Nil.eval(&mut scope), None);
		assert_eq!(
			ast::Constant::Boolean(true).eval(&mut scope),
			Some(true.into())
		);
		assert_eq!(ast::Constant::Number(42).eval(&mut scope), Some(42.into()));
		assert_eq!(
			ast::Constant::String("foobar".into()).eval(&mut scope),
			Some("foobar".into())
		);
	}

	#[test]
	fn test_vars() {
		let mut scope = Scope::default();
		let foo = ast::Location::Variable { name: "foo".into() };
		let bar = ast::Location::Variable { name: "bar".into() };
		assert_eq!(foo.eval(&mut scope), None);
		assert_eq!(bar.eval(&mut scope), None);
		scope.variables.insert("bar".into(), 42.into());
		assert_eq!(foo.eval(&mut scope), None);
		assert_eq!(bar.eval(&mut scope), Some(42.into()));
	}

	#[test]
	fn test_fields() {
		let mut scope = Scope::default();
		let loc = ast::Location::Field {
			table: Box::new(ast::Location::Variable { name: "foo".into() }),
			index: Box::new(ast::Expression::Constant(ast::Constant::String(
				"bar".into(),
			))),
		};
		let foo = types::Table::default();
		let bar = types::Table::default();
		scope.variables.insert("foo".into(), foo.clone().into());
		assert_eq!(loc.eval(&mut scope), None);
		scope.variables.insert("bar".into(), bar.clone().into());
		assert_eq!(loc.eval(&mut scope), None);
		bar.set("foo".into(), Some(666.into()));
		assert_eq!(loc.eval(&mut scope), None);
		foo.set("bar".into(), Some(42.into()));
		assert_eq!(loc.eval(&mut scope), Some(42.into()));
	}

	#[test]
	fn test_call() {
		let mut scope = Scope::default();
		scope.add_fn("sqr", |_, args| {
			let [arg] = args.as_slice() else {
				panic!("exactly one argument expected");
			};
			let Some(ValueInner::Number(arg)) = arg else {
				panic!("number expected");
			};
			Some(ValueInner::Number(arg * arg))
		});
		assert_eq!(
			ast::Expression::Call {
				callee: "sqr".into(),
				args: vec![ast::Expression::Constant(ast::Constant::Number(7))],
			}
			.eval(&mut scope),
			Some(49.into())
		);
	}

	#[test]
	fn test_var_assign() {
		let mut scope = Scope::default();
		assert_eq!(scope.get("foo"), None);
		assert_eq!(scope.get("bar"), None);
		ast::Statement::Assign {
			target: ast::Location::Variable { name: "foo".into() },
			source: Box::new(ast::Expression::Constant(ast::Constant::Number(42))),
		}
		.exec(&mut scope);
		assert_eq!(scope.get("foo"), Some(42.into()));
		assert_eq!(scope.get("bar"), None);
	}

	#[test]
	fn test_call_statement() {
		let mut scope = Scope::default();
		scope.add_fn("set", |scope, args| {
			let [key, value] = args.as_slice() else {
				panic!("exactly two arguments expected");
			};
			let Some(ValueInner::String(key)) = key else {
				panic!("string expected");
			};
			scope.set(key.clone(), value.clone());
			None
		});
		assert_eq!(scope.get("foo"), None);
		assert_eq!(scope.get("bar"), None);
		ast::Statement::Call {
			callee: "set".into(),
			args: vec![
				ast::Expression::Constant(ast::Constant::String("foo".into())),
				ast::Expression::Constant(ast::Constant::Number(42)),
			],
		}
		.exec(&mut scope);
		assert_eq!(scope.get("foo"), Some(42.into()));
		assert_eq!(scope.get("bar"), None);
	}
}