From 19760b58355808e4d9ba0170e41271b337d0dd6b Mon Sep 17 00:00:00 2001
From: Sergey Shulepov <s.pepyakin@gmail.com>
Date: Wed, 9 Jun 2021 19:22:40 +0200
Subject: [PATCH] Fix counting of the local variables

The code assumed that the number of `Local` and number of locals is the
same thing. In reality though it is not. `Local` actually represents a
group of locals with the same type. The group can declare more than one
variable and the number of them is returned by `Local::count`.

In this PR we acknowledge this fact. Along the way we add a checked
arithmetic for locals_count and max_stack_height summation.
---
 src/stack_height/mod.rs                       | 11 ++++++++--
 tests/diff.rs                                 |  1 +
 .../expectations/stack-height/many_locals.wat | 21 +++++++++++++++++++
 tests/fixtures/stack-height/many_locals.wat   | 10 +++++++++
 4 files changed, 41 insertions(+), 2 deletions(-)
 create mode 100644 tests/expectations/stack-height/many_locals.wat
 create mode 100644 tests/fixtures/stack-height/many_locals.wat

diff --git a/src/stack_height/mod.rs b/src/stack_height/mod.rs
index f841864..64bdb01 100644
--- a/src/stack_height/mod.rs
+++ b/src/stack_height/mod.rs
@@ -199,7 +199,13 @@ fn compute_stack_cost(func_idx: u32, module: &elements::Module) -> Result<u32, E
 		.bodies()
 		.get(defined_func_idx as usize)
 		.ok_or_else(|| Error("Function body is out of bounds".into()))?;
-	let locals_count = body.locals().len() as u32;
+
+	let mut locals_count: u32 = 0;
+	for local_group in body.locals() {
+		locals_count = locals_count
+			.checked_add(local_group.count())
+			.ok_or_else(|| Error("Overflow in local count".into()))?;
+	}
 
 	let max_stack_height =
 		max_height::compute(
@@ -207,7 +213,8 @@ fn compute_stack_cost(func_idx: u32, module: &elements::Module) -> Result<u32, E
 			module
 		)?;
 
-	Ok(locals_count + max_stack_height)
+	locals_count.checked_add(max_stack_height)
+		.ok_or_else(|| Error("Overflow in adding locals_count and max_stack_height".into()))
 }
 
 fn instrument_functions(ctx: &mut Context, module: &mut elements::Module) -> Result<(), Error> {
diff --git a/tests/diff.rs b/tests/diff.rs
index b57a4bb..4b41db2 100644
--- a/tests/diff.rs
+++ b/tests/diff.rs
@@ -93,6 +93,7 @@ mod stack_height {
 	def_stack_height_test!(table);
 	def_stack_height_test!(global);
 	def_stack_height_test!(imports);
+	def_stack_height_test!(many_locals);
 }
 
 mod gas {
diff --git a/tests/expectations/stack-height/many_locals.wat b/tests/expectations/stack-height/many_locals.wat
new file mode 100644
index 0000000..e5fee38
--- /dev/null
+++ b/tests/expectations/stack-height/many_locals.wat
@@ -0,0 +1,21 @@
+(module
+  (type (;0;) (func))
+  (func (;0;) (type 0)
+    (local i64 i64 i32))
+  (func (;1;) (type 0)
+    global.get 0
+    i32.const 3
+    i32.add
+    global.set 0
+    global.get 0
+    i32.const 1024
+    i32.gt_u
+    if  ;; label = @1
+      unreachable
+    end
+    call 0
+    global.get 0
+    i32.const 3
+    i32.sub
+    global.set 0)
+  (global (;0;) (mut i32) (i32.const 0)))
diff --git a/tests/fixtures/stack-height/many_locals.wat b/tests/fixtures/stack-height/many_locals.wat
new file mode 100644
index 0000000..289ed31
--- /dev/null
+++ b/tests/fixtures/stack-height/many_locals.wat
@@ -0,0 +1,10 @@
+(module
+	(func $one-group-many-locals
+		(local i64) (local i64) (local i32)
+	)
+	(func $main
+		(call
+			$one-group-many-locals
+		)
+  	)
+)