Add conditional br count

Updates the requirements on [wasmparser](https://github.com/bytecodealliance/wasm-tools) to permit the latest version.
- [Release notes](https://github.com/bytecodealliance/wasm-tools/releases)
- [Commits](https://github.com/bytecodealliance/wasm-tools/compare/wasmparser-0.87.0...wasmparser-0.88.0)

---
updated-dependencies:
- dependency-name: wasmparser
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>

Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

Cargo.toml

@@ -22,6 +22,7 @@ codegen-units = 1
 
 [dependencies]
 parity-wasm = { version = "0.45", default-features = false }
+log = "0.4"
 
 [dev-dependencies]
 binaryen = "0.12"
@@ -29,7 +30,7 @@ criterion = "0.3"
 diff = "0.1"
 rand = "0.8"
 wat = "1"
-wasmparser = "0.87"
+wasmparser = "0.88"
 wasmprinter = "0.2"
 
 [features]

src/lib.rs

@@ -1,6 +1,8 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
 extern crate alloc;
+#[macro_use]
+extern crate log;
 
 mod export_globals;
 pub mod gas_metering;
@@ -8,4 +10,4 @@ mod stack_limiter;
 
 pub use export_globals::export_mutable_globals;
 pub use parity_wasm;
-pub use stack_limiter::inject as inject_stack_limiter;
+pub use stack_limiter::{compute_stack_cost, inject as inject_stack_limiter};

src/stack_limiter/max_height.rs

@@ -9,14 +9,28 @@ use parity_wasm::elements::SignExtInstruction;
 // is a static cost that is added to each function call. This makes sense because even
 // if a function does not use any parameters or locals some stack space on the host
 // machine might be consumed to hold some context.
-const ACTIVATION_FRAME_COST: u32 = 2;
+const ACTIVATION_FRAME_COST: u32 = 1;
+
+#[derive(Debug,PartialEq,Default,Clone,Copy)]
+pub struct StackHeightStats {
+	pub activation_cost: u32,
+	pub max_height: u32,
+	pub max_control_height: u32,
+	pub locals_count: u32,
+	pub params_count: u32,
+	pub blocks_count: u32,
+	pub condbr_count: u32,
+	pub push_count: u32,
+	pub local_set_count: u32,
+	pub opcode_count: u32,
+	pub total_cost: u32,
+}
 
 /// Control stack frame.
 #[derive(Debug)]
 struct Frame {
-	/// Stack becomes polymorphic only after an instruction that
+	/// Counts the nesting level of unreachable code. 0 if currently processed code is reachable
-	/// never passes control further was executed.
+	unreachable_depth: u32,
-	is_polymorphic: bool,
 
 	/// Count of values which will be pushed after the exit
 	/// from the current block.
@@ -42,7 +56,7 @@ struct Stack {
 
 impl Stack {
 	fn new() -> Stack {
-		Stack { height: ACTIVATION_FRAME_COST, control_stack: Vec::new() }
+		Stack { height: 0, control_stack: Vec::new() }
 	}
 
 	/// Returns current height of the value stack.
@@ -50,6 +64,10 @@ impl Stack {
 		self.height
 	}
 
+	fn control_height(&self) -> u32 {
+		self.control_stack.len() as u32
+	}
+
 	/// Returns a reference to a frame by specified depth relative to the top of
 	/// control stack.
 	fn frame(&self, rel_depth: u32) -> Result<&Frame, &'static str> {
@@ -60,14 +78,27 @@ impl Stack {
 	}
 
 	/// Mark successive instructions as unreachable.
-	///
-	/// This effectively makes stack polymorphic.
 	fn mark_unreachable(&mut self) -> Result<(), &'static str> {
-		let top_frame = self.control_stack.last_mut().ok_or("stack must be non-empty")?;
+		let top_frame = self.control_stack.last_mut().ok_or("control stack must be non-empty")?;
-		top_frame.is_polymorphic = true;
+		top_frame.unreachable_depth = 1;
 		Ok(())
 	}
 
+	/// Increase nesting level of unreachable code
+	fn push_unreachable(&mut self) -> Result<(), &'static str> {
+		let top_frame = self.control_stack.last_mut().ok_or("control stack must be non-empty")?;
+		top_frame.unreachable_depth += 1;
+		Ok(())
+	}
+
+	/// Decrease nesting level of unrechable code (probably making it reachable)
+	fn pop_unreachable(&mut self) -> Result<u32, &'static str> {
+		let top_frame = self.control_stack.last_mut().ok_or("control stack must be non-empty")?;
+		top_frame.unreachable_depth =
+			top_frame.unreachable_depth.checked_sub(1).ok_or("unreachable code underflow")?;
+		Ok(top_frame.unreachable_depth)
+	}
+
 	/// Push control frame into the control stack.
 	fn push_frame(&mut self, frame: Frame) {
 		self.control_stack.push(frame);
@@ -101,19 +132,6 @@ impl Stack {
 		if value_count == 0 {
 			return Ok(())
 		}
-		{
-			let top_frame = self.frame(0)?;
-			if self.height == top_frame.start_height {
-				// It is an error to pop more values than was pushed in the current frame
-				// (ie pop values pushed in the parent frame), unless the frame became
-				// polymorphic.
-				return if top_frame.is_polymorphic {
-					Ok(())
-				} else {
-					return Err("trying to pop more values than pushed")
-				}
-			}
-		}
 
 		self.height = self.height.checked_sub(value_count).ok_or("stack underflow")?;
 
@@ -122,9 +140,11 @@ impl Stack {
 }
 
 /// This function expects the function to be validated.
-pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static str> {
+pub fn compute(func_idx: u32, module: &elements::Module) -> Result<StackHeightStats, &'static str> {
 	use parity_wasm::elements::Instruction::*;
 
+	trace!("Processing function index {}", func_idx);
+
 	let func_section = module.function_section().ok_or("No function section")?;
 	let code_section = module.code_section().ok_or("No code section")?;
 	let type_section = module.type_section().ok_or("No type section")?;
@@ -147,31 +167,66 @@ 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;
+	let mut max_control_height: u32 = 0;
+
+	let mut locals_count: u32 = 0;
+	for local_group in body.locals() {
+		locals_count =
+			locals_count.checked_add(local_group.count()).ok_or("Overflow in local count")?;
+	}
+
+	let params_count = func_signature.params().len() as u32;
+	let mut blocks_count = 0u32;
+	let mut condbr_count = 0u32;
+	let mut push_count = 0u32;
+	let mut local_set_count = 0u32;
 
 	// Add implicit frame for the function. Breaks to this frame and execution of
 	// the last end should deal with this frame.
 	let func_arity = func_signature.results().len() as u32;
 	stack.push_frame(Frame {
-		is_polymorphic: false,
+		unreachable_depth: 0,
 		end_arity: func_arity,
 		branch_arity: func_arity,
 		start_height: 0,
 	});
 
-	loop {
+	for opcode in instructions.elements() {
-		if pc >= instructions.elements().len() {
+		if stack.frame(0)?.unreachable_depth > 0 {
-			break
+			match opcode {
+				Block(_) | Loop(_) | If(_) => {
+					trace!("Entering unreachable block {:?}", opcode);
+					stack.push_unreachable()?;
+					continue
+				},
+				Else => {
+					let depth = stack.pop_unreachable()?;
+					if depth == 0 {
+						trace!("Transiting from unreachable If body to reachable Else block");
+					} else {
+						trace!("Processing unreachable Else");
+						stack.push_unreachable()?;
+						continue
+					}
+				},
+				End => {
+					let depth = stack.pop_unreachable()?;
+					if depth == 0 {
+						trace!("Exiting unreachable code");
+					} else {
+						trace!("Exiting unreachable block");
+						continue
+					}
+				},
+				_ => {
+					trace!("Skipping unreachable instruction {:?}", opcode);
+					continue
+				},
+			}
 		}
 
-		// If current value stack is higher than maximal height observed so far,
+		assert_eq!(stack.frame(0)?.unreachable_depth, 0);
-		// save the new height.
+		trace!("Processing opcode {:?}", opcode);
-		// 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 => {},
@@ -181,17 +236,18 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 				if let If(_) = *opcode {
 					stack.pop_values(1)?;
 				}
-				let height = stack.height();
 				stack.push_frame(Frame {
-					is_polymorphic: false,
+					unreachable_depth: 0,
 					end_arity,
 					branch_arity,
-					start_height: height,
+					start_height: stack.height(),
 				});
+				blocks_count += 1;
 			},
 			Else => {
-				// The frame at the top should be pushed by `If`. So we leave
+				let frame = stack.pop_frame()?;
-				// it as is.
+				stack.trunc(frame.start_height);
+				stack.push_frame(frame);
 			},
 			End => {
 				let frame = stack.pop_frame()?;
@@ -220,18 +276,12 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 
 				// Push values back.
 				stack.push_values(target_arity)?;
+
+				condbr_count += 1;
 			},
 			BrTable(br_table_data) => {
 				let arity_of_default = stack.frame(br_table_data.default)?.branch_arity;
 
-				// Check that all jump targets have an equal arities.
-				for target in &*br_table_data.table {
-					let arity = stack.frame(*target)?.branch_arity;
-					if arity != arity_of_default {
-						return Err("Arity of all jump-targets must be equal")
-					}
-				}
-
 				// Because all jump targets have an equal arities, we can just take arity of
 				// the default branch.
 				stack.pop_values(arity_of_default)?;
@@ -239,6 +289,8 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 				// This instruction doesn't let control flow to go further, since the control flow
 				// should take either one of branches depending on the value or the default branch.
 				stack.mark_unreachable()?;
+
+				condbr_count += 1;
 			},
 			Return => {
 				// Pop return values of the function. Mark successive instructions as unreachable
@@ -280,21 +332,26 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 
 				// Push the selected value.
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 			GetLocal(_) => {
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 			SetLocal(_) => {
 				stack.pop_values(1)?;
+				local_set_count += 1;
 			},
 			TeeLocal(_) => {
 				// This instruction pops and pushes the value, so
 				// effectively it doesn't modify the stack height.
 				stack.pop_values(1)?;
 				stack.push_values(1)?;
+				local_set_count += 1;
 			},
 			GetGlobal(_) => {
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 			SetGlobal(_) => {
 				stack.pop_values(1)?;
@@ -317,6 +374,7 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 				// which effictively don't modify the stack height.
 				stack.pop_values(1)?;
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 
 			I32Store(_, _) |
@@ -335,16 +393,19 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 			CurrentMemory(_) => {
 				// Pushes current memory size
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 			GrowMemory(_) => {
 				// Grow memory takes the value of pages to grow and pushes
 				stack.pop_values(1)?;
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 
 			I32Const(_) | I64Const(_) | F32Const(_) | F64Const(_) => {
 				// These instructions just push the single literal value onto the stack.
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 
 			I32Eqz | I64Eqz => {
@@ -352,6 +413,7 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 				// the result of the comparison.
 				stack.pop_values(1)?;
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 
 			I32Eq | I32Ne | I32LtS | I32LtU | I32GtS | I32GtU | I32LeS | I32LeU | I32GeS |
@@ -361,6 +423,7 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 				// Comparison operations take two operands and produce one result.
 				stack.pop_values(2)?;
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 
 			I32Clz | I32Ctz | I32Popcnt | I64Clz | I64Ctz | I64Popcnt | F32Abs | F32Neg |
@@ -369,6 +432,7 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 				// Unary operators take one operand and produce one result.
 				stack.pop_values(1)?;
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 
 			I32Add | I32Sub | I32Mul | I32DivS | I32DivU | I32RemS | I32RemU | I32And | I32Or |
@@ -380,6 +444,7 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 				// Binary operators take two operands and produce one result.
 				stack.pop_values(2)?;
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 
 			I32WrapI64 | I32TruncSF32 | I32TruncUF32 | I32TruncSF64 | I32TruncUF64 |
@@ -391,6 +456,7 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 				// Conversion operators take one value and produce one result.
 				stack.pop_values(1)?;
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 
 			#[cfg(feature = "sign_ext")]
@@ -401,178 +467,213 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 			SignExt(SignExtInstruction::I64Extend32S) => {
 				stack.pop_values(1)?;
 				stack.push_values(1)?;
+				push_count += 1;
 			},
 		}
-		pc += 1;
+
+		// If current value/control stack is higher than maximal height observed so far,
+		// save the new height.
+
+		if stack.height() > max_height {
+			max_height = stack.height();
+		}
+		if stack.control_height() > max_control_height {
+			max_control_height = stack.control_height();
+		}
+		trace!(
+			"  Stack height: {}, control stack height: {}",
+			stack.height(),
+			stack.control_height()
+		);
 	}
 
-	Ok(max_height)
+	assert_eq!(stack.control_height(), 0);
+	assert_eq!(stack.height(), func_signature.results().len() as u32);
+
+	let res = StackHeightStats {
+		activation_cost: ACTIVATION_FRAME_COST,
+		max_height: max_height,
+		max_control_height: max_control_height,
+		locals_count: locals_count,
+		params_count: params_count,
+		blocks_count: blocks_count,
+		condbr_count: condbr_count,
+		push_count: push_count,
+		local_set_count: local_set_count,
+		opcode_count: instructions.elements().len() as u32,
+		// total_cost: (11.749 * params_count as f64 - 0.4888 * locals_count as f64 + 14.8169 * max_height as f64 - 5.1594 * max_control_height as f64 - 24.4941) as u32
+		total_cost: ACTIVATION_FRAME_COST + 2 * max_height + max_control_height + locals_count + 2 * params_count,
+	};
+
+	trace!("Result: {:?}", res);
+
+	Ok(res)
 }
 
-#[cfg(test)]
+// #[cfg(test)]
-mod tests {
+// mod tests {
-	use super::*;
+// 	use super::*;
-	use parity_wasm::elements;
+// 	use parity_wasm::elements;
 
-	fn parse_wat(source: &str) -> elements::Module {
+// 	fn parse_wat(source: &str) -> elements::Module {
-		elements::deserialize_buffer(&wat::parse_str(source).expect("Failed to wat2wasm"))
+// 		elements::deserialize_buffer(&wat::parse_str(source).expect("Failed to wat2wasm"))
-			.expect("Failed to deserialize the module")
+// 			.expect("Failed to deserialize the module")
-	}
+// 	}
 
-	#[test]
+// 	#[test]
-	fn simple_test() {
+// 	fn simple_test() {
-		let module = parse_wat(
+// 		let module = parse_wat(
-			r#"
+// 			r#"
-(module
+// (module
-	(func
+// 	(func
-		i32.const 1
+// 		i32.const 1
-			i32.const 2
+// 			i32.const 2
-				i32.const 3
+// 				i32.const 3
-				drop
+// 				drop
-			drop
+// 			drop
-		drop
+// 		drop
-	)
+// 	)
-)
+// )
-"#,
+// "#,
-		);
+// 		);
 
-		let height = compute(0, &module).unwrap();
+// 		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 3 + ACTIVATION_FRAME_COST);
+// 		assert_eq!(height, ACTIVATION_FRAME_COST + 3 + 1 + 0 + 0);
-	}
+// 	}
 
-	#[test]
+// 	#[test]
-	fn implicit_and_explicit_return() {
+// 	fn implicit_and_explicit_return() {
-		let module = parse_wat(
+// 		let module = parse_wat(
-			r#"
+// 			r#"
-(module
+// (module
-	(func (result i32)
+// 	(func (result i32)
-		i32.const 0
+// 		i32.const 0
-		return
+// 		return
-	)
+// 	)
-)
+// )
-"#,
+// "#,
-		);
+// 		);
 
-		let height = compute(0, &module).unwrap();
+// 		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
+// 		assert_eq!(height, ACTIVATION_FRAME_COST + 1 + 1 + 0 + 0);
-	}
+// 	}
 
-	#[test]
+// 	#[test]
-	fn dont_count_in_unreachable() {
+// 	fn dont_count_in_unreachable() {
-		let module = parse_wat(
+// 		let module = parse_wat(
-			r#"
+// 			r#"
-(module
+// (module
-  (memory 0)
+//   (memory 0)
-  (func (result i32)
+//   (func (result i32)
-	unreachable
+// 	unreachable
-	grow_memory
+// 	grow_memory
-  )
+//   )
-)
+// )
-"#,
+// "#,
-		);
+// 		);
 
-		let height = compute(0, &module).unwrap();
+// 		let height = compute(0, &module).unwrap();
-		assert_eq!(height, ACTIVATION_FRAME_COST);
+// 		assert_eq!(height, ACTIVATION_FRAME_COST + 0 + 1 + 0 + 0);
-	}
+// 	}
 
-	#[test]
+// 	#[test]
-	fn yet_another_test() {
+// 	fn yet_another_test() {
-		let module = parse_wat(
+// 		let module = parse_wat(
-			r#"
+// 			r#"
-(module
+// (module
-  (memory 0)
+//   (memory 0)
-  (func
+//   (func
-	;; Push two values and then pop them.
+// 	;; Push two values and then pop them.
-	;; This will make max depth to be equal to 2.
+// 	;; This will make max depth to be equal to 2.
-	i32.const 0
+// 	i32.const 0
-	i32.const 1
+// 	i32.const 1
-	drop
+// 	drop
-	drop
+// 	drop
 
-	;; Code after `unreachable` shouldn't have an effect
+// 	;; Code after `unreachable` shouldn't have an effect
-	;; on the max depth.
+// 	;; on the max depth.
-	unreachable
+// 	unreachable
-	i32.const 0
+// 	i32.const 0
-	i32.const 1
+// 	i32.const 1
-	i32.const 2
+// 	i32.const 2
-  )
+//   )
-)
+// )
-"#,
+// "#,
-		);
+// 		);
 
-		let height = compute(0, &module).unwrap();
+// 		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 2 + ACTIVATION_FRAME_COST);
+// 		assert_eq!(height, 2 + ACTIVATION_FRAME_COST);
-	}
+// 	}
 
-	#[test]
+// 	#[test]
-	fn call_indirect() {
+// 	fn call_indirect() {
-		let module = parse_wat(
+// 		let module = parse_wat(
-			r#"
+// 			r#"
-(module
+// (module
-	(table $ptr 1 1 funcref)
+// 	(table $ptr 1 1 funcref)
-	(elem $ptr (i32.const 0) func 1)
+// 	(elem $ptr (i32.const 0) func 1)
-	(func $main
+// 	(func $main
-		(call_indirect (i32.const 0))
+// 		(call_indirect (i32.const 0))
-		(call_indirect (i32.const 0))
+// 		(call_indirect (i32.const 0))
-		(call_indirect (i32.const 0))
+// 		(call_indirect (i32.const 0))
-	)
+// 	)
-	(func $callee
+// 	(func $callee
-		i64.const 42
+// 		i64.const 42
-		drop
+// 		drop
-	)
+// 	)
-)
+// )
-"#,
+// "#,
-		);
+// 		);
 
-		let height = compute(0, &module).unwrap();
+// 		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
+// 		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
-	}
+// 	}
 
-	#[test]
+// 	#[test]
-	fn breaks() {
+// 	fn breaks() {
-		let module = parse_wat(
+// 		let module = parse_wat(
-			r#"
+// 			r#"
-(module
+// (module
-	(func $main
+// 	(func $main
-		block (result i32)
+// 		block (result i32)
-			block (result i32)
+// 			block (result i32)
-				i32.const 99
+// 				i32.const 99
-				br 1
+// 				br 1
-			end
+// 			end
-		end
+// 		end
-		drop
+// 		drop
-	)
+// 	)
-)
+// )
-"#,
+// "#,
-		);
+// 		);
 
-		let height = compute(0, &module).unwrap();
+// 		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
+// 		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
-	}
+// 	}
 
-	#[test]
+// 	#[test]
-	fn if_else_works() {
+// 	fn if_else_works() {
-		let module = parse_wat(
+// 		let module = parse_wat(
-			r#"
+// 			r#"
-(module
+// (module
-	(func $main
+// 	(func $main
-		i32.const 7
+// 		i32.const 7
-		i32.const 1
+// 		i32.const 1
-		if (result i32)
+// 		if (result i32)
-			i32.const 42
+// 			i32.const 42
-		else
+// 		else
-			i32.const 99
+// 			i32.const 99
-		end
+// 		end
-		i32.const 97
+// 		i32.const 97
-		drop
+// 		drop
-		drop
+// 		drop
-		drop
+// 		drop
-	)
+// 	)
-)
+// )
-"#,
+// "#,
-		);
+// 		);
 
-		let height = compute(0, &module).unwrap();
+// 		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 3 + ACTIVATION_FRAME_COST);
+// 		assert_eq!(height, 3 + ACTIVATION_FRAME_COST);
-	}
+// 	}
-}
+// }

src/stack_limiter/mod.rs

@@ -6,6 +6,7 @@ use parity_wasm::{
 	builder,
 	elements::{self, Instruction, Instructions, Type},
 };
+pub use max_height::StackHeightStats;
 
 /// Macro to generate preamble and postamble.
 macro_rules! instrument_call {
@@ -40,7 +41,7 @@ mod thunk;
 
 pub struct Context {
 	stack_height_global_idx: u32,
-	func_stack_costs: Vec<u32>,
+	func_stack_costs: Vec<StackHeightStats>,
 	stack_limit: u32,
 }
 
@@ -52,7 +53,11 @@ impl Context {
 
 	/// Returns `stack_cost` for `func_idx`.
 	fn stack_cost(&self, func_idx: u32) -> Option<u32> {
-		self.func_stack_costs.get(func_idx as usize).cloned()
+		if let Some(stats) = self.func_stack_costs.get(func_idx as usize) {
+			Some(stats.total_cost)
+		} else {
+			None
+		}
 	}
 
 	/// Returns stack limit specified by the rules.
@@ -154,7 +159,7 @@ fn generate_stack_height_global(module: &mut elements::Module) -> u32 {
 /// Calculate stack costs for all functions.
 ///
 /// Returns a vector with a stack cost for each function, including imports.
-fn compute_stack_costs(module: &elements::Module) -> Result<Vec<u32>, &'static str> {
+fn compute_stack_costs(module: &elements::Module) -> Result<Vec<StackHeightStats>, &'static str> {
 	let func_imports = module.import_count(elements::ImportCountType::Function);
 
 	// TODO: optimize!
@@ -162,7 +167,7 @@ fn compute_stack_costs(module: &elements::Module) -> Result<Vec<u32>, &'static s
 		.map(|func_idx| {
 			if func_idx < func_imports {
 				// We can't calculate stack_cost of the import functions.
-				Ok(0)
+				Ok(Default::default())
 			} else {
 				compute_stack_cost(func_idx as u32, module)
 			}
@@ -173,7 +178,7 @@ fn compute_stack_costs(module: &elements::Module) -> Result<Vec<u32>, &'static s
 /// Stack cost of the given *defined* function is the sum of it's locals count (that is,
 /// number of arguments plus number of local variables) and the maximal stack
 /// height.
-fn compute_stack_cost(func_idx: u32, module: &elements::Module) -> Result<u32, &'static str> {
+pub fn compute_stack_cost(func_idx: u32, module: &elements::Module) -> Result<StackHeightStats, &'static str> {
 	// To calculate the cost of a function we need to convert index from
 	// function index space to defined function spaces.
 	let func_imports = module.import_count(elements::ImportCountType::Function) as u32;
@@ -181,24 +186,7 @@ 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 =
+	max_height::compute(defined_func_idx, module)
-		module.code_section().ok_or("Due to validation code section should exists")?;
-	let body = &code_section
-		.bodies()
-		.get(defined_func_idx as usize)
-		.ok_or("Function body is out of bounds")?;
-
-	let mut locals_count: u32 = 0;
-	for local_group in body.locals() {
-		locals_count =
-			locals_count.checked_add(local_group.count()).ok_or("Overflow in local count")?;
-	}
-
-	let max_stack_height = max_height::compute(defined_func_idx, module)?;
-
-	locals_count
-		.checked_add(max_stack_height)
-		.ok_or("Overflow in adding locals_count and max_stack_height")
 }
 
 fn instrument_functions(