Add indentation to Luau output

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} ")?;
+			line!(mng, w, "desired = {level}")?;
-			write!(w, "break ")
+			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 ")?;
+	line!(mng, w, "if not br_map[{id}] then")?;
-	write!(w, "br_map[{id}] = (function() return {{[0] =")?;
+	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()))?;
 
-	write!(w, "}} end)()")?;
+	writeln!(w, "}}")?;
-	write!(w, "end ")?;
+	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<()> {
+fn write_br_parent(mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-	if label_list.iter().all(Option::is_none) {
+	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() {
+	if let Some(last) = *mng.label_list().last().unwrap() {
-		let level = label_list.len() - 1;
+		let level = mng.label_list().len() - 1;
 
-		write!(w, "if desired == {level} then ")?;
+		line!(mng, w, "if desired == {level} then")?;
-		write!(w, "desired = nil ")?;
+		mng.indent();
+		line!(mng, w, "desired = nil")?;
 
-		if *last == LabelType::Backward {
+		if last == LabelType::Backward {
-			write!(w, "continue ")?;
+			line!(mng, w, "continue")?;
 		}
 
-		write!(w, "end ")?;
+		mng.dedent();
+		line!(mng, w, "end")?;
 	}
 
-	write!(w, "break ")?;
+	line!(mng, w, "break")?;
-	write!(w, "end ")
+
+	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::Call(s) => write_stat(s, mng, w),
-			Self::CallIndirect(s) => s.write(w),
+			Self::CallIndirect(s) => write_stat(s, mng, w),
-			Self::SetTemporary(s) => s.write(w),
+			Self::SetTemporary(s) => write_stat(s, mng, w),
-			Self::SetLocal(s) => s.write(w),
+			Self::SetLocal(s) => write_stat(s, mng, w),
-			Self::SetGlobal(s) => s.write(w),
+			Self::SetGlobal(s) => write_stat(s, mng, w),
-			Self::StoreAt(s) => s.write(w),
+			Self::StoreAt(s) => write_stat(s, mng, w),
-			Self::MemoryGrow(s) => s.write(w),
+			Self::MemoryGrow(s) => write_stat(s, mng, w),
 		}
 	}
 }
@@ -296,10 +327,10 @@ 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) {
@@ -312,17 +343,17 @@ fn write_variable_list(ast: &FuncData, w: &mut dyn Write) -> Result<()> {
 
 		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")
 	}
 }

codegen/luau/src/translator.rs

@@ -46,7 +46,7 @@ fn write_named_array(name: &str, len: usize, w: &mut dyn Write) -> Result<()> {
 		None => return Ok(()),
 	};
 
-	write!(w, "local {name} = table.create({len})")
+	writeln!(w, "local {name} = table.create({len})")
 }
 
 fn write_constant(init: &InitExpr, type_info: &TypeInfo, w: &mut dyn Write) -> Result<()> {
@@ -56,7 +56,7 @@ fn write_constant(init: &InitExpr, type_info: &TypeInfo, w: &mut dyn Write) -> R
 	if let Some(Statement::SetTemporary(stat)) = func.code().code().last() {
 		stat.value().write(w)
 	} else {
-		write!(w, r#"error("Valueless constant")"#)
+		writeln!(w, r#"error("Valueless constant")"#)
 	}
 }
 
@@ -69,7 +69,8 @@ fn write_import_of(list: &[Import], wanted: External, w: &mut dyn Write) -> Resu
 		.filter(|v| External::from(v.ty) == wanted)
 		.enumerate()
 	{
-		write!(w, r#"{upper}[{i}] = wasm["{module}"].{lower}["{name}"]"#)?;
+		write!(w, "\t")?;
+		writeln!(w, r#"{upper}[{i}] = wasm["{module}"].{lower}["{name}"]"#)?;
 	}
 
 	Ok(())
@@ -79,13 +80,14 @@ fn write_export_of(list: &[Export], wanted: External, w: &mut dyn Write) -> Resu
 	let lower = wanted.as_ie_name();
 	let upper = lower.to_uppercase();
 
-	write!(w, "{lower} = {{")?;
+	writeln!(w, "\t\t{lower} = {{")?;
 
 	for Export { name, index, .. } in list.iter().filter(|v| External::from(v.kind) == wanted) {
-		write!(w, r#"["{name}"] = {upper}[{index}],"#)?;
+		write!(w, "\t\t\t")?;
+		writeln!(w, r#"["{name}"] = {upper}[{index}],"#)?;
 	}
 
-	write!(w, "}},")
+	writeln!(w, "\t\t}},")
 }
 
 fn write_import_list(list: &[Import], w: &mut dyn Write) -> Result<()> {
@@ -111,9 +113,9 @@ fn write_table_list(wasm: &Module, w: &mut dyn Write) -> Result<()> {
 		let min = table.initial;
 		let max = table.maximum.unwrap_or(0xFFFF);
 
-		write!(
+		writeln!(
 			w,
-			"TABLE_LIST[{index}] = {{ min = {min}, max = {max}, data = {{}} }}"
+			"\tTABLE_LIST[{index}] = {{ min = {min}, max = {max}, data = {{}} }}"
 		)?;
 	}
 
@@ -129,7 +131,7 @@ fn write_memory_list(wasm: &Module, w: &mut dyn Write) -> Result<()> {
 		let min = ty.initial;
 		let max = ty.maximum.unwrap_or(0xFFFF);
 
-		write!(w, "MEMORY_LIST[{index}] = rt.allocator.new({min}, {max})")?;
+		writeln!(w, "\tMEMORY_LIST[{index}] = rt.allocator.new({min}, {max})")?;
 	}
 
 	Ok(())
@@ -142,9 +144,9 @@ fn write_global_list(wasm: &Module, type_info: &TypeInfo, w: &mut dyn Write) ->
 	for (i, global) in global.iter().enumerate() {
 		let index = offset + i;
 
-		write!(w, "GLOBAL_LIST[{index}] = {{ value =")?;
+		write!(w, "\tGLOBAL_LIST[{index}] = {{ value = ")?;
 		write_constant(&global.init_expr, type_info, w)?;
-		write!(w, "}}")?;
+		writeln!(w, " }}")?;
 	}
 
 	Ok(())
@@ -160,13 +162,14 @@ fn write_element_list(list: &[Element], type_info: &TypeInfo, w: &mut dyn Write)
 			_ => unimplemented!(),
 		};
 
-		write!(w, "do ")?;
+		writeln!(w, "\tdo")?;
-		write!(w, "local target = TABLE_LIST[{index}].data ")?;
+		writeln!(w, "\t\tlocal target = TABLE_LIST[{index}].data")?;
-		write!(w, "local offset =")?;
+		write!(w, "\t\tlocal offset = ")?;
 
 		write_constant(&init, type_info, w)?;
 
-		write!(w, "local data = {{")?;
+		writeln!(w)?;
+		write!(w, "\t\tlocal data = {{ ")?;
 
 		for item in element.items.get_items_reader().unwrap() {
 			match item.unwrap() {
@@ -175,11 +178,9 @@ fn write_element_list(list: &[Element], type_info: &TypeInfo, w: &mut dyn Write)
 			}?;
 		}
 
-		write!(w, "}}")?;
+		writeln!(w, " }}")?;
-
+		writeln!(w, "\t\ttable.move(data, 1, #data, offset, target)")?;
-		write!(w, "table.move(data, 1, #data, offset, target)")?;
+		writeln!(w, "\tend")?;
-
-		write!(w, "end ")?;
 	}
 
 	Ok(())
@@ -195,10 +196,9 @@ fn write_data_list(list: &[Data], type_info: &TypeInfo, w: &mut dyn Write) -> Re
 			} => (memory_index, init_expr),
 		};
 
-		write!(w, "rt.store.string(")?;
+		write!(w, "\trt.store.string(MEMORY_LIST[{index}], ")?;
-		write!(w, "MEMORY_LIST[{index}],")?;
 		write_constant(&init, type_info, w)?;
-		write!(w, r#","{}")"#, data.data.escape_ascii())?;
+		writeln!(w, r#","{}")"#, data.data.escape_ascii())?;
 	}
 
 	Ok(())
@@ -224,7 +224,9 @@ fn write_local_operation(head: &str, tail: &str, w: &mut dyn Write) -> Result<()
 		("clz", "i32") => write!(w, "bit32.countlz"),
 		("ctz", "i32") => write!(w, "bit32.countrz"),
 		_ => write!(w, "rt.{head}.{tail}"),
-	}
+	}?;
+
+	writeln!(w)
 }
 
 fn write_localize_used(
@@ -258,7 +260,7 @@ fn write_localize_used(
 	}
 
 	for mem in &mem_set {
-		write!(w, "local memory_at_{mem} ")?;
+		writeln!(w, "local memory_at_{mem}")?;
 	}
 
 	Ok(mem_set)
@@ -291,29 +293,30 @@ fn write_module_start(
 	mem_set: &BTreeSet<usize>,
 	w: &mut dyn Write,
 ) -> Result<()> {
-	write!(w, "local function run_init_code()")?;
+	writeln!(w, "local function run_init_code()")?;
 	write_table_list(wasm, w)?;
 	write_memory_list(wasm, w)?;
 	write_global_list(wasm, type_info, w)?;
 	write_element_list(wasm.element_section(), type_info, w)?;
 	write_data_list(wasm.data_section(), type_info, w)?;
-	write!(w, "end ")?;
+	writeln!(w, "end")?;
 
-	write!(w, "return function(wasm)")?;
+	writeln!(w, "return function(wasm)")?;
 	write_import_list(wasm.import_section(), w)?;
-	write!(w, "run_init_code()")?;
+	writeln!(w, "\trun_init_code()")?;
 
 	for mem in mem_set {
-		write!(w, "memory_at_{mem} = MEMORY_LIST[{mem}]")?;
+		writeln!(w, "\tmemory_at_{mem} = MEMORY_LIST[{mem}]")?;
 	}
 
 	if let Some(start) = wasm.start_section() {
-		write!(w, "FUNC_LIST[{start}]()")?;
+		writeln!(w, "\tFUNC_LIST[{start}]()")?;
 	}
 
-	write!(w, "return {{")?;
+	writeln!(w, "\treturn {{")?;
 	write_export_list(wasm.export_section(), w)?;
-	write!(w, "}} end ")
+	writeln!(w, "\t}}")?;
+	writeln!(w, "end")
 }
 
 /// # Errors

dev-test/tests/luau_translate.rs

@@ -105,7 +105,7 @@ impl Target for Luau {
 				};
 				let data = Module::from_data(&bytes);
 
-				write!(w, "assert_trap((function() ")?;
+				writeln!(w, "assert_trap((function()")?;
 				codegen_luau::from_module_untyped(&data, w)?;
 				writeln!(w, "end)(), linked)")
 			}
@@ -154,10 +154,10 @@ impl Target for Luau {
 		let numeric = codegen_luau::NUMERIC;
 
 		writeln!(w, "local Integer = (function()")?;
-		writeln!(w, "{numeric}")?;
+		write!(w, "{numeric}")?;
 		writeln!(w, "end)()")?;
 		writeln!(w, "local rt = (function()")?;
-		writeln!(w, "{runtime}")?;
+		write!(w, "{runtime}")?;
 		writeln!(w, "end)()")?;
 
 		writeln!(w, "{ASSERTION}")
@@ -166,7 +166,7 @@ impl Target for Luau {
 	fn write_module(data: &Module, name: Option<&str>, w: &mut dyn Write) -> Result<()> {
 		let type_info = TypeInfo::from_module(data);
 
-		write!(w, r#"loaded["temp"] = (function() "#)?;
+		writeln!(w, r#"loaded["temp"] = (function()"#)?;
 		codegen_luau::from_module_typed(data, &type_info, w)?;
 		writeln!(w, "end)()(linked)")?;