Unify local and memory finding step

codegen-luajit/src/analyzer/localize.rs

@@ -1,33 +1,36 @@
 use std::collections::BTreeSet;
 
 use wasm_ast::{
-	node::{BinOp, CmpOp, FuncData, LoadAt, StoreAt, UnOp},
+	node::{BinOp, CmpOp, FuncData, LoadAt, MemoryGrow, MemorySize, StoreAt, UnOp},
 	visit::{Driver, Visitor},
 };
 
 use super::operator::{bin_symbol_of, cmp_symbol_of};
 
 struct Visit {
-	result: BTreeSet<(&'static str, &'static str)>,
+	local_set: BTreeSet<(&'static str, &'static str)>,
+	memory_set: BTreeSet<usize>,
 }
 
 impl Visitor for Visit {
 	fn visit_load_at(&mut self, v: &LoadAt) {
 		let name = v.what.as_name();
 
-		self.result.insert(("load", name));
+		self.memory_set.insert(0);
+		self.local_set.insert(("load", name));
 	}
 
 	fn visit_store_at(&mut self, v: &StoreAt) {
 		let name = v.what.as_name();
 
-		self.result.insert(("store", name));
+		self.memory_set.insert(0);
+		self.local_set.insert(("store", name));
 	}
 
 	fn visit_un_op(&mut self, v: &UnOp) {
 		let name = v.op.as_name();
 
-		self.result.insert(name);
+		self.local_set.insert(name);
 	}
 
 	fn visit_bin_op(&mut self, v: &BinOp) {
@@ -37,7 +40,7 @@ impl Visitor for Visit {
 
 		let name = v.op.as_name();
 
-		self.result.insert(name);
+		self.local_set.insert(name);
 	}
 
 	fn visit_cmp_op(&mut self, v: &CmpOp) {
@@ -47,16 +50,25 @@ impl Visitor for Visit {
 
 		let name = v.op.as_name();
 
-		self.result.insert(name);
+		self.local_set.insert(name);
+	}
+
+	fn visit_memory_size(&mut self, m: &MemorySize) {
+		self.memory_set.insert(m.memory);
+	}
+
+	fn visit_memory_grow(&mut self, m: &MemoryGrow) {
+		self.memory_set.insert(m.memory);
 	}
 }
 
-pub fn visit(ast: &FuncData) -> BTreeSet<(&'static str, &'static str)> {
+pub fn visit(ast: &FuncData) -> (BTreeSet<(&'static str, &'static str)>, BTreeSet<usize>) {
 	let mut visit = Visit {
-		result: BTreeSet::new(),
+		local_set: BTreeSet::new(),
+		memory_set: BTreeSet::new(),
 	};
 
 	ast.accept(&mut visit);
 
-	visit.result
+	(visit.local_set, visit.memory_set)
 }

codegen-luajit/src/analyzer/memory.rs

@@ -1,38 +0,0 @@
-use std::collections::BTreeSet;
-
-use wasm_ast::{
-	node::{FuncData, LoadAt, MemoryGrow, MemorySize, StoreAt},
-	visit::{Driver, Visitor},
-};
-
-struct Visit {
-	result: BTreeSet<usize>,
-}
-
-impl Visitor for Visit {
-	fn visit_store_at(&mut self, _: &StoreAt) {
-		self.result.insert(0);
-	}
-
-	fn visit_load_at(&mut self, _: &LoadAt) {
-		self.result.insert(0);
-	}
-
-	fn visit_memory_size(&mut self, m: &MemorySize) {
-		self.result.insert(m.memory);
-	}
-
-	fn visit_memory_grow(&mut self, m: &MemoryGrow) {
-		self.result.insert(m.memory);
-	}
-}
-
-pub fn visit(ast: &FuncData) -> BTreeSet<usize> {
-	let mut visit = Visit {
-		result: BTreeSet::new(),
-	};
-
-	ast.accept(&mut visit);
-
-	visit.result
-}

codegen-luajit/src/analyzer/mod.rs

@@ -1,3 +1,2 @@
 pub mod localize;
-pub mod memory;
 pub mod operator;

codegen-luajit/src/translator.rs

@@ -13,7 +13,7 @@ use wasm_ast::{
 };
 
 use crate::{
-	analyzer::{localize, memory},
+	analyzer::localize,
 	backend::manager::{Driver, Manager},
 };
 
@@ -241,23 +241,24 @@ fn write_local_operation(head: &str, tail: &str, w: &mut dyn Write) -> Result<()
 	}
 }
 
-fn write_localize_used(func_list: &[FuncData], w: &mut dyn Write) -> Result<()> {
-	let loc_set: BTreeSet<_> = func_list.iter().flat_map(localize::visit).collect();
+fn write_localize_used(func_list: &[FuncData], w: &mut dyn Write) -> Result<BTreeSet<usize>> {
+	let mut loc_set = BTreeSet::new();
+	let mut mem_set = BTreeSet::new();
 
-	loc_set
-		.into_iter()
-		.try_for_each(|(a, b)| write_local_operation(a, b, w))
+	for (loc, mem) in func_list.iter().map(localize::visit) {
+		loc_set.extend(loc);
+		mem_set.extend(mem);
 	}
 
-fn write_memory_used(func_list: &[FuncData], w: &mut dyn Write) -> Result<Vec<usize>> {
-	let mem_set: BTreeSet<_> = func_list.iter().flat_map(memory::visit).collect();
-	let list: Vec<_> = mem_set.into_iter().collect();
+	for loc in loc_set {
+		write_local_operation(loc.0, loc.1, w)?;
+	}
 
-	for mem in &list {
+	for mem in &mem_set {
 		write!(w, "local memory_at_{mem} ")?;
 	}
 
-	Ok(list)
+	Ok(mem_set)
 }
 
 fn write_func_start(wasm: &Module, index: u32, w: &mut dyn Write) -> Result<()> {
@@ -293,7 +294,7 @@ fn write_func_list(
 fn write_module_start(
 	wasm: &Module,
 	type_info: &TypeInfo,
-	mem_list: &[usize],
+	mem_set: &BTreeSet<usize>,
 	w: &mut dyn Write,
 ) -> Result<()> {
 	write!(w, "local function run_init_code()")?;
@@ -308,7 +309,7 @@ fn write_module_start(
 	write_import_list(wasm, w)?;
 	write!(w, "run_init_code()")?;
 
-	for mem in mem_list {
+	for mem in mem_set {
 		write!(w, "memory_at_{mem} = MEMORY_LIST[{mem}]")?;
 	}
 
@@ -333,10 +334,7 @@ pub fn from_inst_list(code: &[Instruction], type_info: &TypeInfo, w: &mut dyn Wr
 /// Returns `Err` if writing to `Write` failed.
 pub fn from_module_typed(wasm: &Module, type_info: &TypeInfo, w: &mut dyn Write) -> Result<()> {
 	let func_list = build_func_list(wasm, type_info);
-
-	write_localize_used(&func_list, w)?;
-
-	let mem_list = write_memory_used(&func_list, w)?;
+	let mem_set = write_localize_used(&func_list, w)?;
 
 	write!(w, "local table_new = require(\"table.new\")")?;
 	write_named_array("FUNC_LIST", wasm.functions_space(), w)?;
@@ -345,7 +343,7 @@ pub fn from_module_typed(wasm: &Module, type_info: &TypeInfo, w: &mut dyn Write)
 	write_named_array("GLOBAL_LIST", wasm.globals_space(), w)?;
 
 	write_func_list(wasm, type_info, &func_list, w)?;
-	write_module_start(wasm, type_info, &mem_list, w)
+	write_module_start(wasm, type_info, &mem_set, w)
 }
 
 /// # Errors

codegen-luau/src/analyzer/localize.rs

@@ -2,27 +2,30 @@ use std::collections::BTreeSet;
 
 use parity_wasm::elements::ValueType;
 use wasm_ast::{
-	node::{BinOp, CmpOp, FuncData, LoadAt, StoreAt, UnOp, Value},
+	node::{BinOp, CmpOp, FuncData, LoadAt, MemoryGrow, MemorySize, StoreAt, UnOp, Value},
 	visit::{Driver, Visitor},
 };
 
 use super::operator::{bin_symbol_of, cmp_symbol_of};
 
 struct Visit {
-	result: BTreeSet<(&'static str, &'static str)>,
+	local_set: BTreeSet<(&'static str, &'static str)>,
+	memory_set: BTreeSet<usize>,
 }
 
 impl Visitor for Visit {
 	fn visit_load_at(&mut self, v: &LoadAt) {
 		let name = v.what.as_name();
 
-		self.result.insert(("load", name));
+		self.memory_set.insert(0);
+		self.local_set.insert(("load", name));
 	}
 
 	fn visit_store_at(&mut self, v: &StoreAt) {
 		let name = v.what.as_name();
 
-		self.result.insert(("store", name));
+		self.memory_set.insert(0);
+		self.local_set.insert(("store", name));
 	}
 
 	fn visit_value(&mut self, v: &Value) {
@@ -33,13 +36,13 @@ impl Visitor for Visit {
 			_ => return,
 		};
 
-		self.result.insert(("i64", name));
+		self.local_set.insert(("i64", name));
 	}
 
 	fn visit_un_op(&mut self, v: &UnOp) {
 		let name = v.op.as_name();
 
-		self.result.insert(name);
+		self.local_set.insert(name);
 	}
 
 	fn visit_bin_op(&mut self, v: &BinOp) {
@@ -49,7 +52,7 @@ impl Visitor for Visit {
 
 		let name = v.op.as_name();
 
-		self.result.insert(name);
+		self.local_set.insert(name);
 	}
 
 	fn visit_cmp_op(&mut self, v: &CmpOp) {
@@ -59,13 +62,22 @@ impl Visitor for Visit {
 
 		let name = v.op.as_name();
 
-		self.result.insert(name);
+		self.local_set.insert(name);
+	}
+
+	fn visit_memory_size(&mut self, m: &MemorySize) {
+		self.memory_set.insert(m.memory);
+	}
+
+	fn visit_memory_grow(&mut self, m: &MemoryGrow) {
+		self.memory_set.insert(m.memory);
 	}
 }
 
-pub fn visit(ast: &FuncData) -> BTreeSet<(&'static str, &'static str)> {
+pub fn visit(ast: &FuncData) -> (BTreeSet<(&'static str, &'static str)>, BTreeSet<usize>) {
 	let mut visit = Visit {
-		result: BTreeSet::new(),
+		local_set: BTreeSet::new(),
+		memory_set: BTreeSet::new(),
 	};
 
 	if ast
@@ -73,10 +85,10 @@ pub fn visit(ast: &FuncData) -> BTreeSet<(&'static str, &'static str)> {
 		.iter()
 		.any(|v| v.value_type() == ValueType::I64)
 	{
-		visit.result.insert(("i64", "K_ZERO"));
+		visit.local_set.insert(("i64", "K_ZERO"));
 	}
 
 	ast.accept(&mut visit);
 
-	visit.result
+	(visit.local_set, visit.memory_set)
 }

codegen-luau/src/analyzer/memory.rs

@@ -1,38 +0,0 @@
-use std::collections::BTreeSet;
-
-use wasm_ast::{
-	node::{FuncData, LoadAt, MemoryGrow, MemorySize, StoreAt},
-	visit::{Driver, Visitor},
-};
-
-struct Visit {
-	result: BTreeSet<usize>,
-}
-
-impl Visitor for Visit {
-	fn visit_store_at(&mut self, _: &StoreAt) {
-		self.result.insert(0);
-	}
-
-	fn visit_load_at(&mut self, _: &LoadAt) {
-		self.result.insert(0);
-	}
-
-	fn visit_memory_size(&mut self, m: &MemorySize) {
-		self.result.insert(m.memory);
-	}
-
-	fn visit_memory_grow(&mut self, m: &MemoryGrow) {
-		self.result.insert(m.memory);
-	}
-}
-
-pub fn visit(ast: &FuncData) -> BTreeSet<usize> {
-	let mut visit = Visit {
-		result: BTreeSet::new(),
-	};
-
-	ast.accept(&mut visit);
-
-	visit.result
-}

codegen-luau/src/analyzer/mod.rs

@@ -1,3 +1,2 @@
 pub mod localize;
-pub mod memory;
 pub mod operator;

codegen-luau/src/translator.rs

@@ -13,7 +13,7 @@ use wasm_ast::{
 };
 
 use crate::{
-	analyzer::{localize, memory},
+	analyzer::localize,
 	backend::manager::{Driver, Manager},
 };
 
@@ -237,23 +237,24 @@ fn write_local_operation(head: &str, tail: &str, w: &mut dyn Write) -> Result<()
 	}
 }
 
-fn write_localize_used(func_list: &[FuncData], w: &mut dyn Write) -> Result<()> {
-	let loc_set: BTreeSet<_> = func_list.iter().flat_map(localize::visit).collect();
+fn write_localize_used(func_list: &[FuncData], w: &mut dyn Write) -> Result<BTreeSet<usize>> {
+	let mut loc_set = BTreeSet::new();
+	let mut mem_set = BTreeSet::new();
 
-	loc_set
-		.into_iter()
-		.try_for_each(|(a, b)| write_local_operation(a, b, w))
+	for (loc, mem) in func_list.iter().map(localize::visit) {
+		loc_set.extend(loc);
+		mem_set.extend(mem);
 	}
 
-fn write_memory_used(func_list: &[FuncData], w: &mut dyn Write) -> Result<Vec<usize>> {
-	let mem_set: BTreeSet<_> = func_list.iter().flat_map(memory::visit).collect();
-	let list: Vec<_> = mem_set.into_iter().collect();
+	for loc in loc_set {
+		write_local_operation(loc.0, loc.1, w)?;
+	}
 
-	for mem in &list {
+	for mem in &mem_set {
 		write!(w, "local memory_at_{mem} ")?;
 	}
 
-	Ok(list)
+	Ok(mem_set)
 }
 
 fn write_func_start(wasm: &Module, index: u32, w: &mut dyn Write) -> Result<()> {
@@ -289,7 +290,7 @@ fn write_func_list(
 fn write_module_start(
 	wasm: &Module,
 	type_info: &TypeInfo,
-	mem_list: &[usize],
+	mem_set: &BTreeSet<usize>,
 	w: &mut dyn Write,
 ) -> Result<()> {
 	write!(w, "local function run_init_code()")?;
@@ -304,7 +305,7 @@ fn write_module_start(
 	write_import_list(wasm, w)?;
 	write!(w, "run_init_code()")?;
 
-	for mem in mem_list {
+	for mem in mem_set {
 		write!(w, "memory_at_{mem} = MEMORY_LIST[{mem}]")?;
 	}
 
@@ -329,10 +330,7 @@ pub fn from_inst_list(code: &[Instruction], type_info: &TypeInfo, w: &mut dyn Wr
 /// Returns `Err` if writing to `Write` failed.
 pub fn from_module_typed(wasm: &Module, type_info: &TypeInfo, w: &mut dyn Write) -> Result<()> {
 	let func_list = build_func_list(wasm, type_info);
-
-	write_localize_used(&func_list, w)?;
-
-	let mem_list = write_memory_used(&func_list, w)?;
+	let mem_set = write_localize_used(&func_list, w)?;
 
 	write_named_array("FUNC_LIST", wasm.functions_space(), w)?;
 	write_named_array("TABLE_LIST", wasm.table_space(), w)?;
@@ -340,7 +338,7 @@ pub fn from_module_typed(wasm: &Module, type_info: &TypeInfo, w: &mut dyn Write)
 	write_named_array("GLOBAL_LIST", wasm.globals_space(), w)?;
 
 	write_func_list(wasm, type_info, &func_list, w)?;
-	write_module_start(wasm, type_info, &mem_list, w)
+	write_module_start(wasm, type_info, &mem_set, w)
 }
 
 /// # Errors