Reduce code duplication and add function names

wasm/src/ast/builder.rs

@@ -1,5 +1,5 @@
 use parity_wasm::elements::{
-	BlockType, External, FunctionType, ImportEntry, Instruction, Local, Module, Type,
+	BlockType, External, FuncBody, FunctionType, ImportEntry, Instruction, Local, Module, Type,
 };
 
 use super::{
@@ -109,10 +109,6 @@ pub struct Builder<'a> {
 	last_stack: usize,
 }
 
-fn local_sum(list: &[Local]) -> u32 {
-	list.iter().map(Local::count).sum()
-}
-
 fn is_else_stat(inst: &Instruction) -> bool {
 	inst == &Instruction::Else
 }
@@ -124,6 +120,14 @@ fn is_dead_precursor(inst: &Instruction) -> bool {
 	)
 }
 
+fn local_sum(body: &FuncBody) -> u32 {
+	body.locals().iter().map(Local::count).sum()
+}
+
+fn load_func_at(wasm: &Module, index: usize) -> &FuncBody {
+	&wasm.code_section().unwrap().bodies()[index]
+}
+
 impl<'a> Builder<'a> {
 	pub fn new(wasm: &'a Module, other: &'a Arities) -> Builder<'a> {
 		Builder {
@@ -136,18 +140,22 @@ impl<'a> Builder<'a> {
 		}
 	}
 
-	pub fn consume(mut self, name: usize) -> Function {
+	pub fn consume(mut self, index: usize) -> Function {
-		let arity = &self.other.in_arity[name];
+		let func = load_func_at(self.wasm, index);
+		let arity = &self.other.in_arity[index];
+
+		let num_param = arity.num_param;
+		let num_local = local_sum(func);
 
 		self.num_result = arity.num_result;
 
-		let func = &self.wasm.code_section().unwrap().bodies()[name];
 		let body = self.new_forward(&mut func.code().elements());
+		let num_stack = self.last_stack.try_into().unwrap();
 
 		Function {
-			num_param: arity.num_param,
+			num_param,
-			num_local: local_sum(func.locals()),
+			num_local,
-			num_stack: u32::try_from(self.last_stack).unwrap(),
+			num_stack,
 			body,
 		}
 	}
@@ -331,11 +339,13 @@ impl<'a> Builder<'a> {
 		}
 	}
 
-	fn new_body(&mut self, list: &mut &[Instruction]) -> Vec<Statement> {
+	fn new_stored_body(&mut self, list: &mut &[Instruction]) -> Vec<Statement> {
 		use Instruction as Inst;
 
 		let mut stat = Vec::new();
 
+		self.save_pending();
+
 		loop {
 			let inst = &list[0];
 
@@ -509,17 +519,9 @@ impl<'a> Builder<'a> {
 			}
 		}
 
-		stat
-	}
-
-	fn new_stored_body(&mut self, list: &mut &[Instruction]) -> Vec<Statement> {
-		self.save_pending();
-
-		let body = self.new_body(list);
-
 		self.load_pending();
 
-		body
+		stat
 	}
 
 	fn new_else(&mut self, list: &mut &[Instruction]) -> Else {

wasm/src/main.rs

@@ -13,6 +13,15 @@ fn lang_from_string<'a>(name: &str, wasm: &'a Module) -> Box<dyn Transpiler<'a>
 	}
 }
 
