From 5f47092a65c71b2286d60869962ebab199598176 Mon Sep 17 00:00:00 2001
From: Rerumu <rerumu@protonmail.com>
Date: Sun, 19 Jun 2022 05:08:11 -0400
Subject: [PATCH] Refactor stack leak code

---
 wasm-ast/src/builder.rs | 16 ++------------
 wasm-ast/src/stack.rs   | 46 ++++++++++++++++++++++-------------------
 2 files changed, 27 insertions(+), 35 deletions(-)

diff --git a/wasm-ast/src/builder.rs b/wasm-ast/src/builder.rs
index 92082d3..8a41ffc 100644
--- a/wasm-ast/src/builder.rs
+++ b/wasm-ast/src/builder.rs
@@ -18,11 +18,7 @@ macro_rules! leak_on {
 		fn $name(&mut self, id: usize) {
 			let read = ReadType::$variant(id);
 
-			for i in 0..self.stack.len() {
-				if self.stack.has_read_at(i, read) {
-					self.leak_at(i);
-				}
-			}
+			self.stack.leak_into(&mut self.code, |v| v.has_read(read))
 		}
 	};
 }
@@ -172,16 +168,8 @@ impl StatList {
 		Self::default()
 	}
 
-	fn leak_at(&mut self, index: usize) {
-		if let Some(set) = self.stack.leak_at(index) {
-			self.code.push(set);
-		}
-	}
-
 	fn leak_all(&mut self) {
-		for i in 0..self.stack.len() {
-			self.leak_at(i);
-		}
+		self.stack.leak_into(&mut self.code, |_| true);
 	}
 
 	leak_on!(leak_local_write, Local);
diff --git a/wasm-ast/src/stack.rs b/wasm-ast/src/stack.rs
index 606a8b3..0318509 100644
--- a/wasm-ast/src/stack.rs
+++ b/wasm-ast/src/stack.rs
@@ -20,6 +20,10 @@ impl Slot {
 	fn is_temporary(&self, id: usize) -> bool {
 		matches!(self.data, Expression::GetTemporary(ref v) if v.var == id)
 	}
+
+	pub fn has_read(&self, id: ReadType) -> bool {
+		self.read.contains(&id)
+	}
 }
 
 #[derive(Default)]
@@ -97,10 +101,6 @@ impl Stack {
 		range
 	}
 
-	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 {
@@ -115,24 +115,28 @@ impl Stack {
 
 	// 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> {
-		let old = &mut self.var_list[index];
-		let var = self.previous + index;
+	pub fn leak_into<P>(&mut self, code: &mut Vec<Statement>, predicate: P)
+	where
+		P: Fn(&Slot) -> bool,
+	{
+		for (i, old) in self.var_list.iter_mut().enumerate() {
+			let var = self.previous + i;
 
-		if old.is_temporary(var) {
-			return None;
+			if old.is_temporary(var) || !predicate(old) {
+				continue;
+			}
+
+			old.read.clear();
+
+			let get = Expression::GetTemporary(GetTemporary { var });
+			let set = Statement::SetTemporary(SetTemporary {
+				var,
+				value: std::mem::replace(&mut old.data, get),
+			});
+
+			self.capacity = self.capacity.max(var + 1);
+
+			code.push(set);
 		}
-
-		old.read.clear();
-
-		let get = Expression::GetTemporary(GetTemporary { var });
-		let set = Statement::SetTemporary(SetTemporary {
-			var,
-			value: std::mem::replace(&mut old.data, get),
-		});
-
-		self.capacity = self.capacity.max(var + 1);
-
-		Some(set)
 	}
 }