Refactor Luau codegen

codegen-luau/src/backend/expression.rs

@@ -0,0 +1,147 @@
+use std::io::{Result, Write};
+
+use wasm_ast::node::{
+	BinOp, CmpOp, Expression, GetGlobal, GetLocal, LoadAt, MemoryGrow, MemorySize, Recall, Select,
+	UnOp, Value,
+};
+
+use super::manager::{write_f32, write_f64, write_separated, write_variable, Driver, Manager};
+
+impl Driver for Recall {
+	fn write(&self, _: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "reg_{} ", self.var)
+	}
+}
+
+impl Driver for Select {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "(")?;
+		self.cond.write(mng, w)?;
+		write!(w, "~= 0 and ")?;
+		self.a.write(mng, w)?;
+		write!(w, "or ")?;
+		self.b.write(mng, w)?;
+		write!(w, ")")
+	}
+}
+
+impl Driver for GetLocal {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write_variable(self.var, mng, w)
+	}
+}
+
+impl Driver for GetGlobal {
+	fn write(&self, _: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "GLOBAL_LIST[{}].value ", self.var)
+	}
+}
+
+impl Driver for LoadAt {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "load_{}(memory_at_0, ", self.what.as_name())?;
+		self.pointer.write(mng, w)?;
+		write!(w, "+ {})", self.offset)
+	}
+}
+
+impl Driver for MemorySize {
+	fn write(&self, _: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "memory_at_{}.min ", self.memory)
+	}
+}
+
+impl Driver for MemoryGrow {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "rt.allocator.grow(memory_at_{}, ", self.memory)?;
+		self.value.write(mng, w)?;
+		write!(w, ")")
+	}
+}
+
+impl Driver for Value {
+	fn write(&self, _: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		match self {
+			Self::I32(i) => write!(w, "{i} "),
+			Self::I64(i) => write!(w, "{i} "),
+			Self::F32(f) => write_f32(*f, w),
+			Self::F64(f) => write_f64(*f, w),
+		}
+	}
+}
+
+impl Driver for UnOp {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		let (a, b) = self.op.as_name();
+
+		write!(w, "{a}_{b}(")?;
+		self.rhs.write(mng, w)?;
+		write!(w, ")")
+	}
+}
+
+fn write_bin_op(bin_op: &BinOp, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+	let op = bin_op.op.as_operator().unwrap();
+
+	write!(w, "(")?;
+	bin_op.lhs.write(mng, w)?;
+	write!(w, "{op} ")?;
+	bin_op.rhs.write(mng, w)?;
+	write!(w, ")")
+}
+
+fn write_bin_op_call(bin_op: &BinOp, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+	let (a, b) = bin_op.op.as_name();
+
+	write!(w, "{a}_{b}(")?;
+	bin_op.lhs.write(mng, w)?;
+	write!(w, ", ")?;
+	bin_op.rhs.write(mng, w)?;
+	write!(w, ")")
+}
+
+impl Driver for BinOp {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		if self.op.as_operator().is_some() {
+			write_bin_op(self, mng, w)
+		} else {
+			write_bin_op_call(self, mng, w)
+		}
+	}
+}
+
+impl Driver for CmpOp {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		let (a, b) = self.op.as_name();
+
+		write!(w, "{a}_{b}(")?;
+		self.lhs.write(mng, w)?;
+		write!(w, ", ")?;
+		self.rhs.write(mng, w)?;
+		write!(w, ")")
+	}
+}
+
+impl Driver for Expression {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		match self {
+			Self::Recall(e) => e.write(mng, w),
+			Self::Select(e) => e.write(mng, w),
+			Self::GetLocal(e) => e.write(mng, w),
+			Self::GetGlobal(e) => e.write(mng, w),
+			Self::LoadAt(e) => e.write(mng, w),
+			Self::MemorySize(e) => e.write(mng, w),
+			Self::MemoryGrow(e) => e.write(mng, w),
+			Self::Value(e) => e.write(mng, w),
+			Self::UnOp(e) => e.write(mng, w),
+			Self::BinOp(e) => e.write(mng, w),
+			Self::CmpOp(e) => e.write(mng, w),
+		}
+	}
+}
+
+impl Driver for &[Expression] {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write_separated(self.iter(), |e, w| e.write(mng, w), w)
+	}
+}