+fn parse_module(name: &str) -> Module {
+	let wasm = deserialize_file(name).expect("Failed to parse Wasm file");
+
+	match wasm.parse_names() {
+		Ok(n) => n,
+		Err(n) => n.1,
+	}
+}
+
 fn main() {
 	let mut args = std::env::args().skip(1);
 	let name = args.next().expect("No language specified");
@@ -20,7 +29,7 @@ fn main() {
 	let output = std::io::stdout();
 
 	for v in args {
-		let wasm = deserialize_file(v).unwrap();
+		let wasm = parse_module(&v);
 		let module = lang_from_string(&name, &wasm);
 
 		module.transpile(&mut output.lock()).unwrap();

wasm/src/writer/luajit.rs

@@ -12,13 +12,12 @@ use crate::{
 			MemorySize, Recall, Return, Select, SetGlobal, SetLocal, Statement, Value,
 		},
 	},
-	writer::shared::write_result_list,
 };
 
 use super::{
 	shared::{
-		aux_internal_index, write_f32, write_f64, write_in_order, write_memory_init,
+		aux_internal_index, write_f32, write_f64, write_func_name, write_memory_init,
-		write_nil_array, write_table_init, write_variable_list,
+		write_parameter_list, write_result_list, write_table_init, write_variable_list,
 	},
 	visit::{Transpiler, Writer},
 };
@@ -448,11 +447,7 @@ impl Driver for Statement {
 
 impl Driver for Function {
 	fn visit(&self, v: &mut Visitor, w: Writer) -> Result<()> {
-		write!(w, "function(")?;
+		write_parameter_list(self, w)?;
-		write_in_order("param", self.num_param, w)?;
-		write!(w, ")")?;
-
-		v.num_param = self.num_param;
 
 		for v in memory::visit(self) {
 			write!(w, "local memory_at_{0} = MEMORY_LIST[{0}]", v)?;
@@ -461,6 +456,7 @@ impl Driver for Function {
 		write!(w, "local temp ")?;
 		write_variable_list(self, w)?;
 
+		v.num_param = self.num_param;
 		self.body.visit(v, w)?;
 
 		write!(w, "end ")
@@ -688,16 +684,24 @@ impl<'a> LuaJIT<'a> {
 	/// # Errors
 	/// Returns `Err` if writing to `Writer` failed.
 	pub fn gen_func_list(&self, func_list: &[Function], w: Writer) -> Result<()> {
-		let offset = self.arity.len_ex();
+		let o = self.arity.len_ex();
 
 		func_list.iter().enumerate().try_for_each(|(i, v)| {
-			write!(w, "FUNC_LIST[{}] =", i + offset)?;
+			write_func_name(self.wasm, i as u32, o as u32, w)?;
 
 			v.visit(&mut Visitor::default(), w)
 		})
 	}
 }
 
+fn write_list(name: &str, len: usize, w: Writer) -> Result<()> {
+	if len == 0 {
+		return Ok(());
+	}
+
+	write!(w, "local {} = {{[0] = {}}}", name, "nil, ".repeat(len))
+}
+
 impl<'a> Transpiler<'a> for LuaJIT<'a> {
 	fn new(wasm: &'a Module) -> Self {
 		let arity = Arities::new(wasm);
@@ -712,10 +716,10 @@ impl<'a> Transpiler<'a> for LuaJIT<'a> {
 
 		Self::gen_localize(&func_list, w)?;
 
-		write_nil_array("FUNC_LIST", self.wasm.functions_space(), w)?;
+		write_list("FUNC_LIST", self.wasm.functions_space(), w)?;
-		write_nil_array("TABLE_LIST", self.wasm.table_space(), w)?;
+		write_list("TABLE_LIST", self.wasm.table_space(), w)?;
-		write_nil_array("MEMORY_LIST", self.wasm.memory_space(), w)?;
+		write_list("MEMORY_LIST", self.wasm.memory_space(), w)?;
-		write_nil_array("GLOBAL_LIST", self.wasm.globals_space(), w)?;
+		write_list("GLOBAL_LIST", self.wasm.globals_space(), w)?;
 
 		self.gen_func_list(&func_list, w)?;
 		self.gen_start_point(w)

wasm/src/writer/luau.rs

@@ -16,8 +16,8 @@ use crate::{
 
 use super::{
 	shared::{
-		aux_internal_index, write_f32, write_f64, write_in_order, write_memory_init,
+		aux_internal_index, write_f32, write_f64, write_func_name, write_memory_init,
-		write_nil_array, write_result_list, write_table_init, write_variable_list,
+		write_parameter_list, write_result_list, write_table_init, write_variable_list,
 	},
 	visit::{Transpiler, Writer},
 };
@@ -459,11 +459,7 @@ impl Driver for Statement {
 
 impl Driver for Function {
 	fn visit(&self, v: &mut Visitor, w: Writer) -> Result<()> {
-		write!(w, "function(")?;
+		write_parameter_list(self, w)?;
-		write_in_order("param", self.num_param, w)?;
-		write!(w, ")")?;
-
-		v.num_param = self.num_param;
 
 		for v in memory::visit(self) {
 			write!(w, "local memory_at_{0} = MEMORY_LIST[{0}]", v)?;
@@ -471,6 +467,7 @@ impl Driver for Function {
 
 		write_variable_list(self, w)?;
 
+		v.num_param = self.num_param;
 		self.body.visit(v, w)?;
 
 		write!(w, "end ")
@@ -694,16 +691,22 @@ impl<'a> Luau<'a> {
 	}
 
 	fn gen_func_list(&self, func_list: &[Function], w: Writer) -> Result<()> {
-		let offset = self.arity.len_ex();
+		let o = self.arity.len_ex();
 
 		func_list.iter().enumerate().try_for_each(|(i, v)| {
-			write!(w, "FUNC_LIST[{}] =", i + offset)?;
+			write_func_name(self.wasm, i as u32, o as u32, w)?;
 
 			v.visit(&mut Visitor::default(), w)
 		})
 	}
 }
 
+fn write_list(name: &str, len: usize, w: Writer) -> Result<()> {
+	let len = len.saturating_sub(1);
+
+	write!(w, "local {} = table.create({})", name, len)
+}
+
 impl<'a> Transpiler<'a> for Luau<'a> {
 	fn new(wasm: &'a Module) -> Self {
 		let arity = Arities::new(wasm);
@@ -718,10 +721,10 @@ impl<'a> Transpiler<'a> for Luau<'a> {
 
 		Self::gen_localize(&func_list, w)?;
 
-		write_nil_array("FUNC_LIST", self.wasm.functions_space(), w)?;
+		write_list("FUNC_LIST", self.wasm.functions_space(), w)?;
-		write_nil_array("TABLE_LIST", self.wasm.table_space(), w)?;
+		write_list("TABLE_LIST", self.wasm.table_space(), w)?;
-		write_nil_array("MEMORY_LIST", self.wasm.memory_space(), w)?;
+		write_list("MEMORY_LIST", self.wasm.memory_space(), w)?;
-		write_nil_array("GLOBAL_LIST", self.wasm.globals_space(), w)?;
+		write_list("GLOBAL_LIST", self.wasm.globals_space(), w)?;
 
 		self.gen_func_list(&func_list, w)?;
 		self.gen_start_point(w)

wasm/src/writer/shared.rs

@@ -1,6 +1,6 @@
 use std::{io::Result, ops::Range};
 
-use parity_wasm::elements::{Internal, ResizableLimits};
+use parity_wasm::elements::{Internal, Module, ResizableLimits};
 
 use crate::ast::node::Function;
 
@@ -12,7 +12,7 @@ pub fn aux_internal_index(internal: Internal) -> u32 {
 	}
 }
 
-pub fn new_limit_max(limits: &ResizableLimits) -> String {
+fn new_limit_max(limits: &ResizableLimits) -> String {
 	match limits.maximum() {
 		Some(v) => v.to_string(),
 		None => "0xFFFF".to_string(),
@@ -33,12 +33,19 @@ pub fn write_memory_init(limit: &ResizableLimits, w: Writer) -> Result<()> {
 	write!(w, "rt.memory.new({}, {})", a, b)
 }
 
-pub fn write_nil_array(name: &str, len: usize, w: Writer) -> Result<()> {
+pub fn write_func_name(wasm: &Module, index: u32, offset: u32, w: Writer) -> Result<()> {
-	if len == 0 {
+	let opt = wasm
-		return Ok(());
+		.names_section()
+		.and_then(|v| v.functions())
+		.and_then(|v| v.names().get(index));
+
+	write!(w, "FUNC_LIST")?;
+
+	if let Some(name) = opt {
+		write!(w, "--[[{}]]", name)?;
 	}
 
-	write!(w, "local {} = {{[0] = {}}}", name, "nil, ".repeat(len))
+	write!(w, "[{}] =", index + offset)
 }
 
 pub fn write_in_order(prefix: &str, len: u32, w: Writer) -> Result<()> {
@@ -74,6 +81,12 @@ pub fn write_f64(f: f64, w: Writer) -> Result<()> {
 	}
 }
 
+pub fn write_parameter_list(func: &Function, w: Writer) -> Result<()> {
+	write!(w, "function(")?;
+	write_in_order("param", func.num_param, w)?;
+	write!(w, ")")
+}
+
 pub fn write_result_list(range: Range<u32>, w: Writer) -> Result<()> {
 	if range.is_empty() {
 		return Ok(());