From f7978ec402f4b6d79d60012a99183b35d71edd68 Mon Sep 17 00:00:00 2001
From: Rerumu <rerumu@protonmail.com>
Date: Thu, 19 May 2022 04:19:53 -0400
Subject: [PATCH] Optimize memory instance caching

---
 codegen-luau/src/backend/statement.rs |  7 ------
 codegen-luau/src/translator.rs        | 34 ++++++++++++++++++++-------
 2 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/codegen-luau/src/backend/statement.rs b/codegen-luau/src/backend/statement.rs
index 5b4e042..87e54d3 100644
--- a/codegen-luau/src/backend/statement.rs
+++ b/codegen-luau/src/backend/statement.rs
@@ -8,8 +8,6 @@ use wasm_ast::node::{
 	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<()> {
@@ -286,11 +284,6 @@ fn write_variable_list(ir: &Intermediate, w: &mut dyn Write) -> Result<()> {
 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;
diff --git a/codegen-luau/src/translator.rs b/codegen-luau/src/translator.rs
index 74337fb..4c3f22b 100644
--- a/codegen-luau/src/translator.rs
+++ b/codegen-luau/src/translator.rs
@@ -13,7 +13,7 @@ use wasm_ast::{
 };
 
 use crate::{
-	analyzer::localize,
+	analyzer::{localize, memory},
 	backend::manager::{Driver, Manager},
 };
 
@@ -230,17 +230,24 @@ fn build_func_list(wasm: &Module, type_info: &TypeInfo) -> Vec<Intermediate> {
 }
 
 fn write_localize_used(func_list: &[Intermediate], w: &mut dyn Write) -> Result<()> {
-	let mut loc_set = BTreeSet::new();
-
-	for func in func_list {
-		loc_set.extend(localize::visit(func));
-	}
+	let loc_set: BTreeSet<_> = func_list.iter().flat_map(localize::visit).collect();
 
 	loc_set
 		.into_iter()
 		.try_for_each(|(a, b)| write!(w, "local {a}_{b} = rt.{a}.{b} "))
 }
 
+fn write_memory_used(func_list: &[Intermediate], 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 mem in &list {
+		write!(w, "local memory_at_{mem} ")?;
+	}
+
+	Ok(list)
+}
+
 fn write_func_start(wasm: &Module, index: u32, offset: u32, w: &mut dyn Write) -> Result<()> {
 	let opt = wasm
 		.names_section()
@@ -271,7 +278,12 @@ fn write_func_list(
 	})
 }
 
-fn write_module_start(wasm: &Module, type_info: &TypeInfo, w: &mut dyn Write) -> Result<()> {
+fn write_module_start(
+	wasm: &Module,
+	type_info: &TypeInfo,
+	mem_list: &[usize],
+	w: &mut dyn Write,
+) -> Result<()> {
 	write!(w, "local function run_init_code()")?;
 	write_table_list(wasm, w)?;
 	write_memory_list(wasm, w)?;
@@ -284,6 +296,10 @@ fn write_module_start(wasm: &Module, type_info: &TypeInfo, w: &mut dyn Write) ->
 	write_import_list(wasm, w)?;
 	write!(w, "run_init_code()")?;
 
+	for mem in mem_list {
+		write!(w, "memory_at_{mem} = MEMORY_LIST[{mem}]")?;
+	}
+
 	if let Some(start) = wasm.start_section() {
 		write!(w, "FUNC_LIST[{start}]()")?;
 	}
@@ -300,11 +316,13 @@ pub fn translate(wasm: &Module, type_info: &TypeInfo, w: &mut dyn Write) -> Resu
 
 	write_localize_used(&func_list, w)?;
 
+	let mem_list = write_memory_used(&func_list, w)?;
+
 	write_named_array("FUNC_LIST", wasm.functions_space(), w)?;
 	write_named_array("TABLE_LIST", wasm.table_space(), w)?;
 	write_named_array("MEMORY_LIST", wasm.memory_space(), w)?;
 	write_named_array("GLOBAL_LIST", wasm.globals_space(), w)?;
 
 	write_func_list(wasm, type_info, &func_list, w)?;
-	write_module_start(wasm, type_info, w)
+	write_module_start(wasm, type_info, &mem_list, w)
 }