Fix param count and rearrange code

src/stack_limiter/max_height.rs

@@ -147,7 +147,6 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 
 	let mut stack = Stack::new();
 	let mut max_height: u32 = 0;
-	let mut pc = 0;
 
 	// Add implicit frame for the function. Breaks to this frame and execution of
 	// the last end should deal with this frame.
@@ -159,19 +158,7 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 		start_height: 0,
 	});
 
-	loop {
+	for opcode in instructions.elements() {
-		if pc >= instructions.elements().len() {
-			break
-		}
-
-		// If current value stack is higher than maximal height observed so far,
-		// save the new height.
-		// However, we don't increase maximal value in unreachable code.
-		if stack.height() > max_height && !stack.frame(0)?.is_polymorphic {
-			max_height = stack.height();
-		}
-
-		let opcode = &instructions.elements()[pc];
 
 		match opcode {
 			Nop => {},
@@ -403,7 +390,13 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 				stack.push_values(1)?;
 			},
 		}
-		pc += 1;
+
+		// If current value stack is higher than maximal height observed so far,
+		// save the new height.
+		// However, we don't increase maximal value in unreachable code.
+		if stack.height() > max_height && !stack.frame(0)?.is_polymorphic {
+			max_height = stack.height();
+		}
 	}
 
 	Ok(max_height)

src/stack_limiter/mod.rs

@@ -181,8 +181,19 @@ fn compute_stack_cost(func_idx: u32, module: &elements::Module) -> Result<u32, &
 		.checked_sub(func_imports)
 		.ok_or("This should be a index of a defined function")?;
 
-	let code_section =
+	let func_section = module.function_section().ok_or("No function section")?;
-		module.code_section().ok_or("Due to validation code section should exists")?;
+	let code_section = module.code_section().ok_or("No code section")?;
+	let type_section = module.type_section().ok_or("No type section")?;
+
+	let func_sig_idx = func_section
+		.entries()
+		.get(func_idx as usize)
+		.ok_or("Function is not found in func section")?
+		.type_ref();
+	let Type::Function(func_signature) = type_section
+		.types()
+		.get(func_sig_idx as usize)
+		.ok_or("Function is not found in func section")?;
 	let body = &code_section
 		.bodies()
 		.get(defined_func_idx as usize)
@@ -193,6 +204,7 @@ fn compute_stack_cost(func_idx: u32, module: &elements::Module) -> Result<u32, &
 		locals_count =
 			locals_count.checked_add(local_group.count()).ok_or("Overflow in local count")?;
 	}
+	locals_count = locals_count.checked_add(func_signature.params().len() as u32).ok_or("Overflow in parameter count")?;
 
 	let max_stack_height = max_height::compute(defined_func_idx, module)?;