Add indentation to Luau output

codegen/luau/src/backend/expression.rs

@@ -16,11 +16,11 @@ macro_rules! impl_write_number {
 	($name:tt, $numeric:ty) => {
 		fn $name(number: $numeric, w: &mut dyn Write) -> Result<()> {
 			match (number.classify(), number.is_sign_negative()) {
-				(FpCategory::Nan, true) => write!(w, "(0.0 / 0.0) "),
-				(FpCategory::Nan, false) => write!(w, "-(0.0 / 0.0) "),
-				(FpCategory::Infinite, true) => write!(w, "-math.huge "),
-				(FpCategory::Infinite, false) => write!(w, "math.huge "),
-				_ => write!(w, "{number:e} "),
+				(FpCategory::Nan, true) => write!(w, "(0.0 / 0.0)"),
+				(FpCategory::Nan, false) => write!(w, "-(0.0 / 0.0)"),
+				(FpCategory::Infinite, true) => write!(w, "-math.huge"),
+				(FpCategory::Infinite, false) => write!(w, "math.huge"),
+				_ => write!(w, "{number:e}"),
 			}
 		}
 	};
@@ -30,9 +30,9 @@ impl DriverNoContext for Select {
 	fn write(&self, w: &mut dyn Write) -> Result<()> {
 		write!(w, "(")?;
 		Condition(self.condition()).write(w)?;
-		write!(w, "and ")?;
+		write!(w, " and ")?;
 		self.on_true().write(w)?;
-		write!(w, "or ")?;
+		write!(w, " or ")?;
 		self.on_false().write(w)?;
 		write!(w, ")")
 	}
@@ -40,19 +40,19 @@ impl DriverNoContext for Select {
 
 impl DriverNoContext for GetTemporary {
 	fn write(&self, w: &mut dyn Write) -> Result<()> {
-		write!(w, "reg_{} ", self.var())
+		write!(w, "reg_{}", self.var())
 	}
 }
 
 impl DriverNoContext for GetLocal {
 	fn write(&self, w: &mut dyn Write) -> Result<()> {
-		write!(w, "loc_{} ", self.var())
+		write!(w, "loc_{}", self.var())
 	}
 }
 
 impl DriverNoContext for GetGlobal {
 	fn write(&self, w: &mut dyn Write) -> Result<()> {
-		write!(w, "GLOBAL_LIST[{}].value ", self.var())
+		write!(w, "GLOBAL_LIST[{}].value", self.var())
 	}
 }
 
@@ -62,7 +62,7 @@ impl DriverNoContext for LoadAt {
 		self.pointer().write(w)?;
 
 		if self.offset() != 0 {
-			write!(w, "+ {}", self.offset())?;
+			write!(w, " + {}", self.offset())?;
 		}
 
 		write!(w, ")")
@@ -71,26 +71,26 @@ impl DriverNoContext for LoadAt {
 
 impl DriverNoContext for MemorySize {
 	fn write(&self, w: &mut dyn Write) -> Result<()> {
-		write!(w, "memory_at_{}.min ", self.memory())
+		write!(w, "memory_at_{}.min", self.memory())
 	}
 }
 
 pub fn write_i32(number: i32, w: &mut dyn Write) -> Result<()> {
 	let list = number.to_ne_bytes();
 
-	write!(w, "{} ", u32::from_ne_bytes(list))
+	write!(w, "{}", u32::from_ne_bytes(list))
 }
 
 fn write_i64(number: i64, w: &mut dyn Write) -> Result<()> {
 	match number {
-		0 => write!(w, "i64_ZERO "),
-		1 => write!(w, "i64_ONE "),
+		0 => write!(w, "i64_ZERO"),
+		1 => write!(w, "i64_ONE"),
 		_ => {
 			let list = number.to_ne_bytes();
 			let a = u32::from_ne_bytes(list[0..4].try_into().unwrap());
 			let b = u32::from_ne_bytes(list[4..8].try_into().unwrap());
 
-			write!(w, "i64_from_u32({a}, {b}) ")
+			write!(w, "i64_from_u32({a}, {b})")
 		}
 	}
 }
@@ -124,7 +124,7 @@ impl DriverNoContext for BinOp {
 		if let Some(symbol) = self.op_type().as_symbol() {
 			write!(w, "(")?;
 			self.lhs().write(w)?;
-			write!(w, "{symbol} ")?;
+			write!(w, " {symbol} ")?;
 			self.rhs().write(w)?;
 			write!(w, ")")
 		} else {
@@ -147,7 +147,7 @@ impl DriverNoContext for CmpOpBoolean<'_> {
 
 		if let Some(symbol) = cmp.op_type().as_symbol() {
 			cmp.lhs().write(w)?;
-			write!(w, "{symbol} ")?;
+			write!(w, " {symbol} ")?;
 			cmp.rhs().write(w)
 		} else {
 			let (head, tail) = cmp.op_type().as_name();
@@ -165,7 +165,7 @@ impl DriverNoContext for CmpOp {
 	fn write(&self, w: &mut dyn Write) -> Result<()> {
 		write!(w, "(")?;
 		CmpOpBoolean(self).write(w)?;
-		write!(w, "and 1 or 0)")
+		write!(w, " and 1 or 0)")
 	}
 }
 
@@ -177,7 +177,7 @@ impl DriverNoContext for Condition<'_> {
 			CmpOpBoolean(node).write(w)
 		} else {
 			self.0.write(w)?;
-			write!(w, "~= 0 ")
+			write!(w, " ~= 0")
 		}
 	}
 }

codegen/luau/src/backend/manager.rs

@@ -6,10 +6,36 @@ use std::{
 
 use wasm_ast::node::{BrTable, LabelType};
 
+#[macro_export]
+macro_rules! indentation {
+	($mng:tt, $w:tt) => {{
+		let mut iter = 0..$mng.indentation();
+
+		iter.try_for_each(|_| write!($w, "\t"))
+	}};
+}
+
+#[macro_export]
+macro_rules! indented {
+	($mng:tt, $w:tt, $($args:tt)*) => {{
+		indentation!($mng, $w)?;
+		write!($w, $($args)*)
+	}};
+}
+
+#[macro_export]
+macro_rules! line {
+	($mng:tt, $w:tt, $($args:tt)*) => {{
+		indentation!($mng, $w)?;
+		writeln!($w, $($args)*)
+	}};
+}
+
 #[derive(Default)]
 pub struct Manager {
 	table_map: HashMap<usize, usize>,
 	label_list: Vec<Option<LabelType>>,
+	indentation: usize,
 }
 
 impl Manager {
@@ -34,6 +60,18 @@ impl Manager {
 	pub fn pop_label(&mut self) {
 		self.label_list.pop().unwrap();
 	}
+
+	pub fn indentation(&self) -> usize {
+		self.indentation
+	}
+
+	pub fn indent(&mut self) {
+		self.indentation += 1;
+	}
+
+	pub fn dedent(&mut self) {
+		self.indentation -= 1;
+	}
 }
 
 pub trait Driver {

codegen/luau/src/backend/statement.rs

@@ -9,7 +9,7 @@ use wasm_ast::node::{
 };
 use wasmparser::ValType;
 
-use crate::analyzer::br_target;
+use crate::{analyzer::br_target, indentation, indented, line};
 
 use super::{
 	expression::Condition,
@@ -19,23 +19,24 @@ use super::{
 impl Driver for Br {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		if !self.align().is_aligned() {
+			indentation!(mng, w)?;
 			write_ascending("reg", self.align().new_range(), w)?;
 			write!(w, " = ")?;
 			write_ascending("reg", self.align().old_range(), w)?;
-			write!(w, " ")?;
+			writeln!(w)?;
 		}
 
 		if self.target() == 0 {
 			if let Some(Some(LabelType::Backward)) = mng.label_list().last() {
-				write!(w, "continue ")
+				line!(mng, w, "continue")
 			} else {
-				write!(w, "break ")
+				line!(mng, w, "break")
 			}
 		} else {
 			let level = mng.label_list().len() - 1 - self.target();
 
-			write!(w, "desired = {level} ")?;
-			write!(w, "break ")
+			line!(mng, w, "desired = {level}")?;
+			line!(mng, w, "break")
 		}
 	}
 }
@@ -62,39 +63,51 @@ fn write_search_layer(
 	let br = list[center];
 
 	if range.start != center {
-		write!(w, "if temp < {} then ", br.target())?;
+		line!(mng, w, "if temp < {} then", br.target())?;
+		mng.indent();
 		write_search_layer(range.start..center, list, mng, w)?;
-		write!(w, "else")?;
+		mng.dedent();
+		indented!(mng, w, "else")?;
 	}
 
 	if range.end != center + 1 {
-		write!(w, "if temp > {} then ", br.target())?;
+		writeln!(w, "if temp > {} then", br.target())?;
+		mng.indent();
 		write_search_layer(center + 1..range.end, list, mng, w)?;
-		write!(w, "else")?;
+		mng.dedent();
+		indented!(mng, w, "else")?;
 	}
 
-	write!(w, " ")?;
+	writeln!(w)?;
+	mng.indent();
 	br.write(mng, w)?;
-	write!(w, "end ")
+	mng.dedent();
+	line!(mng, w, "end")
 }
 
 fn write_table_setup(table: &BrTable, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 	let id = mng.get_table_index(table);
 
-	write!(w, "if not br_map[{id}] then ")?;
-	write!(w, "br_map[{id}] = (function() return {{[0] =")?;
+	line!(mng, w, "if not br_map[{id}] then")?;
+	mng.indent();
+	line!(mng, w, "br_map[{id}] = (function()")?;
+	mng.indent();
+	indented!(mng, w, "return {{ [0] = ")?;
 
 	table
 		.data()
 		.iter()
-		.try_for_each(|v| write!(w, "{},", v.target()))?;
+		.try_for_each(|v| write!(w, "{}, ", v.target()))?;
 
-	write!(w, "}} end)()")?;
-	write!(w, "end ")?;
+	writeln!(w, "}}")?;
+	mng.dedent();
+	line!(mng, w, "end)()")?;
+	mng.dedent();
+	line!(mng, w, "end")?;
 
-	write!(w, "local temp = br_map[{id}][")?;
+	indented!(mng, w, "temp = br_map[{id}][")?;
 	table.condition().write(w)?;
-	write!(w, "] or {} ", table.default().target())
+	writeln!(w, "] or {}", table.default().target())
 }
 
 impl Driver for BrTable {
@@ -119,82 +132,94 @@ impl Driver for BrTable {
 impl Driver for Terminator {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		match self {
-			Self::Unreachable => write!(w, "error(\"out of code bounds\")"),
+			Self::Unreachable => line!(mng, w, r#"error("out of code bounds")"#),
 			Self::Br(s) => s.write(mng, w),
 			Self::BrTable(s) => s.write(mng, w),
 		}
 	}
 }
 
-fn write_br_parent(label_list: &[Option<LabelType>], w: &mut dyn Write) -> Result<()> {
-	if label_list.iter().all(Option::is_none) {
+fn write_br_parent(mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+	if mng.label_list().iter().all(Option::is_none) {
 		return Ok(());
 	}
 
-	write!(w, "if desired then ")?;
+	line!(mng, w, "if desired then")?;
+	mng.indent();
 
-	if let Some(last) = label_list.last().unwrap() {
-		let level = label_list.len() - 1;
+	if let Some(last) = *mng.label_list().last().unwrap() {
+		let level = mng.label_list().len() - 1;
 
-		write!(w, "if desired == {level} then ")?;
-		write!(w, "desired = nil ")?;
+		line!(mng, w, "if desired == {level} then")?;
+		mng.indent();
+		line!(mng, w, "desired = nil")?;
 
-		if *last == LabelType::Backward {
-			write!(w, "continue ")?;
+		if last == LabelType::Backward {
+			line!(mng, w, "continue")?;
 		}
 
-		write!(w, "end ")?;
+		mng.dedent();
+		line!(mng, w, "end")?;
 	}
 
-	write!(w, "break ")?;
-	write!(w, "end ")
+	line!(mng, w, "break")?;
+
+	mng.dedent();
+	line!(mng, w, "end")
 }
 
 impl Driver for Block {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		mng.push_label(self.label_type());
 
-		write!(w, "while true do ")?;
+		line!(mng, w, "while true do")?;
+		mng.indent();
 
 		self.code().iter().try_for_each(|s| s.write(mng, w))?;
 
 		match self.last() {
 			Some(v) => v.write(mng, w)?,
-			None => write!(w, "break ")?,
+			None => line!(mng, w, "break")?,
 		}
 
-		write!(w, "end ")?;
+		mng.dedent();
+		line!(mng, w, "end")?;
 
 		mng.pop_label();
-		write_br_parent(mng.label_list(), w)
+		write_br_parent(mng, w)
 	}
 }
 
 impl Driver for BrIf {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		write!(w, "if ")?;
+		indented!(mng, w, "if ")?;
 		Condition(self.condition()).write(w)?;
-		write!(w, "then ")?;
+		writeln!(w, " then")?;
+		mng.indent();
 		self.target().write(mng, w)?;
-		write!(w, "end ")
+		mng.dedent();
+		line!(mng, w, "end")
 	}
 }
 
 impl Driver for If {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		write!(w, "if ")?;
+		indented!(mng, w, "if ")?;
 		Condition(self.condition()).write(w)?;
-		write!(w, "then ")?;
+		writeln!(w, " then")?;
 
+		mng.indent();
 		self.on_true().write(mng, w)?;
+		mng.dedent();
 
 		if let Some(v) = self.on_false() {
-			write!(w, "else ")?;
-
+			line!(mng, w, "else")?;
+			mng.indent();
 			v.write(mng, w)?;
+			mng.dedent();
 		}
 
-		write!(w, "end ")
+		line!(mng, w, "end")
 	}
 }
 
@@ -276,19 +301,25 @@ impl DriverNoContext for MemoryGrow {
 	}
 }
 
+fn write_stat(stat: &dyn DriverNoContext, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+	indentation!(mng, w)?;
+	stat.write(w)?;
+	writeln!(w)
+}
+
 impl Driver for Statement {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		match self {
 			Self::Block(s) => s.write(mng, w),
 			Self::BrIf(s) => s.write(mng, w),
 			Self::If(s) => s.write(mng, w),
-			Self::Call(s) => s.write(w),
-			Self::CallIndirect(s) => s.write(w),
-			Self::SetTemporary(s) => s.write(w),
-			Self::SetLocal(s) => s.write(w),
-			Self::SetGlobal(s) => s.write(w),
-			Self::StoreAt(s) => s.write(w),
-			Self::MemoryGrow(s) => s.write(w),
+			Self::Call(s) => write_stat(s, mng, w),
+			Self::CallIndirect(s) => write_stat(s, mng, w),
+			Self::SetTemporary(s) => write_stat(s, mng, w),
+			Self::SetLocal(s) => write_stat(s, mng, w),
+			Self::SetGlobal(s) => write_stat(s, mng, w),
+			Self::StoreAt(s) => write_stat(s, mng, w),
+			Self::MemoryGrow(s) => write_stat(s, mng, w),
 		}
 	}
 }
@@ -296,33 +327,33 @@ impl Driver for Statement {
 fn write_parameter_list(ast: &FuncData, w: &mut dyn Write) -> Result<()> {
 	write!(w, "function(")?;
 	write_ascending("loc", 0..ast.num_param(), w)?;
-	write!(w, ")")
+	writeln!(w, ")")
 }
 
-fn write_variable_list(ast: &FuncData, w: &mut dyn Write) -> Result<()> {
+fn write_variable_list(ast: &FuncData, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 	let mut total = ast.num_param();
 
 	for data in ast.local_data().iter().filter(|v| v.0 != 0) {
 		let range = total..total + usize::try_from(data.0).unwrap();
 		let zero = if data.1 == ValType::I64 {
-			"i64_ZERO "
+			"i64_ZERO"
 		} else {
-			"0 "
+			"0"
 		};
 
 		total = range.end;
 
-		write!(w, "local ")?;
+		indented!(mng, w, "local ")?;
 		write_ascending("loc", range.clone(), w)?;
 		write!(w, " = ")?;
 		write_separated(range, |_, w| w.write_all(zero.as_bytes()), w)?;
-		write!(w, " ")?;
+		writeln!(w)?;
 	}
 
 	if ast.num_stack() != 0 {
-		write!(w, "local ")?;
+		indented!(mng, w, "local ")?;
 		write_ascending("reg", 0..ast.num_stack(), w)?;
-		write!(w, " ")?;
+		writeln!(w)?;
 	}
 
 	Ok(())
@@ -332,26 +363,30 @@ impl Driver for FuncData {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		let br_data = br_target::visit(self);
 
+		mng.indent();
+
 		write_parameter_list(self, w)?;
-		write_variable_list(self, w)?;
+		write_variable_list(self, mng, w)?;
 
 		if br_data.1 {
-			write!(w, "local desired ")?;
+			line!(mng, w, "local desired")?;
 		}
 
 		if !br_data.0.is_empty() {
-			write!(w, "local br_map = {{}} ")?;
+			line!(mng, w, "local br_map = {{}}")?;
 		}
 
 		mng.set_table_map(br_data.0);
 		self.code().write(mng, w)?;
 
 		if self.num_result() != 0 {
-			write!(w, "return ")?;
+			indented!(mng, w, "return ")?;
 			write_ascending("reg", 0..self.num_result(), w)?;
-			write!(w, " ")?;
+			writeln!(w)?;
 		}
 
-		write!(w, "end ")
+		mng.dedent();
+
+		line!(mng, w, "end")
 	}
 }