Re-classify MemoryGrow as a Statement

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_all();
+				self.target.leak_memory_write(memory);
+				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 {
-		self.read.contains(&read)
-	}
-
-	pub fn has_side_effect(&self) -> bool {
-		matches!(self.data, Expression::MemoryGrow(_))
-	}
-
-	pub fn is_temporary(&self, id: usize) -> bool {
+	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);