From ce59f77f22d639f6b300196c3b860fb3734b1593 Mon Sep 17 00:00:00 2001
From: Rerumu <rerumu@protonmail.com>
Date: Fri, 26 Nov 2021 16:05:03 -0500
Subject: [PATCH] Improve runtime localizing

---
 src/backend/translator/data.rs   | 36 +++++++++++++++++++++++++-------
 src/backend/translator/writer.rs | 22 +++++--------------
 2 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/src/backend/translator/data.rs b/src/backend/translator/data.rs
index 32eef39..4dfc807 100644
--- a/src/backend/translator/data.rs
+++ b/src/backend/translator/data.rs
@@ -1,10 +1,17 @@
-use std::io::{Result, Write};
+use std::{
+	collections::BTreeSet,
+	io::{Result, Write},
+};
 
 use parity_wasm::elements::{
 	External, ImportCountType, Instruction, Internal, Module as WasmModule, ResizableLimits,
 };
 
-use crate::backend::{ast::transformer::Transformer, edition::data::Edition};
+use crate::backend::{
+	ast::{data::Function, transformer::Transformer},
+	edition::data::Edition,
+	visitor::localize,
+};
 
 use super::{arity::List as ArityList, writer::Data};
 
@@ -261,9 +268,27 @@ impl<'a> Module<'a> {
 		write!(w, "}} end ")
 	}
 
+	fn gen_localize(&self, func_list: &[Function], w: &mut dyn Write) -> Result<()> {
+		let mut loc_set = BTreeSet::new();
+
+		for func in func_list {
+			loc_set.extend(localize::visit(func));
+		}
+
+		loc_set
+			.into_iter()
+			.try_for_each(|(a, b)| write!(w, "local {0}_{1} = rt.{0}.{1} ", a, b))
+	}
+
 	pub fn translate(&self, ed: &dyn Edition, w: &mut dyn Write) -> Result<()> {
 		write!(w, "local rt = require({})", ed.runtime())?;
 
+		let func_list: Vec<_> = (0..self.arity.in_arity.len())
+			.map(|i| Transformer::new(self.wasm, &self.arity, i).consume())
+			.collect();
+
+		self.gen_localize(&func_list, w)?;
+
 		gen_nil_array("FUNC_LIST", self.wasm.functions_space(), w)?;
 		gen_nil_array("TABLE_LIST", self.wasm.table_space(), w)?;
 		gen_nil_array("MEMORY_LIST", self.wasm.memory_space(), w)?;
@@ -271,13 +296,10 @@ impl<'a> Module<'a> {
 
 		let offset = self.arity.ex_arity.len();
 
-		for i in 0..self.arity.in_arity.len() {
-			let func = Transformer::new(self.wasm, &self.arity, i).consume();
-			let data = &mut Data::new(func.num_param, ed);
-
+		for (i, v) in func_list.into_iter().enumerate() {
 			write!(w, "FUNC_LIST[{}] =", i + offset)?;
 
-			func.output(data, w)?;
+			v.output(&mut Data::new(v.num_param, ed), w)?;
 		}
 
 		self.gen_start_point(w)
diff --git a/src/backend/translator/writer.rs b/src/backend/translator/writer.rs
index 5cc3210..e2b88f6 100644
--- a/src/backend/translator/writer.rs
+++ b/src/backend/translator/writer.rs
@@ -10,7 +10,7 @@ use crate::backend::{
 		Return, Select, SetGlobal, SetLocal, Statement, Value,
 	},
 	edition::data::Edition,
-	visitor::{localize, memory},
+	visitor::memory,
 };
 
 fn write_in_order(prefix: &'static str, len: u32, w: &mut dyn Write) -> Result<()> {
@@ -483,21 +483,6 @@ impl Function {
 		Ok(())
 	}
 
-	fn write_visitor_list(&self, w: &mut dyn Write) -> Result<()> {
-		let memory = memory::visit(self);
-		let localize = localize::visit(self);
-
-		for v in memory {
-			write!(w, "local memory_at_{0} = MEMORY_LIST[{0}]", v)?;
-		}
-
-		for (a, b) in localize {
-			write!(w, "local {0}_{1} = rt.{0}.{1} ", a, b)?;
-		}
-
-		Ok(())
-	}
-
 	pub fn output(&self, d: &mut Data, w: &mut dyn Write) -> Result<()> {
 		write!(w, "function(")?;
 
@@ -505,7 +490,10 @@ impl Function {
 
 		write!(w, ")")?;
 
-		self.write_visitor_list(w)?;
+		for v in memory::visit(self) {
+			write!(w, "local memory_at_{0} = MEMORY_LIST[{0}]", v)?;
+		}
+
 		self.write_variable_list(w)?;
 
 		self.body.output(d, w)?;