WIP: Known values

Cargo.toml

@@ -22,6 +22,9 @@ 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"
@@ -36,3 +39,4 @@ wasmprinter = "0.2"
 default = ["std"]
 std = ["parity-wasm/std"]
 sign_ext = ["parity-wasm/sign_ext"]
+trace-log = ["dep:log", "dep:test-log", "dep:env_logger"]

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,15 +1,30 @@
 use super::resolve_func_type;
-use alloc::vec::Vec;
-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;
 
+#[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)]
@@ -17,37 +32,50 @@ struct Frame {
 	/// Stack becomes polymorphic only after an instruction that
 	/// never passes control further was executed.
 	is_polymorphic: bool,
+	unreachable_depth: u32,
 
-	/// Count of values which will be pushed after the exit
-	/// from the current block.
-	end_arity: u32,
+	/// Type of value which will be pushed after exiting
+	/// the current block or `None` if block does not return a result.
+	result_type: Option<ValueType>,
 
-	/// Count of values which should be poped upon a branch to
-	/// this frame.
+	/// Type of value which should be poped upon a branch to
+	/// 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
+#[derive(Clone)]
+struct StackValue (ValueType, bool);
+
+/// 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<StackValue>,
 	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() }
+	}
+
+	// fn new_from(stack: &Stack) -> Stack {
+	// 	Stack { values: stack.values.clone(), control_stack: stack.control_stack.clone() }
+	// } 
+
+	/// Returns current weight of the value stack.
+	fn weight(&self) -> u32 {
+		self.values.iter().map(|v| value_cost(v.0)).sum()
 	}
 
 	/// Returns current height of the value stack.
-	fn height(&self) -> u32 {
-		self.height
+	fn height(&self) -> usize {
+		self.values.len()
 	}
 
 	/// Returns a reference to a frame by specified depth relative to the top of
@@ -65,64 +93,96 @@ impl Stack {
 	fn mark_unreachable(&mut self) -> Result<(), &'static str> {
 		let top_frame = self.control_stack.last_mut().ok_or("stack must be non-empty")?;
 		top_frame.is_polymorphic = true;
+		top_frame.unreachable_depth = 1;
 		Ok(())
 	}
 
+	fn push_unreachable(&mut self) -> Result<(), &'static str> {
+		let top_frame = self.control_stack.last_mut().ok_or("stack must be non-empty")?;
+		top_frame.unreachable_depth += 1;
+		Ok(())
+	}
+
+	fn pop_unreachable(&mut self) -> Result<u32, &'static str> {
+		let top_frame = self.control_stack.last_mut().ok_or("stack must be non-empty")?;
+		top_frame.unreachable_depth -= 1;
+		Ok(top_frame.unreachable_depth)
+	}
+
 	/// 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) {
-		self.height = new_height;
+	fn trunc(&mut self, new_height: usize) {
+		trace!("  Truncate value stack to {}", new_height);
+		self.values.truncate(new_height);
 	}
 
-	/// Push specified number of values into the value stack.
-	///
-	/// Returns `Err` if the height overflow usize value.
-	fn push_values(&mut self, value_count: u32) -> Result<(), &'static str> {
-		self.height = self.height.checked_add(value_count).ok_or("stack overflow")?;
+	/// Push a value into the value stack.
+	fn push_value(&mut self, value: StackValue) -> Result<(), &'static str> {
+		trace!("  Push {:?} to value stack", value);
+		self.values.push(value);
+		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.
-	fn pop_values(&mut self, value_count: u32) -> Result<(), &'static str> {
-		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")
-				}
+	/// value popped.
+	fn pop_value(&mut self) -> Result<Option<StackValue>, &'static str> {
+		let top_frame = self.frame(0)?;
+		if self.height() == top_frame.start_height {
+			return if top_frame.is_polymorphic {
+				Ok(None)
+			} else {
+				Err("trying to pop more values than pushed")
 			}
 		}
 
-		self.height = self.height.checked_sub(value_count).ok_or("stack underflow")?;
-
-		Ok(())
+		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")
+		}
 	}
 }
 
+fn value_cost(val: ValueType) -> u32 {
+	match val {
+		ValueType::I32 | ValueType::F32 => 4,
+		ValueType::I64 | ValueType::F64 => 8,
+	}
+}
+
+// struct FunctionContext {
+// 	globals: Vec<ValueType>,
+// 	locals: Vec<ValueType>,
+// 	stack: Stack,
+// 	result_type: Option<ValueType>
+// }
+
 /// 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<(u32, u32), &'static str> {
 	use parity_wasm::elements::Instruction::*;
 
 	let func_section = module.function_section().ok_or("No function section")?;
