From dc1de1acd57e1c69cae2c713657fd6e019f35aa7 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Thu, 19 Sep 2024 18:55:52 +0300
Subject: [PATCH 1/2] feat(stack-limiter): support multi-value feature

---
 Cargo.toml                      |  3 ++-
 src/stack_limiter/max_height.rs | 12 ++++++++++--
 src/utils.rs                    |  2 +-
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 582d434..7c870b0 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,7 +22,8 @@ codegen-units = 1
 
 [dependencies]
 # parity-wasm = { version = "0.45.1", git = "https://github.com/gear-tech/parity-wasm", branch = "v0.45.1-sign-ext", default-features = false }
-parity-wasm = { package = "gear-wasm", version = "0.45.1", default-features = false }
+# TODO: add multi_value feature to gear-wasm as default, #[cfg(feature = "multi_value")] -> // #[cfg(feature = "multi_value")]
+parity-wasm = { package = "gear-wasm", version = "0.45.1", features = ["multi_value"], default-features = false }
 
 [dev-dependencies]
 binaryen = "0.12"
diff --git a/src/stack_limiter/max_height.rs b/src/stack_limiter/max_height.rs
index 6e5e1a3..9efb682 100644
--- a/src/stack_limiter/max_height.rs
+++ b/src/stack_limiter/max_height.rs
@@ -283,12 +283,20 @@ where
 		match opcode {
 			Nop => {},
 			Block(ty) | Loop(ty) | If(ty) => {
-				let end_arity = if *ty == BlockType::NoResult { 0 } else { 1 };
+				let (param_arity, end_arity) = match ty {
+					BlockType::NoResult => (0, 0),
+					BlockType::Value(_) => (0, 1),
+					BlockType::TypeIndex(x) => {
+						let Type::Function(ty) =
+							type_section.types().get(*x as usize).ok_or("Type not found")?;
+						(ty.params().len() as u32, ty.results().len() as u32)
+					},
+				};
 				let branch_arity = if let Loop(_) = *opcode { 0 } else { end_arity };
 				if let If(_) = *opcode {
 					stack.pop_values(1)?;
 				}
-				let height = stack.height();
+				let height = stack.height().checked_sub(param_arity).ok_or("Empty stack")?;
 				stack.push_frame(Frame {
 					is_polymorphic: false,
 					end_arity,
diff --git a/src/utils.rs b/src/utils.rs
index 30d0da3..7d21ca5 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -51,7 +51,7 @@ pub fn rewrite_sections_after_insertion(
 					*start_idx += inserted_count
 				},
 			Section::Name(s) =>
-				for functions in s.functions_mut() {
+				if let Some(functions) = s.functions_mut() {
 					*functions.names_mut() =
 						IndexMap::from_iter(functions.names().iter().map(|(mut idx, name)| {
 							if idx >= inserted_index {

From bddf38c65d40636c55f978905cf8d5f5129ce053 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Wed, 2 Oct 2024 10:22:08 +0300
Subject: [PATCH 2/2] add test with multi_value_if

---
 tests/diff.rs                                 |  1 +
 .../stack-height/multi_value_if.wat           | 52 +++++++++++++++++++
 .../fixtures/stack-height/multi_value_if.wat  | 13 +++++
 3 files changed, 66 insertions(+)
 create mode 100644 tests/expectations/stack-height/multi_value_if.wat
 create mode 100644 tests/fixtures/stack-height/multi_value_if.wat

diff --git a/tests/diff.rs b/tests/diff.rs
index aae4ccd..4a8062d 100644
--- a/tests/diff.rs
+++ b/tests/diff.rs
@@ -90,6 +90,7 @@ mod stack_height {
 	def_stack_height_test!(imports);
 	def_stack_height_test!(many_locals);
 	def_stack_height_test!(empty_functions);
+	def_stack_height_test!(multi_value_if);
 }
 
 mod gas {
diff --git a/tests/expectations/stack-height/multi_value_if.wat b/tests/expectations/stack-height/multi_value_if.wat
new file mode 100644
index 0000000..7673b97
--- /dev/null
+++ b/tests/expectations/stack-height/multi_value_if.wat
@@ -0,0 +1,52 @@
+(module
+  (type (;0;) (func (param i32) (result i32)))
+  (type (;1;) (func (param i32 i32) (result i32)))
+  (type (;2;) (func (result i32)))
+  (func $f (;0;) (type 0) (param $param i32) (result i32)
+    i32.const 1
+    i32.const 2
+    local.get $param
+    if (type 1) (param i32 i32) (result i32) ;; label = @1
+      i32.add
+    else
+      i32.sub
+    end
+  )
+  (func $main (;1;) (type 2) (result i32)
+    i32.const 0
+    global.get 0
+    i32.const 5
+    i32.add
+    global.set 0
+    global.get 0
+    i32.const 1024
+    i32.gt_u
+    if ;; label = @1
+      unreachable
+    end
+    call $f
+    global.get 0
+    i32.const 5
+    i32.sub
+    global.set 0
+  )
+  (func (;2;) (type 2) (result i32)
+    global.get 0
+    i32.const 10
+    i32.add
+    global.set 0
+    global.get 0
+    i32.const 1024
+    i32.gt_u
+    if ;; label = @1
+      unreachable
+    end
+    call $main
+    global.get 0
+    i32.const 10
+    i32.sub
+    global.set 0
+  )
+  (global (;0;) (mut i32) i32.const 0)
+  (export "main" (func 2))
+)
\ No newline at end of file
diff --git a/tests/fixtures/stack-height/multi_value_if.wat b/tests/fixtures/stack-height/multi_value_if.wat
new file mode 100644
index 0000000..8f66500
--- /dev/null
+++ b/tests/fixtures/stack-height/multi_value_if.wat
@@ -0,0 +1,13 @@
+(module
+  (func $f (param $param i32) (result i32)
+    (i32.const 1)
+    (i32.const 2)
+    (if (param i32 i32) (result i32) (local.get $param)
+      (then (i32.add))
+      (else (i32.sub))
+    )
+  )
+  (func $main (export "main") (result i32)
+    (call $f (i32.const 0))
+  )
+)