From b2272f39bcd2e836b071e072c46ca9c6e16bcead Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Thei=C3=9Fen?= <alex.theissen@me.com>
Date: Wed, 21 Apr 2021 14:08:50 +0200
Subject: [PATCH] stack_height: `if` instruction should pop one value from the
 stack (#147)

* stack_height: 'if' instruction should pop one value from the stack

* Fix indentation
---
 src/stack_height/max_height.rs | 51 ++++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/src/stack_height/max_height.rs b/src/stack_height/max_height.rs
index 260345e..1bcf46b 100644
--- a/src/stack_height/max_height.rs
+++ b/src/stack_height/max_height.rs
@@ -201,6 +201,9 @@ pub(crate) fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, E
 			Block(ty) | Loop(ty) | If(ty) => {
 				let end_arity = if *ty == BlockType::NoResult { 0 } else { 1 };
 				let branch_arity = if let Loop(_) = *opcode { 0 } else { end_arity };
+				if let If(_) = *opcode {
+					stack.pop_values(1)?;
+				}
 				let height = stack.height();
 				stack.push_frame(Frame {
 					is_polymorphic: false,
@@ -546,4 +549,52 @@ mod tests {
 		let height = compute(0, &module).unwrap();
 		assert_eq!(height, 1);
 	}
+
+#[test]
+	fn breaks() {
+		let module = parse_wat(
+			r#"
+(module
+	(func $main
+		block (result i32)
+			block (result i32)
+				i32.const 99
+				br 1
+			end
+		end
+		drop
+	)
+)
+"#,
+		);
+
+		let height = compute(0, &module).unwrap();
+		assert_eq!(height, 1);
+	}
+
+#[test]
+	fn if_else_works() {
+		let module = parse_wat(
+			r#"
+(module
+	(func $main
+		i32.const 7
+		i32.const 1
+		if (result i32)
+			i32.const 42
+		else
+			i32.const 99
+		end
+		i32.const 97
+		drop
+		drop
+		drop
+	)
+)
+"#,
+		);
+
+		let height = compute(0, &module).unwrap();
+		assert_eq!(height, 3);
+	}
 }