From 25e2730a0ca21e1feb7a6e029359bdbd76a3b7c3 Mon Sep 17 00:00:00 2001
From: Rerumu <rerumu@protonmail.com>
Date: Wed, 22 Dec 2021 00:23:02 -0500
Subject: [PATCH] Add typed locals

---
 wasm/src/ast/builder.rs   | 13 ++++++++-----
 wasm/src/ast/node.rs      |  4 ++--
 wasm/src/writer/luajit.rs |  5 +++++
 wasm/src/writer/luau.rs   |  5 +++++
 wasm/src/writer/shared.rs | 16 ++++++++++++----
 5 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/wasm/src/ast/builder.rs b/wasm/src/ast/builder.rs
index 95d229e..9e66459 100644
--- a/wasm/src/ast/builder.rs
+++ b/wasm/src/ast/builder.rs
@@ -1,5 +1,5 @@
 use parity_wasm::elements::{
-	BlockType, External, FuncBody, FunctionType, ImportEntry, Instruction, Local, Module, Type,
+	BlockType, External, FuncBody, FunctionType, ImportEntry, Instruction, Module, Type, ValueType,
 };
 
 use super::{
@@ -120,8 +120,11 @@ fn is_dead_precursor(inst: &Instruction) -> bool {
 	)
 }
 
-fn local_sum(body: &FuncBody) -> u32 {
-	body.locals().iter().map(Local::count).sum()
+fn load_local_list(func: &FuncBody) -> Vec<ValueType> {
+	func.locals()
+		.iter()
+		.flat_map(|l| std::iter::repeat(l.value_type()).take(l.count().try_into().unwrap()))
+		.collect()
 }
 
 fn load_func_at(wasm: &Module, index: usize) -> &FuncBody {
@@ -144,8 +147,8 @@ impl<'a> Builder<'a> {
 		let func = load_func_at(self.wasm, index);
 		let arity = &self.other.in_arity[index];
 
+		let local_list = load_local_list(func);
 		let num_param = arity.num_param;
-		let num_local = local_sum(func);
 
 		self.num_result = arity.num_result;
 
@@ -153,8 +156,8 @@ impl<'a> Builder<'a> {
 		let num_stack = self.last_stack.try_into().unwrap();
 
 		Function {
+			local_list,
 			num_param,
-			num_local,
 			num_stack,
 			body,
 		}
diff --git a/wasm/src/ast/node.rs b/wasm/src/ast/node.rs
index 4771630..dd833bf 100644
--- a/wasm/src/ast/node.rs
+++ b/wasm/src/ast/node.rs
@@ -1,6 +1,6 @@
 use std::ops::Range;
 
-use parity_wasm::elements::BrTableData;
+use parity_wasm::elements::{BrTableData, ValueType};
 
 use super::tag::{BinOp, CmpOp, Load, Store, UnOp};
 
@@ -182,8 +182,8 @@ pub enum Statement {
 }
 
 pub struct Function {
+	pub local_list: Vec<ValueType>,
 	pub num_param: u32,
-	pub num_local: u32,
 	pub num_stack: u32,
 	pub body: Forward,
 }
diff --git a/wasm/src/writer/luajit.rs b/wasm/src/writer/luajit.rs
index 6d06100..f21fcbb 100644
--- a/wasm/src/writer/luajit.rs
+++ b/wasm/src/writer/luajit.rs
@@ -734,6 +734,11 @@ impl<'a> Transpiler<'a> for LuaJIT<'a> {
 
 		Self::gen_localize(&func_list, w)?;
 
+		write!(w, "local ZERO_i32 = 0 ")?;
+		write!(w, "local ZERO_i64 = 0LL ")?;
+		write!(w, "local ZERO_f32 = 0.0 ")?;
+		write!(w, "local ZERO_f64 = 0.0 ")?;
+
 		write!(w, "local table_new = require(\"table.new\")")?;
 		write_list("FUNC_LIST", self.wasm.functions_space(), w)?;
 		write_list("TABLE_LIST", self.wasm.table_space(), w)?;
diff --git a/wasm/src/writer/luau.rs b/wasm/src/writer/luau.rs
index e992a89..60e6616 100644
--- a/wasm/src/writer/luau.rs
+++ b/wasm/src/writer/luau.rs
@@ -723,6 +723,11 @@ impl<'a> Transpiler<'a> for Luau<'a> {
 
 		Self::gen_localize(&func_list, w)?;
 
+		write!(w, "local ZERO_i32 = 0 ")?;
+		write!(w, "local ZERO_i64 = 0 ")?;
+		write!(w, "local ZERO_f32 = 0.0 ")?;
+		write!(w, "local ZERO_f64 = 0.0 ")?;
+
 		write_list("FUNC_LIST", self.wasm.functions_space(), w)?;
 		write_list("TABLE_LIST", self.wasm.table_space(), w)?;
 		write_list("MEMORY_LIST", self.wasm.memory_space(), w)?;
diff --git a/wasm/src/writer/shared.rs b/wasm/src/writer/shared.rs
index 363d441..0b8d671 100644
--- a/wasm/src/writer/shared.rs
+++ b/wasm/src/writer/shared.rs
@@ -104,12 +104,20 @@ pub fn write_result_list(range: Range<u32>, w: Writer) -> Result<()> {
 }
 
 pub fn write_variable_list(func: &Function, w: Writer) -> Result<()> {
-	if func.num_local != 0 {
-		let list = vec!["0"; func.num_local as usize].join(", ");
+	if !func.local_list.is_empty() {
+		let num_local = func.local_list.len().try_into().unwrap();
 
 		write!(w, "local ")?;
-		write_in_order("loc", func.num_local, w)?;
-		write!(w, " = {} ", list)?;
+		write_in_order("loc", num_local, w)?;
+		write!(w, " = ")?;
+
+		for (i, t) in func.local_list.iter().enumerate() {
+			if i != 0 {
+				write!(w, ", ")?;
+			}
+
+			write!(w, "ZERO_{} ", t)?;
+		}
 	}
 
 	if func.num_stack != 0 {