@@ -145,47 +205,78 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 		.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: Vec<StackValue> = func_signature.params().iter().map(|p| StackValue(*p, false)).collect();
+	locals.extend(body.locals().iter().flat_map(|l| vec![StackValue(l.value_type(), true); l.count() as usize]));
+
 	let mut stack = Stack::new();
-	let mut max_height: u32 = 0;
-	let mut pc = 0;
+	let mut max_weight: u32 = 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,
-		branch_arity: func_arity,
+		unreachable_depth: 0,
+		result_type: func_result_type,
+		branch_type: func_result_type,
 		start_height: 0,
 	});
 
-	loop {
-		if pc >= instructions.elements().len() {
-			break
+	for opcode in instructions.elements() {
+		let current_frame = stack.frame(0)?;
+		if current_frame.is_polymorphic {
+			match opcode {
+				Block(ty) | Loop(ty) | If(ty) => {
+					trace!("Entering unreachable block {:?}", opcode);
+					stack.push_unreachable()?;
+				},
+				End => {
+					let depth = stack.pop_unreachable()?;
+					if depth == 0 {
+						trace!("Exiting unreachable code");
+						stack.pop_frame()?;
+					} else {
+						trace!("Exiting unreachable block");
+					}
+				},
+				_ => {
+					trace!("Skipping unreachable instruction {:?}", opcode);
+				}
+			}
+			continue;
 		}
 
-		// 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];
+		trace!("Processing opcode {:?}", opcode);
 
 		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,
-					end_arity,
-					branch_arity,
+					unreachable_depth: 0,
+					result_type: end_result,
+					branch_type: if let Loop(_) = *opcode { None } else { end_result },
 					start_height: height,
 				});
 			},
@@ -194,47 +285,58 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 				// it as is.
 			},
 			End => {
-				let frame = stack.pop_frame()?;
-				stack.trunc(frame.start_height);
-				stack.push_values(frame.end_arity)?;
+				// let frame = stack.pop_frame()?;
+				// stack.trunc(frame.start_height);
+				// 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;
-				stack.pop_values(target_arity)?;
+				// if stack.frame(*target)?.branch_type.is_some() {
+				// 	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) => {
-				// Pop values for the destination block result.
-				let target_arity = stack.frame(*target)?.branch_arity;
-				stack.pop_values(target_arity)?;
+				// let target_type = stack.frame(*target)?.branch_type;
+				// // Pop values for the destination block result.
+				// if target_type.is_some() {
+				// 	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 arity != arity_of_default {
-						return Err("Arity of all jump-targets must be equal")
-					}
-				}
+				// Check that all jump targets have an equal number of parameters
+				// for target in &*br_table_data.table {
+				// 	if stack.frame(*target)?.branch_type != default_type {
+				// 		return Err("Types 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.
@@ -243,80 +345,118 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 			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(idx) => {
 				let ty = resolve_func_type(*idx, module)?;
 
 				// 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;
-				stack.push_values(callee_arity)?;
+				let callee_results = ty.results();
+				if !callee_results.is_empty() {
+					stack.push_value(StackValue(callee_results[0], false))?;
+				}
 			},
 			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;
-				stack.push_values(callee_arity)?;
+				let callee_results = ty.results();
+				if !callee_results.is_empty() {
+					stack.push_value(StackValue(callee_results[0], false))?;
+				}
 			},
 			Drop => {
-				stack.pop_values(1)?;
+				stack.pop_value()?;
 			},
 			Select => {
 				// Pop two values and one condition.
-				stack.pop_values(2)?;
-				stack.pop_values(1)?;
+				let val = stack.pop_value()?;
+				stack.pop_value()?;
+				stack.pop_value()?;
 
 				// Push the selected value.
-				stack.push_values(1)?;
+				if let Some(vt) = val {
+					stack.push_value(vt)?;
+				}
 			},
-			GetLocal(_) => {
-				stack.push_values(1)?;
+			GetLocal(idx) => {
+				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)?;
-				stack.push_values(1)?;
+				let idx = *idx as usize;
+				if idx >= locals.len() {
+					return Err("Reference to a local is out of bounds")
+				}
+				stack.pop_value()?;
+				stack.push_value(locals[idx])?;
 			},