codegen-luau/src/backend/manager.rs

@@ -0,0 +1,89 @@
+use std::{
+	io::{Result, Write},
+	ops::Range,
+};
+
+#[derive(PartialEq, Eq)]
+pub enum Label {
+	Forward,
+	Backward,
+	If,
+}
+
+#[derive(Default)]
+pub struct Manager {
+	label_list: Vec<Label>,
+	pub num_param: usize,
+}
+
+impl Manager {
+	pub fn label_list(&self) -> &[Label] {
+		&self.label_list
+	}
+
+	pub fn push_label(&mut self, label: Label) -> usize {
+		self.label_list.push(label);
+
+		self.label_list.len() - 1
+	}
+
+	pub fn pop_label(&mut self) {
+		self.label_list.pop().unwrap();
+	}
+}
+
+pub trait Driver {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()>;
+}
+
+pub fn write_separated<I, T, M>(mut iter: I, mut func: M, w: &mut dyn Write) -> Result<()>
+where
+	M: FnMut(T, &mut dyn Write) -> Result<()>,
+	I: Iterator<Item = T>,
+{
+	match iter.next() {
+		Some(first) => func(first, w)?,
+		None => return Ok(()),
+	}
+
+	iter.try_for_each(|v| {
+		write!(w, ", ")?;
+		func(v, w)
+	})
+}
+
+pub fn write_ascending(prefix: &str, range: Range<usize>, w: &mut dyn Write) -> Result<()> {
+	write_separated(range, |i, w| write!(w, "{prefix}_{i}"), w)
+}
+
+pub fn write_f32(number: f32, w: &mut dyn Write) -> Result<()> {
+	let sign = if number.is_sign_negative() { "-" } else { "" };
+
+	if number.is_infinite() {
+		write!(w, "{sign}math.huge ")
+	} else if number.is_nan() {
+		write!(w, "{sign}0/0 ")
+	} else {
+		write!(w, "{number:e} ")
+	}
+}
+
+pub fn write_f64(number: f64, w: &mut dyn Write) -> Result<()> {
+	let sign = if number.is_sign_negative() { "-" } else { "" };
+
+	if number.is_infinite() {
+		write!(w, "{sign}math.huge ")
+	} else if number.is_nan() {
+		write!(w, "{sign}0/0 ")
+	} else {
+		write!(w, "{number:e} ")
+	}
+}
+
+pub fn write_variable(var: usize, mng: &Manager, w: &mut dyn Write) -> Result<()> {
+	if let Some(rem) = var.checked_sub(mng.num_param) {
+		write!(w, "loc_{rem} ")
+	} else {
+		write!(w, "param_{var} ")
+	}
+}

codegen-luau/src/backend/mod.rs

@@ -0,0 +1,4 @@
+pub mod manager;
+
+mod expression;
+mod statement;

codegen-luau/src/backend/statement.rs

