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

---
 codegen-luajit/src/backend/statement.rs |  7 ------
 codegen-luajit/src/translator.rs        | 29 ++++++++++++++++++++++---
 2 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/codegen-luajit/src/backend/statement.rs b/codegen-luajit/src/backend/statement.rs
index 8de1a9c..bb45254 100644
--- a/codegen-luajit/src/backend/statement.rs
+++ b/codegen-luajit/src/backend/statement.rs
@@ -9,8 +9,6 @@ use wasm_ast::node::{
 	Return, SetGlobal, SetLocal, Statement, StoreAt,
 };
 
-use crate::analyzer::memory;
-
 use super::manager::{
 	write_ascending, write_condition, write_separated, write_variable, Driver, Manager,
 };
@@ -281,11 +279,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)?;
 		write!(w, "local temp ")?;
 
diff --git a/codegen-luajit/src/translator.rs b/codegen-luajit/src/translator.rs
index 80c14ca..8e00c61 100644
--- a/codegen-luajit/src/translator.rs
+++ b/codegen-luajit/src/translator.rs
@@ -13,7 +13,7 @@ use wasm_ast::{
 };
 
 use crate::{
-	analyzer::localize,
+	analyzer::{localize, memory},
 	backend::manager::{write_f32, write_f64, Driver, Manager},
 };
 
@@ -250,6 +250,22 @@ fn write_localize_used(func_list: &[Intermediate], w: &mut dyn Write) -> Result<
 		.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 mut mem_set = BTreeSet::new();
+
+	for func in func_list {
+		mem_set.extend(memory::visit(func));
+	}
+
+	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()
@@ -280,19 +296,24 @@ fn write_func_list(
 	})
 }
 
-fn write_module_start(wasm: &Module, w: &mut dyn Write) -> Result<()> {
+fn write_module_start(wasm: &Module, 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)?;
 	write_global_list(wasm, w)?;
 	write_element_list(wasm, w)?;
 	write_data_list(wasm, w)?;
+
 	write!(w, "end ")?;
 
 	write!(w, "return function(wasm)")?;
 	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}]()")?;
 	}
@@ -309,6 +330,8 @@ 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!(w, "local table_new = require(\"table.new\")")?;
 	write_named_array("FUNC_LIST", wasm.functions_space(), w)?;
 	write_named_array("TABLE_LIST", wasm.table_space(), w)?;
@@ -316,5 +339,5 @@ pub fn translate(wasm: &Module, type_info: &TypeInfo, w: &mut dyn Write) -> Resu
 	write_named_array("GLOBAL_LIST", wasm.globals_space(), w)?;
 
 	write_func_list(wasm, type_info, &func_list, w)?;
-	write_module_start(wasm, w)
+	write_module_start(wasm, &mem_list, w)
 }