-			GetGlobal(_) => {
-				stack.push_values(1)?;
+			GetGlobal(idx) => {
+				let idx = *idx as usize;
+				if idx >= globals.len() {
+					return Err("Reference to a global is out of bounds")
+				}
+				stack.push_value(StackValue(globals[idx], false))?;
 			},
 			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(_, _) => {
+				if let Some(sv) = stack.pop_value()? {
+					stack.push_value(StackValue(ValueType::I32, sv.1))?;
+				}
+			},
+			I64Load(_, _) |
 			I64Load8S(_, _) |
 			I64Load8U(_, _) |
 			I64Load16S(_, _) |
 			I64Load16U(_, _) |
 			I64Load32S(_, _) |
 			I64Load32U(_, _) => {
-				// These instructions pop the address and pushes the result,
-				// which effictively don't modify the stack height.
-				stack.pop_values(1)?;
-				stack.push_values(1)?;
+				if let Some(sv) = stack.pop_value()? {
+					stack.push_value(StackValue(ValueType::I64, sv.1))?;
+				}
+			},
+			F32Load(_, _) => {
+				if let Some(sv) = stack.pop_value()? {
+					stack.push_value(StackValue(ValueType::F32, sv.1))?;
+				}
+			},
+			F64Load(_, _) => {
+				if let Some(sv) = stack.pop_value()? {
+					stack.push_value(StackValue(ValueType::F64, sv.1))?;
+				}
 			},
 
 			I32Store(_, _) |
@@ -329,29 +469,39 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 			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(StackValue(ValueType::I32, false))?;
 			},
 			GrowMemory(_) => {
 				// Grow memory takes the value of pages to grow and pushes
-				stack.pop_values(1)?;
-				stack.push_values(1)?;
+				stack.pop_value()?;
+				stack.push_value(StackValue(ValueType::I32, false))?;
 			},
 
-			I32Const(_) | I64Const(_) | F32Const(_) | F64Const(_) => {
-				// These instructions just push the single literal value onto the stack.
-				stack.push_values(1)?;
+			I32Const(_) => {
+				stack.push_value(StackValue(ValueType::I32, true))?;
+			},
+			I64Const(_) => {
+				stack.push_value(StackValue(ValueType::I64, true))?;
+			},
+			F32Const(_) => {
+				stack.push_value(StackValue(ValueType::F32, true))?;
+			},
+			F64Const(_) => {
+				stack.push_value(StackValue(ValueType::F64, true))?;
 			},
 
 			I32Eqz | I64Eqz => {
 				// These instructions pop the value and compare it against zero, and pushes
 				// the result of the comparison.
-				stack.pop_values(1)?;
-				stack.push_values(1)?;
+				if let Some(sv) = stack.pop_value()? {
+					stack.push_value(StackValue(ValueType::I32, sv.1))?;
+				}
 			},
 
 			I32Eq | I32Ne | I32LtS | I32LtU | I32GtS | I32GtU | I32LeS | I32LeU | I32GeS |
@@ -359,16 +509,18 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 			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.push_values(1)?;
+				let Some(op1) = stack.pop_value()?;
+				let Some(op2) = stack.pop_value()?;
+				stack.push_value(StackValue(ValueType::I32, op1.1 && op2.1))?;
 			},
 
 			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)?;
-				stack.push_values(1)?;
+				if let Some(sv) = stack.pop_value()? {
+					stack.push_value(sv)?;
+				}
 			},
 
 			I32Add | I32Sub | I32Mul | I32DivS | I32DivU | I32RemS | I32RemU | I32And | I32Or |
@@ -378,19 +530,36 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 			F32Min | F32Max | F32Copysign | F64Add | F64Sub | F64Mul | F64Div | F64Min |
 			F64Max | F64Copysign => {
 				// Binary operators take two operands and produce one result.
-				stack.pop_values(2)?;
-				stack.push_values(1)?;
+				let Some(op1) = stack.pop_value()?;
+				let Some(op2) = stack.pop_value()?;
+				stack.push_value(StackValue(op1.0, op1.1 && op2.1))?;
 			},
 
+			// Conversion operators take one value and produce one result.
 			I32WrapI64 | I32TruncSF32 | I32TruncUF32 | I32TruncSF64 | I32TruncUF64 |
+			I32ReinterpretF32 => {
+				if let Some(sv) = stack.pop_value()? {
+					stack.push_value(StackValue(ValueType::I32, sv.1))?;
+				}
+			},
 			I64ExtendSI32 | I64ExtendUI32 | I64TruncSF32 | I64TruncUF32 | I64TruncSF64 |
