minilua/src/twopass.rs

use std::cell::RefCell;

use crate::ast;

#[derive(Debug, Clone)]
pub struct Function {
	name: Option<ast::Ident>,
	args: Vec<ast::Ident>,
	locals: Vec<ast::Ident>,
	upvalues: Vec<(ast::Ident, Ident)>,
	body: Vec<Statement>,
	ret: Box<Expression>,
}

#[derive(Debug, Clone, Copy)]
enum Ident {
	Upvalue(usize),
	Local(usize),
	Argument(usize),
}

#[derive(Debug, Clone)]
enum Statement {
	Assign {
		target: Location,
		source: Box<Expression>,
	},
	Call(Call),
}

#[derive(Debug, Clone)]
enum Location {
	Variable {
		id: Ident,
	},
	Field {
		table: Box<Location>,
		index: Box<Expression>,
	},
}

#[derive(Debug, Clone)]
enum Expression {
	Constant(ast::Constant),
	Variable(Location),
	Call(Call),
	Function(Function),
}

#[derive(Debug, Clone)]
struct Call {
	callee: Box<Expression>,
	args: Vec<Expression>,
}

mod scope {
	use std::collections::HashMap;

	use super::Ident;
	use crate::ast;

	#[derive(Debug, Clone, Default)]
	pub struct Scope {
		locals: Vec<ast::Ident>,
		upvalues: Vec<(ast::Ident, Ident)>,
		scope: HashMap<ast::Ident, Ident>,
	}

	impl Scope {
		pub fn new_toplevel() -> Self {
			Self {
				locals: vec!["_ENV".to_string()],
				upvalues: vec![],
				scope: [("_ENV".to_string(), Ident::Local(0))]
					.into_iter()
					.collect(),
			}
		}

		pub fn lookup(&self, name: &ast::Ident) -> Option<Ident> {
			self.scope.get(name).copied()
		}

		fn add(&mut self, name: &ast::Ident, ident: Ident) -> Ident {
			self.scope.insert(name.clone(), ident);
			ident
		}

		pub fn add_arg(&mut self, name: &ast::Ident, index: usize) {
			self.scope.insert(name.clone(), Ident::Argument(index));
		}

		pub fn add_local(&mut self, name: &ast::Ident) -> Ident {
			let index = self.locals.len();
			self.locals.push(name.clone());
			self.add(name, Ident::Local(index))
		}

		pub fn add_upvalue(&mut self, name: &ast::Ident, up_ident: Ident) -> Ident {
			let index = self.upvalues.len();
			self.upvalues.push((name.clone(), up_ident));
			self.add(name, Ident::Upvalue(index))
		}
	}
}
use scope::Scope;

#[derive(Debug)]
struct BuildContext<'a> {
	parent: Option<&'a BuildContext<'a>>, // mustn't be &mut because of variance
	scope: RefCell<Scope>,
}

impl BuildContext<'_> {
	fn request(&self, name: &ast::Ident) -> Option<Ident> {
		if let Some(ident) = self.scope.borrow().lookup(name) {
			return Some(ident);
		}
		if let Some(ident) = self.parent.as_ref()?.request(name) {
			return Some(self.scope.borrow_mut().add_upvalue(name, ident));
		}
		None
	}
}

fn lookup(scope: &BuildContext<'_>, loc: &ast::Location) -> Location {
	match loc {
		ast::Location::Variable { name } => {
			if let Some(id) = scope.request(name) {
				return Location::Variable { id };
			};
			let env = scope
				.request(&"_ENV".into())
				.expect("_ENV must be always available");
			Location::Field {
				table: Box::new(Location::Variable { id: env }),
				index: Box::new(Expression::Constant(ast::Constant::String(name.clone()))),
			}
		}
		ast::Location::Field { table, index } => {
			let table = Box::new(lookup(scope, &table));
			let index = Box::new(build_expression(scope, &index));
			Location::Field { table, index }
		}
	}
}

fn build_expression(scope: &BuildContext<'_>, code: &ast::Expression) -> Expression {
	match code {
		ast::Expression::Constant(c) => Expression::Constant(c.clone()),
		ast::Expression::Variable(loc) => Expression::Variable(lookup(scope, loc)),
		ast::Expression::Call(call) => Expression::Call(build_call(scope, call)),
		ast::Expression::Function(f) => Expression::Function(build_function(scope, f)),
	}
}

fn build_call(scope: &BuildContext<'_>, code: &ast::Call) -> Call {
	Call {
		callee: Box::new(build_expression(
			&scope,
			&ast::Expression::Variable(ast::Location::Variable {
				name: code.callee.clone(),
			}),
		)),
		args: code
			.args
			.iter()
			.map(|arg| build_expression(&scope, arg))
			.collect(),
	}
}

fn build_function(parent: &BuildContext<'_>, code: &ast::Function) -> Function {
	let scope = BuildContext {
		parent: Some(parent),
		scope: Default::default(),
	};
	for (index, name) in code.args.iter().enumerate() {
		scope.scope.borrow_mut().add_arg(name, index);
	}
	let mut body = vec![];
	for statement in &code.body {
		match statement {
			ast::Statement::Assign { target, source } => {
				let target = lookup(&scope, target);
				body.push(Statement::Assign {
					target,
					source: Box::new(build_expression(&scope, source)),
				});
			}
			ast::Statement::Call(call) => {
				body.push(Statement::Call(build_call(&scope, call)));
			}
			ast::Statement::Local { name, init } => {
				let init = init.as_ref().map(|init| build_expression(&scope, &init));
				let id = scope.scope.borrow_mut().add_local(&name);
				if let Some(init) = init {
					body.push(Statement::Assign {
						target: Location::Variable { id },
						source: Box::new(init),
					});
				}
			}
		}
	}
	Function {
		name: code.name.clone(),
		args: code.args.clone(),
		locals: vec![],
		upvalues: vec![],
		body,
		ret: Box::new(Expression::Constant(ast::Constant::Nil)),
	}
}

pub fn build(code: &ast::Function) -> Function {
	let mut root = BuildContext {
		parent: None,
		scope: RefCell::new(Scope::new_toplevel()),
	};
	build_function(&mut root, &code)
}