Backward compatibility and tracing

CHANGELOG.md

@@ -16,17 +16,7 @@ The interface provided to smart contracts will adhere to semver with one excepti
 major version bumps will be backwards compatible with regard to already deployed contracts.
 In other words: Upgrading this pallet will not break pre-existing contracts.
 
-## [v0.3.0]
-
-### Changed
-
-- Use 64bit arithmetic for per-block gas counter
-[#30](https://github.com/paritytech/wasm-instrument/pull/30)
-
 ## [v0.2.0] 2022-06-06
-
-### Changed
-
 - Adjust debug information (if already parsed) when injecting gas metering
 [#16](https://github.com/paritytech/wasm-instrument/pull/16)
 

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "wasm-instrument"
-version = "0.3.0"
+version = "0.2.0"
 edition = "2021"
 rust-version = "1.56.1"
 authors = ["Parity Technologies <admin@parity.io>"]
@@ -22,18 +22,21 @@ codegen-units = 1
 
 [dependencies]
 parity-wasm = { version = "0.45", default-features = false }
+log = { version = "0.4.8", default-features = false, optional = true }
+test-log = { version = "0.2", optional = true }
+env_logger = { version = "0.9", optional = true }
 
 [dev-dependencies]
 binaryen = "0.12"
-criterion = "0.4"
+criterion = "0.3"
 diff = "0.1"
-pretty_assertions = "1"
 rand = "0.8"
 wat = "1"
-wasmparser = "0.92"
+wasmparser = "0.87"
 wasmprinter = "0.2"
 
 [features]
 default = ["std"]
 std = ["parity-wasm/std"]
 sign_ext = ["parity-wasm/sign_ext"]
+trace-log = ["dep:log", "dep:test-log", "dep:env_logger"]

src/gas_metering/mod.rs

@@ -105,7 +105,7 @@ impl Rules for ConstantCostRules {
 /// imported gas metering function.
 ///
 /// The output module imports a function "gas" from the specified module with type signature
-/// [i64] -> []. The argument is the amount of gas required to continue execution. The external
+/// [i32] -> []. The argument is the amount of gas required to continue execution. The external
 /// function is meant to keep track of the total amount of gas used and trap or otherwise halt
 /// execution of the runtime if the gas usage exceeds some allowed limit.
 ///
@@ -144,7 +144,7 @@ pub fn inject<R: Rules>(
 	// Injecting gas counting external
 	let mut mbuilder = builder::from_module(module);
 	let import_sig =
-		mbuilder.push_signature(builder::signature().with_param(ValueType::I64).build_sig());
+		mbuilder.push_signature(builder::signature().with_param(ValueType::I32).build_sig());
 
 	mbuilder.push_import(
 		builder::import()
@@ -284,7 +284,7 @@ struct MeteredBlock {
 	/// Index of the first instruction (aka `Opcode`) in the block.
 	start_pos: usize,
 	/// Sum of costs of all instructions until end of the block.
-	cost: u64,
+	cost: u32,
 }
 
 /// Counter is used to manage state during the gas metering algorithm implemented by
@@ -375,8 +375,7 @@ impl Counter {
 				.expect("last_index is greater than 0; last_index is stack size - 1; qed");
 			let prev_metered_block = &mut prev_control_block.active_metered_block;
 			if closing_metered_block.start_pos == prev_metered_block.start_pos {
-				prev_metered_block.cost =
+				prev_metered_block.cost += closing_metered_block.cost;
-					prev_metered_block.cost.checked_add(closing_metered_block.cost).ok_or(())?;
 				return Ok(())
 			}
 		}
@@ -426,7 +425,7 @@ impl Counter {
 	/// Increment the cost of the current block by the specified value.
 	fn increment(&mut self, val: u32) -> Result<(), ()> {
 		let top_block = self.active_metered_block()?;
-		top_block.cost = top_block.cost.checked_add(val.into()).ok_or(())?;
+		top_block.cost = top_block.cost.checked_add(val).ok_or(())?;
 		Ok(())
 	}
 }
@@ -466,9 +465,8 @@ fn add_grow_counter<R: Rules>(
 			.with_instructions(elements::Instructions::new(vec![
 				GetLocal(0),
 				GetLocal(0),
-				I64ExtendUI32,
+				I32Const(cost as i32),
-				I64Const(i64::from(cost)),
+				I32Mul,
-				I64Mul,
 				// todo: there should be strong guarantee that it does not return anything on
 				// stack?
 				Call(gas_func),
@@ -585,7 +583,7 @@ fn insert_metering_calls(
 		// If there the next block starts at this position, inject metering instructions.
 		let used_block = if let Some(block) = block_iter.peek() {
 			if block.start_pos == original_pos {
-				new_instrs.push(I64Const(block.cost as i64));
+				new_instrs.push(I32Const(block.cost as i32));
 				new_instrs.push(Call(gas_func));
 				true
 			} else {
@@ -614,7 +612,6 @@ fn insert_metering_calls(
 mod tests {
 	use super::*;
 	use parity_wasm::{builder, elements, elements::Instruction::*, serialize};
-	use pretty_assertions::assert_eq;
 
 	fn get_function_body(
 		module: &elements::Module,
@@ -642,20 +639,12 @@ mod tests {
 
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
-			&vec![I64Const(2), Call(0), GetGlobal(0), Call(2), End][..]
+			&vec![I32Const(2), Call(0), GetGlobal(0), Call(2), End][..]
 		);
 		assert_eq!(
 			get_function_body(&injected_module, 1).unwrap(),
-			&vec![
+			&vec![GetLocal(0), GetLocal(0), I32Const(10000), I32Mul, Call(0), GrowMemory(0), End,]
-				GetLocal(0),
+				[..]
-				GetLocal(0),
-				I64ExtendUI32,
-				I64Const(10000),
-				I64Mul,
-				Call(0),
-				GrowMemory(0),
-				End,
-			][..]
 		);
 
 		let binary = serialize(injected_module).expect("serialization failed");
@@ -678,7 +667,7 @@ mod tests {
 
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
-			&vec![I64Const(2), Call(0), GetGlobal(0), GrowMemory(0), End][..]
+			&vec![I32Const(2), Call(0), GetGlobal(0), GrowMemory(0), End][..]
 		);
 
 		assert_eq!(injected_module.functions_space(), 2);
@@ -730,17 +719,17 @@ mod tests {
 		assert_eq!(
 			get_function_body(&injected_module, 1).unwrap(),
 			&vec![
-				I64Const(3),
+				I32Const(3),
 				Call(0),
 				Call(1),
 				If(elements::BlockType::NoResult),
-				I64Const(3),
+				I32Const(3),
 				Call(0),
 				Call(1),
 				Call(1),
 				Call(1),
 				Else,
-				I64Const(2),
+				I32Const(2),
 				Call(0),
 				Call(1),
 				Call(1),
@@ -786,7 +775,7 @@ mod tests {
 		expected = r#"
 		(module
 			(func (result i32)
-				(call 0 (i64.const 1))
+				(call 0 (i32.const 1))
 				(get_global 0)))
 		"#
 	}
@@ -806,7 +795,7 @@ mod tests {
 		expected = r#"
 		(module
 			(func (result i32)
-				(call 0 (i64.const 6))
+				(call 0 (i32.const 6))
 				(get_global 0)
 				(block
 					(get_global 0)
@@ -835,16 +824,16 @@ mod tests {
 		expected = r#"
 		(module
 			(func (result i32)
-				(call 0 (i64.const 3))
+				(call 0 (i32.const 3))
 				(get_global 0)
 				(if
 					(then
-						(call 0 (i64.const 3))
+						(call 0 (i32.const 3))
 						(get_global 0)
 						(get_global 0)
 						(get_global 0))
 					(else
-						(call 0 (i64.const 2))
+						(call 0 (i32.const 2))
 						(get_global 0)
 						(get_global 0)))
 				(get_global 0)))
@@ -868,13 +857,13 @@ mod tests {
 		expected = r#"
 		(module
 			(func (result i32)
-				(call 0 (i64.const 6))
+				(call 0 (i32.const 6))
 				(get_global 0)
 				(block
 					(get_global 0)
 					(drop)
 					(br 0)
-					(call 0 (i64.const 2))
+					(call 0 (i32.const 2))
 					(get_global 0)
 					(drop))
 				(get_global 0)))
@@ -902,18 +891,18 @@ mod tests {
 		expected = r#"
 		(module
 			(func (result i32)
-				(call 0 (i64.const 5))
+				(call 0 (i32.const 5))
 				(get_global 0)
 				(block
 					(get_global 0)
 					(if
 						(then
-							(call 0 (i64.const 4))
+							(call 0 (i32.const 4))
 							(get_global 0)
 							(get_global 0)
 							(drop)
 							(br_if 1)))
-					(call 0 (i64.const 2))
+					(call 0 (i32.const 2))
 					(get_global 0)
 					(drop))
 				(get_global 0)))
@@ -944,18 +933,18 @@ mod tests {
 		expected = r#"
 		(module
 			(func (result i32)
-				(call 0 (i64.const 3))
+				(call 0 (i32.const 3))
 				(get_global 0)
 				(loop
-					(call 0 (i64.const 4))
+					(call 0 (i32.const 4))
 					(get_global 0)
 					(if
 						(then
-							(call 0 (i64.const 2))
+							(call 0 (i32.const 2))
 							(get_global 0)
 							(br_if 0))
 						(else
-							(call 0 (i64.const 4))
+							(call 0 (i32.const 4))
 							(get_global 0)
 							(get_global 0)
 							(drop)
@@ -980,13 +969,13 @@ mod tests {
 		expected = r#"
 		(module
 			(func (result i32)
-				(call 0 (i64.const 2))
+				(call 0 (i32.const 2))
 				(get_global 0)
 				(if
 					(then
-						(call 0 (i64.const 1))
+						(call 0 (i32.const 1))
 						(return)))
-				(call 0 (i64.const 1))
+				(call 0 (i32.const 1))
 				(get_global 0)))
 		"#
 	}
@@ -1009,18 +998,18 @@ mod tests {
 		expected = r#"
 		(module
 			(func (result i32)
-				(call 0 (i64.const 5))
+				(call 0 (i32.const 5))
 				(get_global 0)
 				(block
 					(get_global 0)
 					(if
 						(then
-							(call 0 (i64.const 1))
+							(call 0 (i32.const 1))
 							(br 1))
 						(else
-							(call 0 (i64.const 1))
+							(call 0 (i32.const 1))
 							(br 0)))
-					(call 0 (i64.const 2))
+					(call 0 (i32.const 2))
 					(get_global 0)
 					(drop))
 				(get_global 0)))
@@ -1042,9 +1031,9 @@ mod tests {
 		expected = r#"
 		(module
 			(func
-				(call 0 (i64.const 2))
+				(call 0 (i32.const 2))
 				(loop
-					(call 0 (i64.const 1))
+					(call 0 (i32.const 1))
 					(br 0)
 				)
 				unreachable

src/gas_metering/validation.rs

@@ -23,10 +23,10 @@ struct ControlFlowNode {
 	first_instr_pos: Option<usize>,
 
 	/// The actual gas cost of executing all instructions in the basic block.
-	actual_cost: u64,
+	actual_cost: u32,
 
 	/// The amount of gas charged by the injected metering instructions within this basic block.
-	charged_cost: u64,
+	charged_cost: u32,
 
 	/// Whether there are any other nodes in the graph that loop back to this one. Every cycle in
 	/// the control flow graph contains at least one node with this flag set.
@@ -68,10 +68,10 @@ impl ControlFlowGraph {
 	}
 
 	fn increment_actual_cost(&mut self, node_id: NodeId, cost: u32) {
-		self.get_node_mut(node_id).actual_cost += u64::from(cost);
+		self.get_node_mut(node_id).actual_cost += cost;
 	}
 
-	fn increment_charged_cost(&mut self, node_id: NodeId, cost: u64) {
+	fn increment_charged_cost(&mut self, node_id: NodeId, cost: u32) {
 		self.get_node_mut(node_id).charged_cost += cost;
 	}
 
@@ -267,9 +267,9 @@ fn validate_graph_gas_costs(graph: &ControlFlowGraph) -> bool {
 	fn visit(
 		graph: &ControlFlowGraph,
 		node_id: NodeId,
-		mut total_actual: u64,
+		mut total_actual: u32,
-		mut total_charged: u64,
+		mut total_charged: u32,
-		loop_costs: &mut Map<NodeId, (u64, u64)>,
+		loop_costs: &mut Map<NodeId, (u32, u32)>,
 	) -> bool {
 		let node = graph.get_node(node_id);
 

src/lib.rs

@@ -8,4 +8,6 @@ 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, compute_stack_height_weight, inject as inject_stack_limiter,
+};

src/stack_limiter/max_height.rs

@@ -1,16 +1,30 @@
-use alloc::vec::Vec;
+use super::resolve_func_type;
-use parity_wasm::elements::{self, BlockType, Type};
+use alloc::{vec, vec::Vec};
+use parity_wasm::elements::{self, BlockType, Type, ValueType};
 
 #[cfg(feature = "sign_ext")]
 use parity_wasm::elements::SignExtInstruction;
 
-use super::Context;
+#[cfg(feature = "trace-log")]
+macro_rules! trace {
+    ($($tt:tt)*) => {
+      ::log::trace!($($tt)*);
+    };
+}
+
+#[cfg(not(feature = "trace-log"))]
+macro_rules! trace {
+	($($tt:tt)*) => {};
+}
 
 // The cost in stack items that should be charged per call of a function. This is
 // 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_HEIGHT: u32 = 2;
+
+// Weight of an activation frame.
+const ACTIVATION_FRAME_WEIGHT: u32 = 32;
 
 /// Control stack frame.
 #[derive(Debug)]
@@ -19,36 +33,41 @@ struct Frame {
 	/// never passes control further was executed.
 	is_polymorphic: bool,
 
-	/// Count of values which will be pushed after the exit
+	/// Type of value which will be pushed after exiting
-	/// from the current block.
+	/// the current block or `None` if block does not return a result.
-	end_arity: u32,
+	result_type: Option<ValueType>,
 
-	/// Count of values which should be poped upon a branch to
+	/// Type of value which should be poped upon a branch to
-	/// this frame.
+	/// this frame or `None` if branching shouldn't affect the stack.
 	///
-	/// This might be diffirent from `end_arity` since branch
+	/// This might be diffirent from `result_type` since branch
 	/// to the loop header can't take any values.
-	branch_arity: u32,
+	branch_type: Option<ValueType>,
 
 	/// Stack height before entering in the block.
-	start_height: u32,
+	start_height: usize,
 }
 
-/// This is a compound stack that abstracts tracking height of the value stack
+/// This is a compound stack that abstracts tracking height and weight of the value stack
 /// and manipulation of the control stack.
 struct Stack {
-	height: u32,
+	values: Vec<ValueType>,
 	control_stack: Vec<Frame>,
 }
 
 impl Stack {
 	fn new() -> Stack {
-		Stack { height: ACTIVATION_FRAME_COST, control_stack: Vec::new() }
+		Stack { values: Vec::new(), control_stack: Vec::new() }
+	}
+
+	/// Returns current weight of the value stack.
+	fn weight(&self) -> u32 {
+		self.values.iter().map(|v| value_cost(*v)).sum()
 	}
 
 	/// Returns current height of the value stack.
-	fn height(&self) -> u32 {
+	fn height(&self) -> usize {
-		self.height
+		self.values.len()
 	}
 
 	/// Returns a reference to a frame by specified depth relative to the top of
@@ -71,63 +90,71 @@ impl Stack {
 
 	/// Push control frame into the control stack.
 	fn push_frame(&mut self, frame: Frame) {
+		trace!("  Push control frame {:?}", frame);
 		self.control_stack.push(frame);
 	}
 
 	/// Pop control frame from the control stack.
 	///
 	/// Returns `Err` if the control stack is empty.
+
+	#[allow(clippy::let_and_return)]
 	fn pop_frame(&mut self) -> Result<Frame, &'static str> {
-		self.control_stack.pop().ok_or("stack must be non-empty")
+		trace!("  Pop control frame");
+		let frame = self.control_stack.pop().ok_or("stack must be non-empty");
+		trace!("    {:?}", frame);
+		frame
 	}
 
 	/// Truncate the height of value stack to the specified height.
-	fn trunc(&mut self, new_height: u32) {
+	fn trunc(&mut self, new_height: usize) {
-		self.height = new_height;
+		trace!("  Truncate value stack to {}", new_height);
+		self.values.truncate(new_height);
 	}
 
-	/// Push specified number of values into the value stack.
+	/// Push a value into the value stack.
-	///
+	fn push_value(&mut self, value: ValueType) -> Result<(), &'static str> {
-	/// Returns `Err` if the height overflow usize value.
+		trace!("  Push {:?} to value stack", value);
-	fn push_values(&mut self, value_count: u32) -> Result<(), &'static str> {
+		self.values.push(value);
-		self.height = self.height.checked_add(value_count).ok_or("stack overflow")?;
+		if self.values.len() >= u32::MAX as usize {
+			return Err("stack overflow")
+		}
 		Ok(())
 	}
 
-	/// Pop specified number of values from the value stack.
+	/// Pop a value from the value stack.
 	///
 	/// Returns `Err` if the stack happen to be negative value after
-	/// values popped.
+	/// value popped.
-	fn pop_values(&mut self, value_count: u32) -> Result<(), &'static str> {
+	fn pop_value(&mut self) -> Result<Option<ValueType>, &'static str> {
-		if value_count == 0 {
+		let top_frame = self.frame(0)?;
-			return Ok(())
+		if self.height() == top_frame.start_height {
-		}
+			return if top_frame.is_polymorphic {
-		{
+				Ok(None)
-			let top_frame = self.frame(0)?;
+			} else {
-			if self.height == top_frame.start_height {
+				Err("trying to pop more values than pushed")
-				// 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")?;
+		if self.height() > 0 {
+			let vt = self.values.pop();
+			trace!("Pop {:?} from value stack", vt);
+			Ok(vt)
+		} else {
+			Err("trying to pop more values than pushed")
+		}
+	}
+}
 
-		Ok(())
+fn value_cost(val: ValueType) -> u32 {
+	match val {
+		ValueType::I32 | ValueType::F32 => 4,
+		ValueType::I64 | ValueType::F64 => 8,
 	}
 }
 
 /// This function expects the function to be validated.
-pub fn compute(
+pub fn compute(func_idx: u32, module: &elements::Module) -> Result<(u32, u32), &'static str> {
-	func_idx: u32,
-	ctx: &Context,
-	module: &elements::Module,
-) -> Result<u32, &'static str> {
 	use parity_wasm::elements::Instruction::*;
 
 	let func_section = module.function_section().ok_or("No function section")?;
@@ -150,47 +177,55 @@ pub fn compute(
 		.ok_or("Function body for the index isn't found")?;
 	let instructions = body.code();
 
+	// Get globals to resove their types
+	let globals: Vec<ValueType> = if let Some(global_section) = module.global_section() {
+		global_section
+			.entries()
+			.iter()
+			.map(|g| g.global_type().content_type())
+			.collect()
+	} else {
+		Vec::new()
+	};
+
+	let mut locals = func_signature.params().to_vec();
+	locals.extend(body.locals().iter().flat_map(|l| vec![l.value_type(); l.count() as usize]));
+
 	let mut stack = Stack::new();
-	let mut max_height: u32 = 0;
+	let mut max_weight: u32 = 0;
-	let mut pc = 0;
+	let mut max_height: usize = 0;
 
 	// 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;
+	let func_results = func_signature.results();
+	let param_weight: u32 = func_signature.params().iter().map(|v| value_cost(*v)).sum();
+
+	let func_result_type = if func_results.is_empty() { None } else { Some(func_results[0]) };
+
 	stack.push_frame(Frame {
 		is_polymorphic: false,
-		end_arity: func_arity,
+		result_type: func_result_type,
-		branch_arity: func_arity,
+		branch_type: func_result_type,
 		start_height: 0,
 	});
 
-	loop {
+	for opcode in instructions.elements() {
-		if pc >= instructions.elements().len() {
+		trace!("Processing opcode {:?}", opcode);
-			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 => {},
 			Block(ty) | Loop(ty) | If(ty) => {
-				let end_arity = if *ty == BlockType::NoResult { 0 } else { 1 };
-				let branch_arity = if let Loop(_) = *opcode { 0 } else { end_arity };
 				if let If(_) = *opcode {
-					stack.pop_values(1)?;
+					stack.pop_value()?;
 				}
 				let height = stack.height();
+				let end_result = if let BlockType::Value(vt) = *ty { Some(vt) } else { None };
 				stack.push_frame(Frame {
-					is_polymorphic: false,
+					is_polymorphic: stack.frame(0)?.is_polymorphic, /* Block inside unreachable
-					end_arity,
+					                                                 * code is
-					branch_arity,
+					                                                 * unreachable */
+					result_type: end_result,
+					branch_type: if let Loop(_) = *opcode { None } else { end_result },
 					start_height: height,
 				});
 			},
@@ -201,45 +236,56 @@ pub fn compute(
 			End => {
 				let frame = stack.pop_frame()?;
 				stack.trunc(frame.start_height);
-				stack.push_values(frame.end_arity)?;
+				if let Some(vt) = frame.result_type {
+					stack.push_value(vt)?;
+				}
+				// Push the frame back for now to allow for stack calculations. We'll get rid of it
+				// later
+				stack.push_frame(frame);
 			},
 			Unreachable => {
 				stack.mark_unreachable()?;
 			},
 			Br(target) => {
 				// Pop values for the destination block result.
-				let target_arity = stack.frame(*target)?.branch_arity;
+				if stack.frame(*target)?.branch_type.is_some() {
-				stack.pop_values(target_arity)?;
+					stack.pop_value()?;
+				}
 
 				// This instruction unconditionally transfers control to the specified block,
 				// thus all instruction until the end of the current block is deemed unreachable
 				stack.mark_unreachable()?;
 			},
 			BrIf(target) => {
+				let target_type = stack.frame(*target)?.branch_type;
 				// Pop values for the destination block result.
-				let target_arity = stack.frame(*target)?.branch_arity;
+				if target_type.is_some() {
-				stack.pop_values(target_arity)?;
+					stack.pop_value()?;
+				}
 
 				// Pop condition value.
-				stack.pop_values(1)?;
+				stack.pop_value()?;
 
 				// Push values back.
-				stack.push_values(target_arity)?;
+				if let Some(vt) = target_type {
+					stack.push_value(vt)?;
+				}
 			},
 			BrTable(br_table_data) => {
-				let arity_of_default = stack.frame(br_table_data.default)?.branch_arity;
+				let default_type = stack.frame(br_table_data.default)?.branch_type;
 
 				// Check that all jump targets have an equal arities.
 				for target in &*br_table_data.table {
-					let arity = stack.frame(*target)?.branch_arity;
+					if stack.frame(*target)?.branch_type != default_type {
-					if arity != arity_of_default {
+						return Err("Types of all jump-targets must be equal")
-						return Err("Arity of all jump-targets must be equal")
 					}
 				}
 
-				// Because all jump targets have an equal arities, we can just take arity of
+				// Because all jump targets have equal types, we can just take type of
 				// the default branch.
-				stack.pop_values(arity_of_default)?;
+				if default_type.is_some() {
+					stack.pop_value()?;
+				}
 
 				// 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.
@@ -248,82 +294,114 @@ pub fn compute(
 			Return => {
 				// Pop return values of the function. Mark successive instructions as unreachable
 				// since this instruction doesn't let control flow to go further.
-				stack.pop_values(func_arity)?;
+				if func_result_type.is_some() {
+					stack.pop_value()?;
+				}
 				stack.mark_unreachable()?;
 			},
-			Call(fn_idx) => {
+			Call(idx) => {
-				let ty_idx = ctx.func_type(*fn_idx).ok_or("function idx is not found in the func types list")?;
+				let ty = resolve_func_type(*idx, module)?;
-				let Type::Function(ty) =
-					type_section.types().get(ty_idx as usize).ok_or("Type not found")?;
 
 				// Pop values for arguments of the function.
-				stack.pop_values(ty.params().len() as u32)?;
+				for _ in ty.params() {
+					stack.pop_value()?;
+				}
 
 				// Push result of the function execution to the stack.
-				let callee_arity = ty.results().len() as u32;
+				let callee_results = ty.results();
-				stack.push_values(callee_arity)?;
+				if !callee_results.is_empty() {
+					stack.push_value(callee_results[0])?;
+				}
 			},
 			CallIndirect(x, _) => {
 				let Type::Function(ty) =
 					type_section.types().get(*x as usize).ok_or("Type not found")?;
 
 				// Pop the offset into the function table.
-				stack.pop_values(1)?;
+				stack.pop_value()?;
 
 				// Pop values for arguments of the function.
-				stack.pop_values(ty.params().len() as u32)?;
+				for _ in ty.params() {
+					stack.pop_value()?;
+				}
 
 				// Push result of the function execution to the stack.
-				let callee_arity = ty.results().len() as u32;
+				let callee_results = ty.results();
-				stack.push_values(callee_arity)?;
+				if !callee_results.is_empty() {
+					stack.push_value(callee_results[0])?;
+				}
 			},
 			Drop => {
-				stack.pop_values(1)?;
+				stack.pop_value()?;
 			},
 			Select => {
 				// Pop two values and one condition.
-				stack.pop_values(2)?;
+				let val = stack.pop_value()?;
-				stack.pop_values(1)?;
+				stack.pop_value()?;
+				stack.pop_value()?;
 
 				// Push the selected value.
-				stack.push_values(1)?;
+				if let Some(vt) = val {
+					stack.push_value(vt)?;
+				}
 			},
-			GetLocal(_) => {
+			GetLocal(idx) => {
-				stack.push_values(1)?;
+				let idx = *idx as usize;
+				if idx >= locals.len() {
+					return Err("Reference to a global is out of bounds")
+				}
+				stack.push_value(locals[idx])?;
 			},
 			SetLocal(_) => {
-				stack.pop_values(1)?;
+				stack.pop_value()?;
 			},
-			TeeLocal(_) => {
+			TeeLocal(idx) => {
 				// This instruction pops and pushes the value, so
 				// effectively it doesn't modify the stack height.
-				stack.pop_values(1)?;
+				let idx = *idx as usize;
-				stack.push_values(1)?;
+				if idx >= locals.len() {
+					return Err("Reference to a local is out of bounds")
+				}
+				stack.pop_value()?;
+				stack.push_value(locals[idx])?;
 			},
-			GetGlobal(_) => {
+			GetGlobal(idx) => {
-				stack.push_values(1)?;
+				let idx = *idx as usize;
+				if idx >= globals.len() {
+					return Err("Reference to a global is out of bounds")
+				}
+				stack.push_value(globals[idx])?;
 			},
 			SetGlobal(_) => {
-				stack.pop_values(1)?;
+				stack.pop_value()?;
 			},
+
+			// These instructions pop the address and pushes the result
 			I32Load(_, _) |
-			I64Load(_, _) |
-			F32Load(_, _) |
-			F64Load(_, _) |
 			I32Load8S(_, _) |
 			I32Load8U(_, _) |
 			I32Load16S(_, _) |
-			I32Load16U(_, _) |
+			I32Load16U(_, _) => {
+				stack.pop_value()?;
+				stack.push_value(ValueType::I32)?;
+			},
+			I64Load(_, _) |
 			I64Load8S(_, _) |
 			I64Load8U(_, _) |
 			I64Load16S(_, _) |
 			I64Load16U(_, _) |
 			I64Load32S(_, _) |
 			I64Load32U(_, _) => {
-				// These instructions pop the address and pushes the result,
+				stack.pop_value()?;
-				// which effictively don't modify the stack height.
+				stack.push_value(ValueType::I64)?;
-				stack.pop_values(1)?;
+			},
-				stack.push_values(1)?;
+			F32Load(_, _) => {
+				stack.pop_value()?;
+				stack.push_value(ValueType::F32)?;
+			},
+			F64Load(_, _) => {
+				stack.pop_value()?;
+				stack.push_value(ValueType::F64)?;
 			},
 
 			I32Store(_, _) |
@@ -336,29 +414,38 @@ pub fn compute(
 			I64Store16(_, _) |
 			I64Store32(_, _) => {
 				// These instructions pop the address and the value.
-				stack.pop_values(2)?;
+				stack.pop_value()?;
+				stack.pop_value()?;
 			},
 
 			CurrentMemory(_) => {
 				// Pushes current memory size
-				stack.push_values(1)?;
+				stack.push_value(ValueType::I32)?;
 			},
 			GrowMemory(_) => {
 				// Grow memory takes the value of pages to grow and pushes
-				stack.pop_values(1)?;
+				stack.pop_value()?;
-				stack.push_values(1)?;
+				stack.push_value(ValueType::I32)?;
 			},
 
-			I32Const(_) | I64Const(_) | F32Const(_) | F64Const(_) => {
+			I32Const(_) => {
-				// These instructions just push the single literal value onto the stack.
+				stack.push_value(ValueType::I32)?;
-				stack.push_values(1)?;
+			},
+			I64Const(_) => {
+				stack.push_value(ValueType::I64)?;
+			},
+			F32Const(_) => {
+				stack.push_value(ValueType::F32)?;
+			},
+			F64Const(_) => {
+				stack.push_value(ValueType::F64)?;
 			},
 
 			I32Eqz | I64Eqz => {
 				// These instructions pop the value and compare it against zero, and pushes
 				// the result of the comparison.
-				stack.pop_values(1)?;
+				stack.pop_value()?;
-				stack.push_values(1)?;
+				stack.push_value(ValueType::I32)?;
 			},
 
 			I32Eq | I32Ne | I32LtS | I32LtU | I32GtS | I32GtU | I32LeS | I32LeU | I32GeS |
@@ -366,16 +453,18 @@ pub fn compute(
 			I64GeS | I64GeU | F32Eq | F32Ne | F32Lt | F32Gt | F32Le | F32Ge | F64Eq | F64Ne |
 			F64Lt | F64Gt | F64Le | F64Ge => {
 				// Comparison operations take two operands and produce one result.
-				stack.pop_values(2)?;
+				stack.pop_value()?;
-				stack.push_values(1)?;
+				stack.pop_value()?;
+				stack.push_value(ValueType::I32)?;
 			},
 
 			I32Clz | I32Ctz | I32Popcnt | I64Clz | I64Ctz | I64Popcnt | F32Abs | F32Neg |
 			F32Ceil | F32Floor | F32Trunc | F32Nearest | F32Sqrt | F64Abs | F64Neg | F64Ceil |
 			F64Floor | F64Trunc | F64Nearest | F64Sqrt => {
 				// Unary operators take one operand and produce one result.
-				stack.pop_values(1)?;
+				if let Some(vt) = stack.pop_value()? {
-				stack.push_values(1)?;
+					stack.push_value(vt)?;
+				}
 			},
 
 			I32Add | I32Sub | I32Mul | I32DivS | I32DivU | I32RemS | I32RemU | I32And | I32Or |
@@ -385,19 +474,34 @@ pub fn compute(
 			F32Min | F32Max | F32Copysign | F64Add | F64Sub | F64Mul | F64Div | F64Min |
 			F64Max | F64Copysign => {
 				// Binary operators take two operands and produce one result.
-				stack.pop_values(2)?;
+				let val = stack.pop_value()?;
-				stack.push_values(1)?;
+				stack.pop_value()?;
+				if let Some(vt) = val {
+					stack.push_value(vt)?;
+				}
 			},
 
+			// Conversion operators take one value and produce one result.
 			I32WrapI64 | I32TruncSF32 | I32TruncUF32 | I32TruncSF64 | I32TruncUF64 |
+			I32ReinterpretF32 => {
+				stack.pop_value()?;
+				stack.push_value(ValueType::I32)?;
+			},
 			I64ExtendSI32 | I64ExtendUI32 | I64TruncSF32 | I64TruncUF32 | I64TruncSF64 |
-			I64TruncUF64 | F32ConvertSI32 | F32ConvertUI32 | F32ConvertSI64 | F32ConvertUI64 |
+			I64TruncUF64 | I64ReinterpretF64 => {
-			F32DemoteF64 | F64ConvertSI32 | F64ConvertUI32 | F64ConvertSI64 | F64ConvertUI64 |
+				stack.pop_value()?;
-			F64PromoteF32 | I32ReinterpretF32 | I64ReinterpretF64 | F32ReinterpretI32 |
+				stack.push_value(ValueType::I64)?;
+			},
+			F32ConvertSI32 | F32ConvertUI32 | F32ConvertSI64 | F32ConvertUI64 | F32DemoteF64 |
+			F32ReinterpretI32 => {
+				stack.pop_value()?;
+				stack.push_value(ValueType::F32)?;
+			},
+
+			F64ConvertSI32 | F64ConvertUI32 | F64ConvertSI64 | F64ConvertUI64 | F64PromoteF32 |
 			F64ReinterpretI64 => {
-				// Conversion operators take one value and produce one result.
+				stack.pop_value()?;
-				stack.pop_values(1)?;
+				stack.push_value(ValueType::F64)?;
-				stack.push_values(1)?;
 			},
 
 			#[cfg(feature = "sign_ext")]
@@ -405,37 +509,64 @@ pub fn compute(
 			SignExt(SignExtInstruction::I32Extend16S) |
 			SignExt(SignExtInstruction::I64Extend8S) |
 			SignExt(SignExtInstruction::I64Extend16S) |
-			SignExt(SignExtInstruction::I64Extend32S) => {
+			SignExt(SignExtInstruction::I64Extend32S) =>
-				stack.pop_values(1)?;
+				if let Some(vt) = stack.pop_value()? {
-				stack.push_values(1)?;
+					stack.push_value(vt)?;
-			},
+				},
+		}
+
+		// If current value stack is heavier than maximal weight observed so far,
+		// save the new weight.
+		// However, we don't increase maximal value in unreachable code.
+		if !stack.frame(0)?.is_polymorphic {
+			let (cur_weight, cur_height) = (stack.weight(), stack.height());
+			if cur_weight > max_weight {
+				max_weight = cur_weight;
+				trace!("Max weight is now {}", max_weight);
+			}
+			if cur_height > max_height {
+				max_height = cur_height;
+				trace!("Max height is now {}", max_height);
+			}
+		}
+
+		// Post-execution stage: pop a control frame if block is ended
+		if *opcode == End {
+			stack.pop_frame()?;
 		}
-		pc += 1;
 	}
 
-	Ok(max_height)
+	trace!("Final max stack height: {} + {}", ACTIVATION_FRAME_HEIGHT, max_height);
+	trace!(
+		"Final max stack weight: {} + {} + {}",
+		ACTIVATION_FRAME_WEIGHT,
+		max_weight,
+		param_weight
+	);
+
+	Ok((
+		ACTIVATION_FRAME_HEIGHT + max_height as u32,
+		ACTIVATION_FRAME_WEIGHT + max_weight + param_weight,
+	))
 }
 
 #[cfg(test)]
 mod tests {
-	use super::ACTIVATION_FRAME_COST;
+	use super::*;
-	use crate::stack_limiter::prepare_context;
 	use parity_wasm::elements;
 
+	#[cfg(feature = "trace-log")]
+	use test_log::test;
+
 	fn parse_wat(source: &str) -> elements::Module {
 		elements::deserialize_buffer(&wat::parse_str(source).expect("Failed to wat2wasm"))
 			.expect("Failed to deserialize the module")
 	}
 
-	fn test_compute(func_idx: u32, source: &str) -> u32 {
-		let module = parse_wat(source);
-		let ctx = prepare_context(&module, 0).unwrap();
-		ctx.stack_cost(func_idx).unwrap()
-	}
-
 	#[test]
 	fn simple_test() {
-		let module = r#"
+		let module = parse_wat(
+			r#"
 (module
 	(func
 		i32.const 1
@@ -446,30 +577,34 @@ mod tests {
 		drop
 	)
 )
-"#;
+"#,
+		);
 
-		let height = test_compute(0, module);
+		let res = compute(0, &module).unwrap();
-		assert_eq!(height, 3 + ACTIVATION_FRAME_COST);
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 3, ACTIVATION_FRAME_WEIGHT + 12));
 	}
 
 	#[test]
 	fn implicit_and_explicit_return() {
-		let module = r#"
+		let module = parse_wat(
+			r#"
 (module
 	(func (result i32)
-		i32.const 0
+		i64.const 0
 		return
 	)
 )
-"#;
+"#,
+		);
 
-		let height = test_compute(0, module);
+		let res = compute(0, &module).unwrap();
-		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 1, ACTIVATION_FRAME_WEIGHT + 8));
 	}
 
 	#[test]
 	fn dont_count_in_unreachable() {
-		let module = r#"
+		let module = parse_wat(
+			r#"
 (module
   (memory 0)
   (func (result i32)
@@ -477,15 +612,17 @@ mod tests {
 	grow_memory
   )
 )
-"#;
+"#,
+		);
 
-		let height = test_compute(0, module);
+		let res = compute(0, &module).unwrap();
-		assert_eq!(height, ACTIVATION_FRAME_COST);
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT, ACTIVATION_FRAME_WEIGHT));
 	}
 
 	#[test]
 	fn yet_another_test() {
-		let module = r#"
+		let module = parse_wat(
+			r#"
 (module
   (memory 0)
   (func
@@ -504,15 +641,17 @@ mod tests {
 	i32.const 2
   )
 )
-"#;
+"#,
+		);
 
-		let height = test_compute(0, module);
+		let res = compute(0, &module).unwrap();
-		assert_eq!(height, 2 + ACTIVATION_FRAME_COST);
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 2, ACTIVATION_FRAME_WEIGHT + 8));
 	}
 
 	#[test]
 	fn call_indirect() {
-		let module = r#"
+		let module = parse_wat(
+			r#"
 (module
 	(table $ptr 1 1 funcref)
 	(elem $ptr (i32.const 0) func 1)
@@ -526,15 +665,17 @@ mod tests {
 		drop
 	)
 )
-"#;
+"#,
+		);
 
-		let height = test_compute(0, module);
+		let res = compute(0, &module).unwrap();
-		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 1, ACTIVATION_FRAME_WEIGHT + 4));
 	}
 
 	#[test]
 	fn breaks() {
-		let module = r#"
+		let module = parse_wat(
+			r#"
 (module
 	(func $main
 		block (result i32)
@@ -546,15 +687,17 @@ mod tests {
 		drop
 	)
 )
-"#;
+"#,
+		);
 
-		let height = test_compute(0, module);
+		let res = compute(0, &module).unwrap();
-		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 1, ACTIVATION_FRAME_WEIGHT + 4));
 	}
 
 	#[test]
 	fn if_else_works() {
-		let module = r#"
+		let module = parse_wat(
+			r#"
 (module
 	(func $main
 		i32.const 7
@@ -570,9 +713,10 @@ mod tests {
 		drop
 	)
 )
-"#;
+"#,
+		);
 
-		let height = test_compute(0, module);
+		let res = compute(0, &module).unwrap();
-		assert_eq!(height, 3 + ACTIVATION_FRAME_COST);
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 3, ACTIVATION_FRAME_WEIGHT + 12));
 	}
 }

src/stack_limiter/mod.rs

@@ -39,15 +39,7 @@ mod max_height;
 mod thunk;
 
 pub struct Context {
-	/// Number of functions that the module imports. Required to convert defined functions indicies
-	/// into the global function index space.
-	func_imports: u32,
-	/// For each function in the function space this vector stores the respective type index.
-	func_types: Vec<u32>,
-	/// The index of the global variable that contains the current stack height.
 	stack_height_global_idx: u32,
-	/// Logical stack costs for each function in the function space. Imported functions have cost
-	/// of 0.
 	func_stack_costs: Vec<u32>,
 	stack_limit: u32,
 }
@@ -63,11 +55,6 @@ impl Context {
 		self.func_stack_costs.get(func_idx as usize).cloned()
 	}
 
-	/// Returns a reference to the function type index given by the index into the function space.
-	fn func_type(&self, func_idx: u32) -> Option<u32> {
-		self.func_types.get(func_idx as usize).copied()
-	}
-
 	/// Returns stack limit specified by the rules.
 	fn stack_limit(&self) -> u32 {
 		self.stack_limit
@@ -128,106 +115,20 @@ pub fn inject(
 	mut module: elements::Module,
 	stack_limit: u32,
 ) -> Result<elements::Module, &'static str> {
-	let mut ctx = prepare_context(&module, stack_limit)?;
+	let mut ctx = Context {
+		stack_height_global_idx: generate_stack_height_global(&mut module),
+		func_stack_costs: compute_stack_costs(&module)?,
+		stack_limit,
+	};
 
-	generate_stack_height_global(&mut ctx.stack_height_global_idx, &mut module)?;
+	instrument_functions(&mut ctx, &mut module)?;
-	instrument_functions(&ctx, &mut module)?;
 	let module = thunk::generate_thunks(&mut ctx, module)?;
 
 	Ok(module)
 }
 
-fn prepare_context(module: &elements::Module, stack_limit: u32) -> Result<Context, &'static str> {
-	let mut ctx = Context {
-		func_imports: module.import_count(elements::ImportCountType::Function) as u32,
-		func_types: Vec::new(),
-		stack_height_global_idx: 0,
-		func_stack_costs: Vec::new(),
-		stack_limit,
-	};
-	collect_func_types(&mut ctx, &module)?;
-	compute_stack_costs(&mut ctx, &module)?;
-	Ok(ctx)
-}
-
-fn collect_func_types(ctx: &mut Context, module: &elements::Module) -> Result<(), &'static str> {
-	let types = module.type_section().map(|ts| ts.types()).unwrap_or(&[]);
-	let functions = module.function_section().map(|fs| fs.entries()).unwrap_or(&[]);
-	let imports = module.import_section().map(|is| is.entries()).unwrap_or(&[]);
-
-	let ensure_ty = |sig_idx: u32| -> Result<(), &'static str> {
-		let Type::Function(_) = types
-			.get(sig_idx as usize)
-			.ok_or("The signature as specified by a function isn't defined")?;
-		Ok(())
-	};
-
-	for import in imports {
-		if let elements::External::Function(sig_idx) = import.external() {
-			ensure_ty(*sig_idx)?;
-			ctx.func_types.push(*sig_idx);
-		}
-	}
-	for def_func_idx in functions {
-		ensure_ty(def_func_idx.type_ref())?;
-		ctx.func_types.push(def_func_idx.type_ref());
-	}
-
-	Ok(())
-}
-
-/// Calculate stack costs for all functions in the function space.
-///
-/// The function space consists of the imported functions followed by defined functions.
-/// All imported functions assumed to have the cost of 0.
-fn compute_stack_costs(ctx: &mut Context, module: &elements::Module) -> Result<(), &'static str> {
-	for _ in 0..ctx.func_imports {
-		ctx.func_stack_costs.push(0);
-	}
-	let def_func_n = module.function_section().map(|fs| fs.entries().len()).unwrap_or(0) as u32;
-	for def_func_idx in 0..def_func_n {
-		let cost = compute_stack_cost(def_func_idx, ctx, module)?;
-		ctx.func_stack_costs.push(cost);
-	}
-	Ok(())
-}
-
-/// Computes the stack cost of a given function. The function is specified by its index in the
-/// declared function space.
-///
-/// Stack cost of a given 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(
-	def_func_idx: u32,
-	ctx: &Context,
-	module: &elements::Module,
-) -> Result<u32, &'static str> {
-	let code_section =
-		module.code_section().ok_or("Due to validation code section should exists")?;
-	let body = &code_section
-		.bodies()
-		.get(def_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(def_func_idx, ctx, module)?;
-
-	locals_count
-		.checked_add(max_stack_height)
-		.ok_or("Overflow in adding locals_count and max_stack_height")
-}
-
 /// Generate a new global that will be used for tracking current stack height.
-fn generate_stack_height_global(
+fn generate_stack_height_global(module: &mut elements::Module) -> u32 {
-	stack_height_global_idx: &mut u32,
-	module: &mut elements::Module,
-) -> Result<(), &'static str> {
 	let global_entry = builder::global()
 		.value_type()
 		.i32()
@@ -239,41 +140,85 @@ fn generate_stack_height_global(
 	for section in module.sections_mut() {
 		if let elements::Section::Global(gs) = section {
 			gs.entries_mut().push(global_entry);
-			*stack_height_global_idx = (gs.entries().len() as u32) - 1;
+			return (gs.entries().len() as u32) - 1
-			return Ok(());
 		}
 	}
 
 	// Existing section not found, create one!
-	//
-	// It's a bit tricky since the sections have a strict prescribed order.
-	let global_section = elements::GlobalSection::with_entries(vec![global_entry]);
-	let prec_index = module
-		.sections()
-		.iter()
-		.rposition(|section| {
-			use elements::Section::*;
-			match section {
-				Type(_) | Import(_) | Function(_) | Table(_) | Memory(_) => true,
-				_ => false,
-			}
-		})
-		.ok_or("generate stack height global hasn't found any preceding sections")?;
-	// now `prec_index` points to the last section preceding the `global_section`. It's guaranteed that at least
-	// one of those functions is present. Therefore, the candidate position for the global section is the following
-	// one. However, technically, custom sections could occupy any place between the well-known sections.
-	//
-	// Now, regarding `+1` here. `insert` panics iff `index > len`. `prec_index + 1` can only be equal to `len`.
 	module
 		.sections_mut()
-		.insert(prec_index + 1, elements::Section::Global(global_section));
+		.push(elements::Section::Global(elements::GlobalSection::with_entries(vec![global_entry])));
-	// First entry in the brand new globals section.
+	0
-	*stack_height_global_idx = 0;
-
-	Ok(())
 }
 
-fn instrument_functions(ctx: &Context, module: &mut elements::Module) -> Result<(), &'static str> {
+/// Calculate stack costs for all functions.
+///
+/// Returns a vector with a stack cost for each function, including imports.
+pub fn compute_stack_costs(module: &elements::Module) -> Result<Vec<u32>, &'static str> {
+	let func_imports = module.import_count(elements::ImportCountType::Function);
+
+	// TODO: optimize!
+	(0..module.functions_space())
+		.map(|func_idx| {
+			if func_idx < func_imports {
+				// We can't calculate stack_cost of the import functions.
+				Ok(0)
+			} else {
+				compute_stack_cost(func_idx as u32, module)
+			}
+		})
+		.collect()
+}
+
+/// 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.
+pub fn compute_stack_cost(func_idx: u32, module: &elements::Module) -> Result<u32, &'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;
+	let defined_func_idx = func_idx
+		.checked_sub(func_imports)
+		.ok_or("This should be a index of a defined function")?;
+
+	let code_section =
+		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_stack_weight) = 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")
+}
+
+/// Stack height is the measurement maximum wasm stack height reached during function execution.
+/// Stack weight is weighted value which approximates a real stack size on x64 architecture.
+pub fn compute_stack_height_weight(
+	func_idx: u32,
+	module: &elements::Module,
+) -> Result<(u32, u32), &'static str> {
+	let func_imports = module.import_count(elements::ImportCountType::Function) as u32;
+	let defined_func_idx = func_idx
+		.checked_sub(func_imports)
+		.ok_or("This should be a index of a defined function")?;
+
+	max_height::compute(defined_func_idx, module)
+}
+
+fn instrument_functions(
+	ctx: &mut Context,
+	module: &mut elements::Module,
+) -> Result<(), &'static str> {
 	for section in module.sections_mut() {
 		if let elements::Section::Code(code_section) = section {
 			for func_body in code_section.bodies_mut() {
@@ -311,7 +256,7 @@ fn instrument_functions(ctx: &Context, module: &mut elements::Module) -> Result<
 ///
 /// drop
 /// ```
-fn instrument_function(ctx: &Context, func: &mut Instructions) -> Result<(), &'static str> {
+fn instrument_function(ctx: &mut Context, func: &mut Instructions) -> Result<(), &'static str> {
 	use Instruction::*;
 
 	struct InstrumentCall {
@@ -378,6 +323,42 @@ fn instrument_function(ctx: &Context, func: &mut Instructions) -> Result<(), &'s
 	Ok(())
 }
 
+fn resolve_func_type(
+	func_idx: u32,
+	module: &elements::Module,
+) -> Result<&elements::FunctionType, &'static str> {
+	let types = module.type_section().map(|ts| ts.types()).unwrap_or(&[]);
+	let functions = module.function_section().map(|fs| fs.entries()).unwrap_or(&[]);
+
+	let func_imports = module.import_count(elements::ImportCountType::Function);
+	let sig_idx = if func_idx < func_imports as u32 {
+		module
+			.import_section()
+			.expect("function import count is not zero; import section must exists; qed")
+			.entries()
+			.iter()
+			.filter_map(|entry| match entry.external() {
+				elements::External::Function(idx) => Some(*idx),
+				_ => None,
+			})
+			.nth(func_idx as usize)
+			.expect(
+				"func_idx is less than function imports count;
+				nth function import must be `Some`;
+				qed",
+			)
+	} else {
+		functions
+			.get(func_idx as usize - func_imports)
+			.ok_or("Function at the specified index is not defined")?
+			.type_ref()
+	};
+	let Type::Function(ty) = types
+		.get(sig_idx as usize)
+		.ok_or("The signature as specified by a function isn't defined")?;
+	Ok(ty)
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;

src/stack_limiter/thunk.rs

@@ -1,17 +1,17 @@
 #[cfg(not(features = "std"))]
 use alloc::collections::BTreeMap as Map;
 use alloc::vec::Vec;
-use parity_wasm::elements::{self, Internal};
+use parity_wasm::{
+	builder,
+	elements::{self, FunctionType, Internal},
+};
 #[cfg(features = "std")]
 use std::collections::HashMap as Map;
 
-use super::Context;
+use super::{resolve_func_type, Context};
 
 struct Thunk {
-	/// The index of the signature in the type section.
+	signature: FunctionType,
-	type_idx: u32,
-	/// The number of parameters the function has.
-	param_num: u32,
 	// Index in function space of this thunk.
 	idx: Option<u32>,
 	callee_stack_cost: u32,
@@ -19,11 +19,10 @@ struct Thunk {
 
 pub fn generate_thunks(
 	ctx: &mut Context,
-	mut module: elements::Module,
+	module: elements::Module,
 ) -> Result<elements::Module, &'static str> {
 	// First, we need to collect all function indices that should be replaced by thunks
 	let mut replacement_map: Map<u32, Thunk> = {
-		let types = module.type_section().map(|ts| ts.types()).unwrap_or(&[]);
 		let exports = module.export_section().map(|es| es.entries()).unwrap_or(&[]);
 		let elem_segments = module.elements_section().map(|es| es.entries()).unwrap_or(&[]);
 		let start_func_idx = module.start_section();
@@ -46,12 +45,14 @@ pub fn generate_thunks(
 
 			// Don't generate a thunk if stack_cost of a callee is zero.
 			if callee_stack_cost != 0 {
-				let type_idx = ctx.func_type(func_idx).ok_or("type idx for thunk not found")?;
+				replacement_map.insert(
-				let elements::Type::Function(func_ty) =
+					func_idx,
-					types.get(type_idx as usize).ok_or("sig for thunk is not found")?;
+					Thunk {
-				let param_num = func_ty.params().len() as u32;
+						signature: resolve_func_type(func_idx, &module)?.clone(),
-				replacement_map
+						idx: None,
-					.insert(func_idx, Thunk { type_idx, param_num, idx: None, callee_stack_cost });
+						callee_stack_cost,
+					},
+				);
 			}
 		}
 
@@ -60,6 +61,10 @@ pub fn generate_thunks(
 
 	// Then, we generate a thunk for each original function.
 
+	// Save current func_idx
+	let mut next_func_idx = module.functions_space() as u32;
+
+	let mut mbuilder = builder::from_module(module);
 	for (func_idx, thunk) in replacement_map.iter_mut() {
 		let instrumented_call = instrument_call!(
 			*func_idx,
@@ -72,23 +77,32 @@ pub fn generate_thunks(
 		//  - instrumented call
 		//  - end
 		let mut thunk_body: Vec<elements::Instruction> =
-			Vec::with_capacity(thunk.param_num as usize + instrumented_call.len() + 1);
+			Vec::with_capacity(thunk.signature.params().len() + instrumented_call.len() + 1);
 
-		for arg_idx in 0..thunk.param_num {
+		for (arg_idx, _) in thunk.signature.params().iter().enumerate() {
-			thunk_body.push(elements::Instruction::GetLocal(arg_idx));
+			thunk_body.push(elements::Instruction::GetLocal(arg_idx as u32));
 		}
 		thunk_body.extend_from_slice(&instrumented_call);
 		thunk_body.push(elements::Instruction::End);
 
-		let func_idx = insert_function(
+		// TODO: Don't generate a signature, but find an existing one.
-			ctx,
+
-			&mut module,
+		mbuilder = mbuilder
-			thunk.type_idx,
+			.function()
-			Vec::new(), // No declared local variables.
+			// Signature of the thunk should match the original function signature.
-			elements::Instructions::new(thunk_body),
+			.signature()
-		)?;
+			.with_params(thunk.signature.params().to_vec())
-		thunk.idx = Some(func_idx);
+			.with_results(thunk.signature.results().to_vec())
+			.build()
+			.body()
+			.with_instructions(elements::Instructions::new(thunk_body))
+			.build()
+			.build();
+
+		thunk.idx = Some(next_func_idx);
+		next_func_idx += 1;
 	}
+	let mut module = mbuilder.build();
 
 	// And finally, fixup thunks in export and table sections.
 
@@ -104,20 +118,18 @@ pub fn generate_thunks(
 
 	for section in module.sections_mut() {
 		match section {
-			elements::Section::Export(export_section) => {
+			elements::Section::Export(export_section) =>
 				for entry in export_section.entries_mut() {
 					if let Internal::Function(function_idx) = entry.internal_mut() {
 						fixup(function_idx)
 					}
-				}
+				},
-			},
+			elements::Section::Element(elem_section) =>
-			elements::Section::Element(elem_section) => {
 				for segment in elem_section.entries_mut() {
 					for function_idx in segment.members_mut() {
 						fixup(function_idx)
 					}
-				}
+				},
-			},
 			elements::Section::Start(start_idx) => fixup(start_idx),
 			_ => {},
 		}
@@ -125,30 +137,3 @@ pub fn generate_thunks(
 
 	Ok(module)
 }
-
-/// Inserts a new function into the module and returns it's index in the function space.
-///
-/// Specifically, inserts entires into the function section and the code section.
-fn insert_function(
-	ctx: &Context,
-	module: &mut elements::Module,
-	type_idx: u32,
-	locals: Vec<elements::Local>,
-	insns: elements::Instructions,
-) -> Result<u32, &'static str> {
-	let funcs = module
-		.function_section_mut()
-		.ok_or("insert function no function section")?
-		.entries_mut();
-	let new_func_idx = ctx
-		.func_imports
-		.checked_add(funcs.len() as u32)
-		.ok_or("insert function func idx overflow")?;
-	funcs.push(elements::Func::new(type_idx));
-
-	let func_bodies =
-		module.code_section_mut().ok_or("insert function no code section")?.bodies_mut();
-	func_bodies.push(elements::FuncBody::new(locals, insns));
-
-	Ok(new_func_idx)
-}

tests/expectations/gas/branch.wat

@@ -1,10 +1,10 @@
 (module
   (type (;0;) (func (result i32)))
-  (type (;1;) (func (param i64)))
+  (type (;1;) (func (param i32)))
   (import "env" "gas" (func (;0;) (type 1)))
   (func $fibonacci_with_break (;1;) (type 0) (result i32)
     (local i32 i32)
-    i64.const 13
+    i32.const 13
     call 0
     block  ;; label = @1
       i32.const 0
@@ -18,7 +18,7 @@
       local.set 1
       i32.const 1
       br_if 0 (;@1;)
-      i64.const 5
+      i32.const 5
       call 0
       local.get 0
       local.get 1

tests/expectations/gas/call.wat

@@ -1,10 +1,10 @@
 (module
   (type (;0;) (func (param i32 i32) (result i32)))
-  (type (;1;) (func (param i64)))
+  (type (;1;) (func (param i32)))
   (import "env" "gas" (func (;0;) (type 1)))
   (func $add_locals (;1;) (type 0) (param $x i32) (param $y i32) (result i32)
     (local i32)
-    i64.const 5
+    i32.const 5
     call 0
     local.get $x
     local.get $y
@@ -13,7 +13,7 @@
     local.get 2
   )
   (func $add (;2;) (type 0) (param i32 i32) (result i32)
-    i64.const 3
+    i32.const 3
     call 0
     local.get 0
     local.get 1

tests/expectations/gas/ifs.wat

@@ -1,19 +1,19 @@
 (module
   (type (;0;) (func (param i32) (result i32)))
-  (type (;1;) (func (param i64)))
+  (type (;1;) (func (param i32)))
   (import "env" "gas" (func (;0;) (type 1)))
   (func (;1;) (type 0) (param i32) (result i32)
-    i64.const 2
+    i32.const 2
     call 0
     i32.const 1
     if (result i32)  ;; label = @1
-      i64.const 3
+      i32.const 3
       call 0
       local.get 0
       i32.const 1
       i32.add
     else
-      i64.const 2
+      i32.const 2
       call 0
       local.get 0
       i32.popcnt

tests/expectations/gas/simple.wat

@@ -1,16 +1,16 @@
 (module
   (type (;0;) (func))
-  (type (;1;) (func (param i64)))
+  (type (;1;) (func (param i32)))
   (import "env" "gas" (func (;0;) (type 1)))
   (func (;1;) (type 0)
-    i64.const 2
+    i32.const 2
     call 0
     i32.const 1
     if  ;; label = @1
-      i64.const 1
+      i32.const 1
       call 0
       loop  ;; label = @2
-        i64.const 2
+        i32.const 2
         call 0
         i32.const 123
         drop
@@ -18,7 +18,7 @@
     end
   )
   (func (;2;) (type 0)
-    i64.const 1
+    i32.const 1
     call 0
     block  ;; label = @1
     end

tests/expectations/gas/start.wat

@@ -1,12 +1,12 @@
 (module
   (type (;0;) (func (param i32 i32)))
   (type (;1;) (func))
-  (type (;2;) (func (param i64)))
+  (type (;2;) (func (param i32)))
   (import "env" "ext_return" (func $ext_return (;0;) (type 0)))
   (import "env" "memory" (memory (;0;) 1 1))
   (import "env" "gas" (func (;1;) (type 2)))
   (func $start (;2;) (type 1)
-    i64.const 4
+    i32.const 4
     call 1
     i32.const 8
     i32.const 4