-			I64TruncUF64 | F32ConvertSI32 | F32ConvertUI32 | F32ConvertSI64 | F32ConvertUI64 |
-			F32DemoteF64 | F64ConvertSI32 | F64ConvertUI32 | F64ConvertSI64 | F64ConvertUI64 |
-			F64PromoteF32 | I32ReinterpretF32 | I64ReinterpretF64 | F32ReinterpretI32 |
+			I64TruncUF64 | I64ReinterpretF64 => {
+				if let Some(sv) = stack.pop_value()? {
+					stack.push_value(StackValue(ValueType::I64, sv.1))?;
+				}
+			},
+			F32ConvertSI32 | F32ConvertUI32 | F32ConvertSI64 | F32ConvertUI64 | F32DemoteF64 |
+			F32ReinterpretI32 => {
+				if let Some(sv) = stack.pop_value()? {
+					stack.push_value(StackValue(ValueType::F32, sv.1))?;
+				}
+			},
+
+			F64ConvertSI32 | F64ConvertUI32 | F64ConvertSI64 | F64ConvertUI64 | F64PromoteF32 |
 			F64ReinterpretI64 => {
-				// Conversion operators take one value and produce one result.
-				stack.pop_values(1)?;
-				stack.push_values(1)?;
+				if let Some(sv) = stack.pop_value()? {
+					stack.push_value(StackValue(ValueType::F64, sv.1))?;
+				}
 			},
 
 			#[cfg(feature = "sign_ext")]
@@ -398,15 +567,45 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 			SignExt(SignExtInstruction::I32Extend16S) |
 			SignExt(SignExtInstruction::I64Extend8S) |
 			SignExt(SignExtInstruction::I64Extend16S) |
-			SignExt(SignExtInstruction::I64Extend32S) => {
-				stack.pop_values(1)?;
-				stack.push_values(1)?;
-			},
+			SignExt(SignExtInstruction::I64Extend32S) =>
+				if let Some(sv) = stack.pop_value()? {
+					stack.push_value(sv)?;
+				},
+		}
+
+		// 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)]
@@ -414,6 +613,9 @@ mod tests {
 	use super::*;
 	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")
@@ -436,8 +638,8 @@ mod tests {
 "#,
 		);
 
-		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 3 + ACTIVATION_FRAME_COST);
+		let res = compute(0, &module).unwrap();
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 3, ACTIVATION_FRAME_WEIGHT + 12));
 	}
 
 	#[test]
@@ -446,15 +648,15 @@ mod tests {
 			r#"
 (module
 	(func (result i32)
-		i32.const 0
+		i64.const 0
 		return
 	)
 )
 "#,
 		);
 
-		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
+		let res = compute(0, &module).unwrap();
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 1, ACTIVATION_FRAME_WEIGHT + 8));
 	}
 
 	#[test]
@@ -471,8 +673,8 @@ mod tests {
 "#,
 		);
 
-		let height = compute(0, &module).unwrap();
-		assert_eq!(height, ACTIVATION_FRAME_COST);
+		let res = compute(0, &module).unwrap();
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT, ACTIVATION_FRAME_WEIGHT));
 	}
 
 	#[test]
@@ -500,8 +702,8 @@ mod tests {
 "#,
 		);
 
-		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 2 + ACTIVATION_FRAME_COST);
+		let res = compute(0, &module).unwrap();
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 2, ACTIVATION_FRAME_WEIGHT + 8));
 	}
 
 	#[test]
@@ -524,8 +726,8 @@ mod tests {
 "#,
 		);
 
-		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
+		let res = compute(0, &module).unwrap();
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 1, ACTIVATION_FRAME_WEIGHT + 4));
 	}
 
 	#[test]
@@ -546,8 +748,8 @@ mod tests {
 "#,
 		);
 
-		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
+		let res = compute(0, &module).unwrap();
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 1, ACTIVATION_FRAME_WEIGHT + 4));
 	}
 
 	#[test]
@@ -572,7 +774,7 @@ mod tests {
 "#,
 		);
 
-		let height = compute(0, &module).unwrap();
-		assert_eq!(height, 3 + ACTIVATION_FRAME_COST);
+		let res = compute(0, &module).unwrap();
+		assert_eq!(res, (ACTIVATION_FRAME_HEIGHT + 3, ACTIVATION_FRAME_WEIGHT + 12));
 	}
 }

src/stack_limiter/mod.rs

@@ -154,7 +154,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> {
+pub fn compute_stack_costs(module: &elements::Module) -> Result<Vec<u32>, &'static str> {
 	let func_imports = module.import_count(elements::ImportCountType::Function);
 
 	// TODO: optimize!
@@ -173,7 +173,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<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;
@@ -194,13 +194,27 @@ fn compute_stack_cost(func_idx: u32, module: &elements::Module) -> Result<u32, &
 			locals_count.checked_add(local_group.count()).ok_or("Overflow in local count")?;
 	}
 
-	let max_stack_height = max_height::compute(defined_func_idx, module)?;
+	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,