From 2f7d7b564cd51a904aa3f3878393424609f15110 Mon Sep 17 00:00:00 2001
From: Rerumu <rerumu@protonmail.com>
Date: Sun, 19 Jun 2022 05:10:12 -0400
Subject: [PATCH] Improve call stack leak behavior

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

diff --git a/wasm-ast/src/builder.rs b/wasm-ast/src/builder.rs
index 8a41ffc..7591999 100644
--- a/wasm-ast/src/builder.rs
+++ b/wasm-ast/src/builder.rs
@@ -172,6 +172,12 @@ impl StatList {
 		self.stack.leak_into(&mut self.code, |_| true);
 	}
 
+	fn leak_pre_call(&mut self) {
+		self.stack.leak_into(&mut self.code, |v| {
+			v.has_global_read() || v.has_memory_read()
+		});
+	}
+
 	leak_on!(leak_local_write, Local);
 	leak_on!(leak_global_write, Global);
 	leak_on!(leak_memory_write, Memory);
@@ -452,7 +458,7 @@ impl<'a> Builder<'a> {
 		let arity = self.type_info.rel_arity_of(func);
 		let param_list = self.target.stack.pop_len(arity.num_param).collect();
 
-		self.target.leak_all();
+		self.target.leak_pre_call();
 
 		let result = self.target.stack.push_temporary(arity.num_result);
 
@@ -470,7 +476,7 @@ impl<'a> Builder<'a> {
 		let index = self.target.stack.pop();
 		let param_list = self.target.stack.pop_len(arity.num_param).collect();
 
-		self.target.leak_all();
+		self.target.leak_pre_call();
 
 		let result = self.target.stack.push_temporary(arity.num_result);
 
diff --git a/wasm-ast/src/stack.rs b/wasm-ast/src/stack.rs
index 0318509..5eef7bf 100644
--- a/wasm-ast/src/stack.rs
+++ b/wasm-ast/src/stack.rs
@@ -24,6 +24,14 @@ impl Slot {
 	pub fn has_read(&self, id: ReadType) -> bool {
 		self.read.contains(&id)
 	}
+
+	pub fn has_global_read(&self) -> bool {
+		self.read.iter().any(|r| matches!(r, ReadType::Global(_)))
+	}
+
+	pub fn has_memory_read(&self) -> bool {
+		self.read.iter().any(|r| matches!(r, ReadType::Memory(_)))
+	}
 }
 
 #[derive(Default)]