Re-classify MemoryGrow as a Statement

codegen-luajit/src/backend/expression.rs

@@ -4,8 +4,8 @@ use std::{
 };
 
 use wasm_ast::node::{
-	BinOp, CmpOp, Expression, GetGlobal, GetLocal, GetTemporary, LoadAt, MemoryGrow, MemorySize,
+	BinOp, CmpOp, Expression, GetGlobal, GetLocal, GetTemporary, LoadAt, MemorySize, Select, UnOp,
-	Select, UnOp, Value,
+	Value,
 };
 
 use crate::analyzer::operator::bin_symbol_of;
@@ -77,14 +77,6 @@ impl Driver for MemorySize {
 	}
 }
 
-impl Driver for MemoryGrow {
-	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		write!(w, "rt.allocator.grow(memory_at_{}, ", self.memory)?;
-		self.value.write(mng, w)?;
-		write!(w, ")")
-	}
-}
-
 impl_write_number!(write_f32, f32);
 impl_write_number!(write_f64, f64);
 
@@ -146,7 +138,6 @@ impl Driver for Expression {
 			Self::GetGlobal(e) => e.write(mng, w),
 			Self::LoadAt(e) => e.write(mng, w),
 			Self::MemorySize(e) => e.write(mng, w),
-			Self::MemoryGrow(e) => e.write(mng, w),
 			Self::Value(e) => e.write(mng, w),
 			Self::UnOp(e) => e.write(mng, w),
 			Self::BinOp(e) => e.write(mng, w),

codegen-luajit/src/backend/statement.rs

@@ -5,8 +5,8 @@ use std::{
 
 use parity_wasm::elements::ValueType;
 use wasm_ast::node::{
-	Backward, Br, BrIf, BrTable, Call, CallIndirect, Forward, FuncData, If, SetGlobal, SetLocal,
+	Backward, Br, BrIf, BrTable, Call, CallIndirect, Forward, FuncData, If, MemoryGrow, SetGlobal,
-	SetTemporary, Statement, StoreAt, Terminator,
+	SetLocal, SetTemporary, Statement, StoreAt, Terminator,
 };
 
 use super::manager::{
@@ -198,6 +198,17 @@ impl Driver for StoreAt {
 	}
 }
 
+impl Driver for MemoryGrow {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		let result = self.result;
+		let memory = self.memory;
+
+		write!(w, "reg_{result} = rt.allocator.grow(memory_at_{memory}, ")?;
+		self.value.write(mng, w)?;
+		write!(w, ")")
+	}
+}
+
 impl Driver for Statement {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		match self {
@@ -211,6 +222,7 @@ impl Driver for Statement {
 			Self::SetLocal(s) => s.write(mng, w),
 			Self::SetGlobal(s) => s.write(mng, w),
 			Self::StoreAt(s) => s.write(mng, w),
+			Self::MemoryGrow(s) => s.write(mng, w),
 		}
 	}
 }

codegen-luau/src/backend/expression.rs

@@ -4,8 +4,8 @@ use std::{
 };
 
 use wasm_ast::node::{
-	BinOp, CmpOp, Expression, GetGlobal, GetLocal, GetTemporary, LoadAt, MemoryGrow, MemorySize,
+	BinOp, CmpOp, Expression, GetGlobal, GetLocal, GetTemporary, LoadAt, MemorySize, Select, UnOp,
-	Select, UnOp, Value,
+	Value,
 };
 
 use crate::analyzer::operator::bin_symbol_of;
@@ -77,14 +77,6 @@ impl Driver for MemorySize {
 	}
 }
 
