Fix bulk memory stores

wasm-ast/src/factory.rs

@@ -4,9 +4,9 @@ use crate::{
 	module::{read_checked, TypeInfo},
 	node::{
 		BinOp, BinOpType, Block, Br, BrIf, BrTable, Call, CallIndirect, CmpOp, CmpOpType,
-		Expression, FuncData, GetGlobal, GetLocal, If, LabelType, LoadAt, LoadType, MemoryCopy,
-		MemoryFill, MemoryGrow, MemorySize, Select, SetGlobal, SetLocal, Statement, StoreAt,
-		StoreType, Terminator, UnOp, UnOpType, Value,
+		Expression, FuncData, GetGlobal, GetLocal, If, LabelType, LoadAt, LoadType, MemoryArgument,
+		MemoryCopy, MemoryFill, MemoryGrow, MemorySize, Select, SetGlobal, SetLocal, Statement,
+		StoreAt, StoreType, Terminator, UnOp, UnOpType, Value,
 	},
 	stack::{ReadType, Stack},
 };
@@ -611,19 +611,43 @@ impl<'a> Factory<'a> {
 			Operator::MemoryCopy { dst_mem, src_mem } => {
 				let size = self.target.stack.pop().into();
 
+				let source = MemoryArgument {
+					memory: src_mem.try_into().unwrap(),
+					pointer: self.target.stack.pop().into(),
+				};
+
+				let destination = MemoryArgument {
+					memory: dst_mem.try_into().unwrap(),
+					pointer: self.target.stack.pop().into(),
+				};
+
+				self.target.leak_memory_write(source.memory);
+				self.target.leak_memory_write(destination.memory);
+
 				let data = Statement::MemoryCopy(MemoryCopy {
-					dst: dst_mem,
-					src: src_mem,
+					destination,
+					source,
 					size,
 				});
 
 				self.target.code.push(data);
 			}
 			Operator::MemoryFill { mem } => {
-				let n = self.target.stack.pop().into();
+				let size = self.target.stack.pop().into();
 				let value = self.target.stack.pop().into();
 
-				let data = Statement::MemoryFill(MemoryFill { mem, value, n });
+				let destination = MemoryArgument {
+					memory: mem.try_into().unwrap(),
+					pointer: self.target.stack.pop().into(),
+				};
+
+				self.target.leak_memory_write(destination.memory);
+
+				let data = Statement::MemoryFill(MemoryFill {
+					destination,
+					size,
+					value,
+				});
 
 				self.target.code.push(data);
 			}

wasm-ast/src/node.rs

@@ -1108,21 +1108,40 @@ impl MemoryGrow {
 	}
 }
 
+pub struct MemoryArgument {
+	pub(crate) memory: usize,
+	pub(crate) pointer: Box<Expression>,
+}
+
+impl MemoryArgument {
+	#[must_use]
+	pub const fn memory(&self) -> usize {
+		self.memory
+	}
+
+	#[must_use]
+	pub fn pointer(&self) -> &Expression {
+		&self.pointer
+	}
+}
+
 pub struct MemoryCopy {
-	pub(crate) dst: u32,
-	pub(crate) src: u32,
+	pub(crate) destination: MemoryArgument,
+	pub(crate) source: MemoryArgument,
 	pub(crate) size: Box<Expression>,
 }
 
 impl MemoryCopy {
 	#[must_use]
-	pub const fn dst(&self) -> u32 {
-		self.dst
+	pub const fn destination(&self) -> &MemoryArgument {
+		&self.destination
 	}
+
 	#[must_use]
-	pub const fn src(&self) -> u32 {
-		self.src
+	pub const fn source(&self) -> &MemoryArgument {
+		&self.source
 	}
+
 	#[must_use]
 	pub fn size(&self) -> &Expression {
 		&self.size
@@ -1130,24 +1149,26 @@ impl MemoryCopy {
 }
 
 pub struct MemoryFill {
-	pub(crate) mem: u32,
+	pub(crate) destination: MemoryArgument,
+	pub(crate) size: Box<Expression>,
 	pub(crate) value: Box<Expression>,
-	pub(crate) n: Box<Expression>,
 }
 
 impl MemoryFill {
 	#[must_use]
-	pub const fn mem(&self) -> u32 {
-		self.mem
+	pub const fn destination(&self) -> &MemoryArgument {
+		&self.destination
 	}
+
+	#[must_use]
+	pub fn size(&self) -> &Expression {
+		&self.size
+	}
+
 	#[must_use]
 	pub fn value(&self) -> &Expression {
 		&self.value
 	}
-	#[must_use]
-	pub fn n(&self) -> &Expression {
-		&self.n
-	}
 }
 
 pub enum Statement {

wasm-ast/src/visit.rs

@@ -110,6 +110,8 @@ impl<T: Visitor> Driver<T> for MemorySize {
 
 impl<T: Visitor> Driver<T> for MemoryCopy {
 	fn accept(&self, visitor: &mut T) {
+		self.destination().pointer().accept(visitor);
+		self.source().pointer().accept(visitor);
 		self.size().accept(visitor);
 
 		visitor.visit_memory_copy(self);
@@ -118,8 +120,9 @@ impl<T: Visitor> Driver<T> for MemoryCopy {
 
 impl<T: Visitor> Driver<T> for MemoryFill {
 	fn accept(&self, visitor: &mut T) {
+		self.destination().pointer().accept(visitor);
+		self.size().accept(visitor);
 		self.value().accept(visitor);
-		self.n().accept(visitor);
 
 		visitor.visit_memory_fill(self);
 	}