Return detailed stats

src/stack_limiter/max_height.rs

@@ -11,6 +11,17 @@ use parity_wasm::elements::SignExtInstruction;
 // machine might be consumed to hold some context.
 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 total_cost: u32,
+}
+
 /// Control stack frame.
 #[derive(Debug)]
 struct Frame {
@@ -125,7 +136,7 @@ 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);
@@ -161,6 +172,7 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 	}
 
 	let params_count = func_signature.params().len() as u32;
+	let mut blocks_count = 0u32;
 
 	// Add implicit frame for the function. Breaks to this frame and execution of
 	// the last end should deal with this frame.
@@ -223,6 +235,7 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 					branch_arity,
 					start_height: stack.height(),
 				});
+				blocks_count += 1;
 			},
 			Else => {
 				let frame = stack.pop_frame()?;
@@ -451,17 +464,19 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 	assert_eq!(stack.control_height(), 0);
 	assert_eq!(stack.height(), func_signature.results().len() as u32);
 
-	trace!(
-		"Result: {} activation + {} stack + {} control stack + {} locals + {} params = {}",
-		ACTIVATION_FRAME_COST,
-		max_height,
-		max_control_height,
-		locals_count,
-		params_count,
-		ACTIVATION_FRAME_COST + max_height + max_control_height + locals_count + params_count
-	);
+	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,
+		total_cost: ACTIVATION_FRAME_COST + max_height + max_control_height + locals_count + params_count,
+	};
 
-	Ok(ACTIVATION_FRAME_COST + max_height + max_control_height + locals_count + params_count)
+	trace!("Result: {:?}", res);
+
+	Ok(res)
 }
 
 #[cfg(test)]

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.
-pub 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;