@@ -0,0 +1,296 @@
+use std::{
+	io::{Result, Write},
+	ops::Range,
+};
+
+use wasm_ast::node::{
+	Backward, Br, BrIf, BrTable, Call, CallIndirect, Else, Forward, If, Intermediate, Memorize,
+	Return, SetGlobal, SetLocal, Statement, StoreAt,
+};
+
+use crate::analyzer::memory;
+
+use super::manager::{write_ascending, write_separated, write_variable, Driver, Label, Manager};
+
+fn br_target(level: usize, in_loop: bool, w: &mut dyn Write) -> Result<()> {
+	write!(w, "if desired then ")?;
+	write!(w, "if desired == {level} then ")?;
+	write!(w, "desired = nil ")?;
+
+	if in_loop {
+		write!(w, "continue ")?;
+	}
+
+	write!(w, "end ")?;
+	write!(w, "break ")?;
+	write!(w, "end ")
+}
+
+fn write_br_gadget(label_list: &[Label], rem: usize, w: &mut dyn Write) -> Result<()> {
+	match label_list.last() {
+		Some(Label::Forward | Label::If) => br_target(rem, false, w),
+		Some(Label::Backward) => br_target(rem, true, w),
+		None => Ok(()),
+	}
+}
+
+impl Driver for Memorize {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "reg_{} = ", self.var)?;
+		self.value.write(mng, w)
+	}
+}
+
+impl Driver for Forward {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		let rem = mng.push_label(Label::Forward);
+
+		write!(w, "while true do ")?;
+
+		self.body.iter().try_for_each(|s| s.write(mng, w))?;
+
+		write!(w, "break ")?;
+		write!(w, "end ")?;
+
+		mng.pop_label();
+		write_br_gadget(mng.label_list(), rem, w)
+	}
+}
+
+impl Driver for Backward {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		let rem = mng.push_label(Label::Backward);
+
+		write!(w, "while true do ")?;
+
+		self.body.iter().try_for_each(|s| s.write(mng, w))?;
+
+		write!(w, "break ")?;
+		write!(w, "end ")?;
+
+		mng.pop_label();
+		write_br_gadget(mng.label_list(), rem, w)
+	}
+}
+
+impl Driver for Else {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "else ")?;
+
+		self.body.iter().try_for_each(|s| s.write(mng, w))
+	}
+}
+
+impl Driver for If {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		let rem = mng.push_label(Label::If);
+
+		write!(w, "while true do ")?;
+		write!(w, "if ")?;
+		self.cond.write(mng, w)?;
+		write!(w, "~= 0 then ")?;
+
+		self.truthy.iter().try_for_each(|s| s.write(mng, w))?;
+
+		if let Some(s) = &self.falsey {
+			s.write(mng, w)?;
+		}
+
+		write!(w, "end ")?;
+		write!(w, "break ")?;
+		write!(w, "end ")?;
+
+		mng.pop_label();
+		write_br_gadget(mng.label_list(), rem, w)
+	}
+}
+
+fn write_br_at(up: usize, mng: &Manager, w: &mut dyn Write) -> Result<()> {
+	write!(w, "do ")?;
+
+	if up == 0 {
+		if let Some(&Label::Backward) = mng.label_list().last() {
+			write!(w, "continue ")?;
+		} else {
+			write!(w, "break ")?;
+		}
+	} else {
+		let level = mng.label_list().len() - 1 - up;
+
+		write!(w, "desired = {level} ")?;
+		write!(w, "break ")?;
+	}
+
+	write!(w, "end ")
+}
+
+impl Driver for Br {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write_br_at(self.target, mng, w)
+	}
+}
+
+impl Driver for BrIf {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "if ")?;
+		self.cond.write(mng, w)?;
+		write!(w, "~= 0 then ")?;
+		write_br_at(self.target, mng, w)?;
+		write!(w, "end ")
+	}
+}
+
+impl Driver for BrTable {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "do ")?;
+		write!(w, "local temp = {{")?;
+
+		if !self.data.table.is_empty() {
+			write!(w, "[0] =")?;
+
+			for d in self.data.table.iter() {
+				write!(w, "{d}, ")?;
+			}
+		}
+
+		write!(w, "}} ")?;
+
+		write!(w, "desired = temp[")?;
+		self.cond.write(mng, w)?;
+		write!(w, "] or {} ", self.data.default)?;
+		write!(w, "break ")?;
+		write!(w, "end ")
+	}
+}
+
+impl Driver for Return {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "do return ")?;
+		self.list.as_slice().write(mng, w)?;
+		write!(w, "end ")
+	}
+}
+
+fn write_call_store(result: Range<usize>, w: &mut dyn Write) -> Result<()> {
+	if result.is_empty() {
+		return Ok(());
+	}
+
+	write_ascending("reg", result, w)?;
+	write!(w, " = ")
+}
+
+impl Driver for Call {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write_call_store(self.result.clone(), w)?;
+
+		write!(w, "FUNC_LIST[{}](", self.func)?;
+		self.param_list.as_slice().write(mng, w)?;
+		write!(w, ")")
+	}
+}
+
+impl Driver for CallIndirect {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write_call_store(self.result.clone(), w)?;
+
+		write!(w, "TABLE_LIST[{}].data[", self.table)?;
+		self.index.write(mng, w)?;
+		write!(w, "](")?;
+		self.param_list.as_slice().write(mng, w)?;
+		write!(w, ")")
+	}
+}
+
+impl Driver for SetLocal {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write_variable(self.var, mng, w)?;
+		write!(w, "= ")?;
+		self.value.write(mng, w)
+	}
+}
+
+impl Driver for SetGlobal {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "GLOBAL_LIST[{}].value = ", self.var)?;
+		self.value.write(mng, w)
+	}
+}
+
+impl Driver for StoreAt {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write!(w, "store_{}(memory_at_0, ", self.what.as_name())?;
+		self.pointer.write(mng, w)?;
+		write!(w, "+ {}, ", self.offset)?;
+		self.value.write(mng, w)?;
+		write!(w, ")")
+	}
+}
+
+impl Driver for Statement {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		match self {
+			Self::Unreachable => write!(w, "error(\"out of code bounds\")"),
+			Self::Memorize(s) => s.write(mng, w),
+			Self::Forward(s) => s.write(mng, w),
+			Self::Backward(s) => s.write(mng, w),
+			Self::If(s) => s.write(mng, w),
+			Self::Br(s) => s.write(mng, w),
+			Self::BrIf(s) => s.write(mng, w),
+			Self::BrTable(s) => s.write(mng, w),
+			Self::Return(s) => s.write(mng, w),
+			Self::Call(s) => s.write(mng, w),
+			Self::CallIndirect(s) => s.write(mng, w),
+			Self::SetLocal(s) => s.write(mng, w),
+			Self::SetGlobal(s) => s.write(mng, w),
+			Self::StoreAt(s) => s.write(mng, w),
+		}
+	}
+}
+
+fn write_parameter_list(ir: &Intermediate, w: &mut dyn Write) -> Result<()> {
+	write!(w, "function(")?;
+	write_ascending("param", 0..ir.num_param, w)?;
+	write!(w, ")")
+}
+
+fn write_variable_list(ir: &Intermediate, w: &mut dyn Write) -> Result<()> {
+	let mut total = 0;
+
+	for data in &ir.local_data {
+		let range = total..total + usize::try_from(data.count()).unwrap();
+		let typed = data.value_type();
+
+		total = range.end;
+
+		write!(w, "local ")?;
+		write_ascending("loc", range.clone(), w)?;
+		write!(w, " = ")?;
+		write_separated(range, |_, w| write!(w, "ZERO_{typed} "), w)?;
+	}
+
+	if ir.num_stack != 0 {
+		write!(w, "local ")?;
+		write_ascending("reg", 0..ir.num_stack, w)?;
+		write!(w, " ")?;
+	}
+
+	Ok(())
+}
+
+impl Driver for Intermediate {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		write_parameter_list(self, w)?;
+
+		for v in memory::visit(self) {
+			write!(w, "local memory_at_{v} = MEMORY_LIST[{v}]")?;
+		}
+
+		write_variable_list(self, w)?;
+
+		mng.num_param = self.num_param;
+		self.code.write(mng, w)?;
+
+		write!(w, "end ")
+	}
+}