-impl Driver for MemoryGrow {
-	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
-		write!(w, "rt.allocator.grow(memory_at_{}, ", self.memory)?;
-		self.value.write(mng, w)?;
-		write!(w, ")")
-	}
-}
-
 pub fn write_i32(number: i32, w: &mut dyn Write) -> Result<()> {
 	let list = number.to_ne_bytes();
 
@@ -166,7 +158,6 @@ impl Driver for Expression {
 			Self::GetGlobal(e) => e.write(mng, w),
 			Self::LoadAt(e) => e.write(mng, w),
 			Self::MemorySize(e) => e.write(mng, w),
-			Self::MemoryGrow(e) => e.write(mng, w),
 			Self::Value(e) => e.write(mng, w),
 			Self::UnOp(e) => e.write(mng, w),
 			Self::BinOp(e) => e.write(mng, w),

codegen-luau/src/backend/statement.rs

@@ -5,8 +5,8 @@ use std::{
 
 use parity_wasm::elements::ValueType;
 use wasm_ast::node::{
-	Backward, Br, BrIf, BrTable, Call, CallIndirect, Forward, FuncData, If, SetGlobal, SetLocal,
+	Backward, Br, BrIf, BrTable, Call, CallIndirect, Forward, FuncData, If, MemoryGrow, SetGlobal,
-	SetTemporary, Statement, StoreAt, Terminator,
+	SetLocal, SetTemporary, Statement, StoreAt, Terminator,
 };
 
 use super::manager::{
@@ -233,6 +233,17 @@ impl Driver for StoreAt {
 	}
 }
 
+impl Driver for MemoryGrow {
+	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
+		let result = self.result;
+		let memory = self.memory;
+
+		write!(w, "reg_{result} = rt.allocator.grow(memory_at_{memory}, ")?;
+		self.value.write(mng, w)?;
+		write!(w, ")")
+	}
+}
+
 impl Driver for Statement {
 	fn write(&self, mng: &mut Manager, w: &mut dyn Write) -> Result<()> {
 		match self {
@@ -246,6 +257,7 @@ impl Driver for Statement {
 			Self::SetLocal(s) => s.write(mng, w),
 			Self::SetGlobal(s) => s.write(mng, w),
 			Self::StoreAt(s) => s.write(mng, w),
+			Self::MemoryGrow(s) => s.write(mng, w),
 		}
 	}
 }

wasm-ast/src/builder.rs

@@ -19,7 +19,7 @@ macro_rules! leak_on {
 			let read = ReadType::$variant(id);
 
 			for i in 0..self.stack.len() {
-				if self.stack.get(i).has_read(read) {
+				if self.stack.has_read_at(i, read) {
 					self.leak_at(i);
 				}
 			}
@@ -607,12 +607,6 @@ impl<'a> Builder<'a> {
 				self.add_call_indirect(i.try_into().unwrap(), t.into());
 			}
 			Inst::Drop => {
-				let last = self.target.stack.len() - 1;
-
-				if self.target.stack.get(last).has_side_effect() {
-					self.target.leak_at(last);
-				}
-
 				self.target.stack.pop();
 			}
 			Inst::Select => {
@@ -705,14 +699,19 @@ impl<'a> Builder<'a> {
 				self.target.stack.push(data);
 			}
 			Inst::GrowMemory(i) => {
+				let value = self.target.stack.pop().into();
+				let result = self.target.stack.len();
 				let memory = i.try_into().unwrap();
-				let data = Expression::MemoryGrow(MemoryGrow {
+
+				let data = Statement::MemoryGrow(MemoryGrow {
+					result,
 					memory,
-					value: self.target.stack.pop().into(),
+					value,
 				});
 
-				self.target.stack.push(data);
+				self.target.leak_memory_write(memory);
-				self.target.leak_all();
+				self.target.stack.push_temporary(1);
+				self.target.code.push(data);
 			}
 			Inst::I32Const(v) => self.target.push_constant(v),
 			Inst::I64Const(v) => self.target.push_constant(v),

wasm-ast/src/node.rs

@@ -560,6 +560,7 @@ pub struct MemorySize {
 }
 
 pub struct MemoryGrow {
+	pub result: usize,
 	pub memory: usize,
 	pub value: Box<Expression>,
 }
@@ -620,7 +621,6 @@ pub enum Expression {
 	GetGlobal(GetGlobal),
 	LoadAt(LoadAt),
 	MemorySize(MemorySize),
-	MemoryGrow(MemoryGrow),
 	Value(Value),
 	UnOp(UnOp),
 	BinOp(BinOp),
@@ -736,6 +736,7 @@ pub enum Statement {
 	SetLocal(SetLocal),
 	SetGlobal(SetGlobal),
 	StoreAt(StoreAt),
+	MemoryGrow(MemoryGrow),
 }
 
 pub struct FuncData {

wasm-ast/src/stack.rs

@@ -17,15 +17,7 @@ pub struct Slot {
 }
 
 impl Slot {
-	pub fn has_read(&self, read: ReadType) -> bool {
+	fn is_temporary(&self, id: usize) -> bool {
-		self.read.contains(&read)
-	}
-
-	pub fn has_side_effect(&self) -> bool {
-		matches!(self.data, Expression::MemoryGrow(_))
-	}
-
-	pub fn is_temporary(&self, id: usize) -> bool {
 		matches!(self.data, Expression::GetTemporary(ref v) if v.var == id)
 	}
 }
@@ -42,10 +34,6 @@ impl Stack {
 		self.var_list.len()
 	}
 
-	pub fn get(&self, index: usize) -> &Slot {
-		&self.var_list[index]
-	}
-
 	pub fn split_last(&mut self, len: usize) -> Self {
 		let desired = self.len() - len;
 		let var_list = self.var_list.split_off(desired);
@@ -106,6 +94,22 @@ impl Stack {
 		self.capacity = self.capacity.max(len + num);
 	}
 
+	pub fn has_read_at(&self, index: usize, read: ReadType) -> bool {
+		self.var_list[index].read.contains(&read)
+	}
+
+	// Return the alignment necessary for this block to branch out to a
+	// another given stack frame
+	pub fn get_br_alignment(&self, par_start: usize, par_result: usize) -> Align {
+		let start = self.len() + self.previous - par_result;
+
+		Align {
+			new: par_start,
+			old: start,
+			length: par_result,
+		}
+	}
+
 	// Try to leak a slot's value to a `SetTemporary` instruction,
 	// adjusting the capacity and old index accordingly
 	pub fn leak_at(&mut self, index: usize) -> Option<Statement> {
@@ -128,16 +132,4 @@ impl Stack {
 
 		Some(set)
 	}
-
-	// Return the alignment necessary for this block to branch out to a
-	// another given stack frame
-	pub fn get_br_alignment(&self, par_start: usize, par_result: usize) -> Align {
-		let start = self.len() + self.previous - par_result;
-
-		Align {
-			new: par_start,
-			old: start,
-			length: par_result,
-		}
-	}
 }

wasm-ast/src/visit.rs

@@ -17,8 +17,6 @@ pub trait Visitor {
 
 	fn visit_memory_size(&mut self, _: &MemorySize) {}
 
-	fn visit_memory_grow(&mut self, _: &MemoryGrow) {}
-
 	fn visit_value(&mut self, _: &Value) {}
 
 	fn visit_un_op(&mut self, _: &UnOp) {}
@@ -57,6 +55,8 @@ pub trait Visitor {
 
 	fn visit_store_at(&mut self, _: &StoreAt) {}
 
+	fn visit_memory_grow(&mut self, _: &MemoryGrow) {}
+
 	fn visit_statement(&mut self, _: &Statement) {}
 }
 
@@ -106,14 +106,6 @@ impl<T: Visitor> Driver<T> for MemorySize {
 	}
 }
 
-impl<T: Visitor> Driver<T> for MemoryGrow {
-	fn accept(&self, visitor: &mut T) {
-		self.value.accept(visitor);
-
-		visitor.visit_memory_grow(self);
-	}
-}
-
 impl<T: Visitor> Driver<T> for Value {
 	fn accept(&self, visitor: &mut T) {
 		visitor.visit_value(self);
@@ -155,7 +147,6 @@ impl<T: Visitor> Driver<T> for Expression {
 			Self::GetGlobal(v) => v.accept(visitor),
 			Self::LoadAt(v) => v.accept(visitor),
 			Self::MemorySize(v) => v.accept(visitor),
-			Self::MemoryGrow(v) => v.accept(visitor),
 			Self::Value(v) => v.accept(visitor),
 			Self::UnOp(v) => v.accept(visitor),
 			Self::BinOp(v) => v.accept(visitor),
@@ -296,6 +287,14 @@ impl<T: Visitor> Driver<T> for StoreAt {
 	}
 }
 
+impl<T: Visitor> Driver<T> for MemoryGrow {
+	fn accept(&self, visitor: &mut T) {
+		self.value.accept(visitor);
+
+		visitor.visit_memory_grow(self);
+	}
+}
+
 impl<T: Visitor> Driver<T> for Statement {
 	fn accept(&self, visitor: &mut T) {
 		match self {
@@ -309,6 +308,7 @@ impl<T: Visitor> Driver<T> for Statement {
 			Self::SetLocal(v) => v.accept(visitor),
 			Self::SetGlobal(v) => v.accept(visitor),
 			Self::StoreAt(v) => v.accept(visitor),
+			Self::MemoryGrow(v) => v.accept(visitor),
 		}
 
 		visitor.visit_statement(self);