Merge pull request #152 from paritytech/ser-fix-locals

Fix counting of the local variables

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> {

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 {

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)))

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
+		)
+  	)
+)