From 4c0f42c6fc6de1f9cf0804e33728dd9ac207d755 Mon Sep 17 00:00:00 2001
From: Jim Posen <jim.posen@gmail.com>
Date: Tue, 2 Jul 2019 15:41:33 +0200
Subject: [PATCH 1/9] Perform gas metering injection step in linear time.

Previously the code was quadratic in the worst case as inserting into
the middle of a vector is a linear-time operation.
---
 src/gas.rs | 51 +++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 43 insertions(+), 8 deletions(-)

diff --git a/src/gas.rs b/src/gas.rs
index b3d377e..27bc4da 100644
--- a/src/gas.rs
+++ b/src/gas.rs
@@ -4,6 +4,7 @@
 //! module into one that charges gas for code to be executed. See function documentation for usage
 //! and details.
 
+use std::mem;
 use std::vec::Vec;
 
 use parity_wasm::{elements, builder};
@@ -176,16 +177,50 @@ pub fn inject_counter(
 		}
 	}
 
-	// Then insert metering calls.
-	let mut cumulative_offset = 0;
-	for block in counter.blocks {
-		let effective_pos = block.start_pos + cumulative_offset;
+	insert_metering_calls(instructions, counter.blocks, gas_func)
+}
 
-		instructions.elements_mut().insert(effective_pos, I32Const(block.cost as i32));
-		instructions.elements_mut().insert(effective_pos+1, Call(gas_func));
+// Then insert metering calls into a sequence of instructions given the block locations and costs.
+fn insert_metering_calls(
+	instructions: &mut elements::Instructions,
+	blocks: Vec<BlockEntry>,
+	gas_func: u32,
+)
+	-> Result<(), ()>
+{
+	use parity_wasm::elements::Instruction::*;
 
-		// Take into account these two inserted instructions.
-		cumulative_offset += 2;
+	// To do this in linear time, construct a new vector of instructions, copying over old
+	// instructions one by one and injecting new ones as required.
+	let new_instrs_len = instructions.elements().len() + 2 * blocks.len();
+	let original_instrs = mem::replace(
+		instructions.elements_mut(), Vec::with_capacity(new_instrs_len)
+	);
+	let new_instrs = instructions.elements_mut();
+
+	let mut original_pos = 0;
+	let mut block_iter = blocks.into_iter().peekable();
+	for instr in original_instrs.into_iter() {
+		// If there the next block starts at this position, inject metering instructions.
+		let used_block = if let Some(ref block) = block_iter.peek() {
+			if block.start_pos == original_pos {
+				new_instrs.push(I32Const(block.cost as i32));
+				new_instrs.push(Call(gas_func));
+				true
+			} else { false }
+		} else { false };
+
+		if used_block {
+			block_iter.next();
+		}
+
+		// Copy over the original instruction.
+		new_instrs.push(instr);
+		original_pos += 1;
+	}
+
+	if block_iter.next().is_some() {
+		return Err(());
 	}
 
 	Ok(())

From de60f491b44a1c4cb735da2f494aff7462190282 Mon Sep 17 00:00:00 2001
From: Jim Posen <jim.posen@gmail.com>
Date: Mon, 1 Jul 2019 17:32:37 +0200
Subject: [PATCH 2/9] Fix ordering of actual and expected arguments in
 assert_eq!.

---
 src/gas.rs | 50 ++++++++++++++++++++++----------------------------
 1 file changed, 22 insertions(+), 28 deletions(-)

diff --git a/src/gas.rs b/src/gas.rs
index 27bc4da..d0a243a 100644
--- a/src/gas.rs
+++ b/src/gas.rs
@@ -342,6 +342,14 @@ mod tests {
 	use super::*;
 	use rules;
 
+	fn get_function_body(module: &elements::Module, index: usize)
+		-> Option<&[elements::Instruction]>
+	{
+		module.code_section()
+			.and_then(|code_section| code_section.bodies().get(index))
+			.map(|func_body| func_body.code().elements())
+	}
+
 	#[test]
 	fn simple_grow() {
 		use parity_wasm::elements::Instruction::*;
@@ -367,18 +375,17 @@ mod tests {
 		let injected_module = inject_gas_counter(module, &rules::Set::default().with_grow_cost(10000)).unwrap();
 
 		assert_eq!(
+			get_function_body(&injected_module, 0).unwrap(),
 			&vec![
 				I32Const(3),
 				Call(0),
 				GetGlobal(0),
 				Call(2),
 				End
-			][..],
-			injected_module
-				.code_section().expect("function section should exist").bodies()[0]
-				.code().elements()
+			][..]
 		);
 		assert_eq!(
+			get_function_body(&injected_module, 1).unwrap(),
 			&vec![
 				GetLocal(0),
 				GetLocal(0),
@@ -387,10 +394,7 @@ mod tests {
 				Call(0),
 				GrowMemory(0),
 				End,
-			][..],
-			injected_module
-				.code_section().expect("function section should exist").bodies()[1]
-				.code().elements()
+			][..]
 		);
 
 		let binary = serialize(injected_module).expect("serialization failed");
@@ -422,16 +426,14 @@ mod tests {
 		let injected_module = inject_gas_counter(module, &rules::Set::default()).unwrap();
 
 		assert_eq!(
+			get_function_body(&injected_module, 0).unwrap(),
 			&vec![
 				I32Const(3),
 				Call(0),
 				GetGlobal(0),
 				GrowMemory(0),
 				End
-			][..],
-			injected_module
-				.code_section().expect("function section should exist").bodies()[0]
-				.code().elements()
+			][..]
 		);
 
 		assert_eq!(injected_module.functions_space(), 2);
@@ -464,15 +466,13 @@ mod tests {
 		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
 
 		assert_eq!(
+			get_function_body(&injected_module, 0).unwrap(),
 			&vec![
 				I32Const(2),
 				Call(0),
 				GetGlobal(0),
 				End
-			][..],
-			injected_module
-				.code_section().expect("function section should exist").bodies()[0]
-				.code().elements()
+			][..]
 		);
 	}
 
@@ -506,6 +506,7 @@ mod tests {
 		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
 
 		assert_eq!(
+			get_function_body(&injected_module, 0).unwrap(),
 			&vec![
 				I32Const(4),
 				Call(0),
@@ -519,10 +520,7 @@ mod tests {
 				End,
 				GetGlobal(0),
 				End
-			][..],
-			injected_module
-				.code_section().expect("function section should exist").bodies()[0]
-				.code().elements()
+			][..]
 		);
 	}
 
@@ -559,6 +557,7 @@ mod tests {
 		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
 
 		assert_eq!(
+			get_function_body(&injected_module, 0).unwrap(),
 			&vec![
 				I32Const(4),
 				Call(0),
@@ -577,10 +576,7 @@ mod tests {
 				End,
 				GetGlobal(0),
 				End
-			][..],
-			injected_module
-				.code_section().expect("function section should exist").bodies()[0]
-				.code().elements()
+			][..]
 		);
 	}
 
@@ -621,6 +617,7 @@ mod tests {
 		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
 
 		assert_eq!(
+			get_function_body(&injected_module, 1).unwrap(),
 			&vec![
 				I32Const(4),
 				Call(0),
@@ -639,10 +636,7 @@ mod tests {
 				End,
 				Call(1),
 				End
-			][..],
-			injected_module
-				.code_section().expect("function section should exist").bodies()[1]
-				.code().elements()
+			][..]
 		);
 	}
 

From 24924f59ec2daa300b66f42d6c4c01d31112e907 Mon Sep 17 00:00:00 2001
From: Jim Posen <jim.posen@gmail.com>
Date: Mon, 1 Jul 2019 16:57:45 +0200
Subject: [PATCH 3/9] Change gas metering injection code to handle branches
 properly.

---
 src/gas.rs | 509 +++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 433 insertions(+), 76 deletions(-)

diff --git a/src/gas.rs b/src/gas.rs
index d0a243a..a26dd69 100644
--- a/src/gas.rs
+++ b/src/gas.rs
@@ -4,6 +4,7 @@
 //! module into one that charges gas for code to be executed. See function documentation for usage
 //! and details.
 
+use std::cmp::min;
 use std::mem;
 use std::vec::Vec;
 
@@ -19,9 +20,9 @@ pub fn update_call_index(instructions: &mut elements::Instructions, inserted_ind
 	}
 }
 
-/// A block of code represented by it's start position and cost.
-///
-/// The block typically starts with instructions such as `loop`, `block`, `if`, etc.
+/// A control flow block is opened with the `block`, `loop`, and `if` instructions and is closed
+/// with `end`. Each block implicitly defines a new label. The control blocks form a stack during
+/// program execution.
 ///
 /// An example of block:
 ///
@@ -38,54 +39,181 @@ pub fn update_call_index(instructions: &mut elements::Instructions, inserted_ind
 /// The start of the block is `i32.const 1`.
 ///
 #[derive(Debug)]
-struct BlockEntry {
+struct ControlBlock {
+	/// The lowest control stack index corresponding to a forward jump targeted by a br, br_if, or
+	/// br_table instruction within this control block. The index must refer to a control block
+	/// that is not a loop, meaning it is a forward jump.
+	///
+	/// This value will always be at most the index of the block itself, even if there is no
+	/// explicit br instruction targeting this control block. This does not affect how the value is
+	/// used in the metering algorithm.
+	lowest_forward_br: usize,
+
+	/// The active metering block that new instructions contribute a gas cost towards.
+	active_metered_block: MeteredBlock,
+
+	/// Whether the control block is a loop. Loops have the distinguishing feature that branches to
+	/// them jump to the beginning of the block, not the end as with the other control blocks.
+	is_loop: bool,
+}
+
+/// A block of code that metering instructions will be inserted at the beginning of. Metered blocks
+/// are constructed with the property that, in the absence of any traps, either all instructions in
+/// the block are executed or none are.
+#[derive(Debug)]
+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: u32,
 }
 
+/// Counter is used to manage state during the gas metering algorithm implemented by
+/// `inject_counter`.
 struct Counter {
-	/// All blocks in the order of theirs start position.
-	blocks: Vec<BlockEntry>,
+	/// A stack of control blocks. This stack grows when new control blocks are opened with
+	/// `block`, `loop`, and `if` and shrinks when control blocks are closed with `end`. The first
+	/// block on the stack corresponds to the function body, not to any labelled block. Therefore
+	/// the actual Wasm label index associated with each control block is 1 less than its position
+	/// in this stack.
+	stack: Vec<ControlBlock>,
 
-	// Stack of blocks. Each element is an index to a `self.blocks` vector.
-	stack: Vec<usize>,
+	/// A list of metered blocks that have been finalized, meaning they will no longer change.
+	finalized_blocks: Vec<MeteredBlock>,
 }
 
 impl Counter {
 	fn new() -> Counter {
 		Counter {
 			stack: Vec::new(),
-			blocks: Vec::new(),
+			finalized_blocks: Vec::new(),
 		}
 	}
 
-	/// Begin a new block.
-	fn begin(&mut self, cursor: usize) {
-		let block_idx = self.blocks.len();
-		self.blocks.push(BlockEntry {
-			start_pos: cursor,
-			cost: 1,
-		});
-		self.stack.push(block_idx);
+	/// Open a new control block. The cursor is the position of the first instruction in the block.
+	fn begin_control_block(&mut self, cursor: usize, is_loop: bool) {
+		let index = self.stack.len();
+		self.stack.push(ControlBlock {
+			lowest_forward_br: index,
+			active_metered_block: MeteredBlock {
+				start_pos: cursor,
+				cost: 0,
+			},
+			is_loop,
+		})
 	}
 
-	/// Finalize the current block.
+	/// Close the last control block. The cursor is the position of the final (pseudo-)instruction
+	/// in the block.
+	fn finalize_control_block(&mut self, cursor: usize) -> Result<(), ()> {
+		// This either finalizes the active metered block or merges its cost into the active
+		// metered block in the previous control block on the stack.
+		self.finalize_metered_block(cursor)?;
+
+		// Pop the control block stack.
+		let closing_control_block = self.stack.pop().ok_or_else(|| ())?;
+		let closing_control_index = self.stack.len();
+
+		if self.stack.is_empty() {
+			return Ok(())
+		}
+
+		// Update the lowest_forward_br for the control block now on top of the stack.
+		{
+			let control_block = self.stack.last_mut().ok_or_else(|| ())?;
+			control_block.lowest_forward_br = min(
+				control_block.lowest_forward_br, closing_control_block.lowest_forward_br
+			);
+		}
+
+		// If there may have been a branch to a lower index, then also finalize the active metered
+		// block for the previous control block. Otherwise, finalize it and begin a new one.
+		//
+		// TODO: handle diverging function calls.
+		let may_br_out = closing_control_block.lowest_forward_br < closing_control_index;
+		if may_br_out {
+			self.finalize_metered_block(cursor)?;
+		}
+
+		Ok(())
+	}
+
+	/// Finalize the current active metered block.
 	///
 	/// Finalized blocks have final cost which will not change later.
-	fn finalize(&mut self) -> Result<(), ()> {
-		self.stack.pop().ok_or_else(|| ())?;
+	fn finalize_metered_block(&mut self, cursor: usize) -> Result<(), ()> {
+		let closing_metered_block = {
+			let control_block = self.stack.last_mut().ok_or_else(|| ())?;
+			mem::replace(
+				&mut control_block.active_metered_block,
+				MeteredBlock {
+					start_pos: cursor + 1,
+					cost: 0,
+				}
+			)
+		};
+
+		// If the block was opened with a `begin`, then its start position will be set to that of
+		// the active metered block in the control block one higher on the stack. This is because
+		// any instructions between a `begin` and the first branch are part of the same basic block
+		// as the preceding instruction. In this case, instead of finalizing the block, merge its
+		// cost into the other active metered block to avoid injecting unnecessary instructions.
+		let last_index = self.stack.len() - 1;
+		if last_index > 0 {
+			let prev_control_block = self.stack.get_mut(last_index - 1)
+				.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 += closing_metered_block.cost;
+				return Ok(())
+			}
+		}
+
+		if closing_metered_block.cost > 0 {
+			self.finalized_blocks.push(closing_metered_block);
+		}
 		Ok(())
 	}
 
+	/// Handle a branch instruction in the program. The cursor is the index of the branch
+	/// instruction in the program. The indices are the stack positions of the target control
+	/// blocks. Recall that the index is 0 for a `return` and relatively indexed from the top of
+	/// the stack by the label of `br`, `br_if`, and `br_table` instructions.
+	fn branch(&mut self, cursor: usize, indices: &[usize]) -> Result<(), ()> {
+		self.finalize_metered_block(cursor)?;
+
+		// Update the lowest_forward_br of the current control block.
+		for index in indices.iter().cloned() {
+			let target_is_loop = {
+				let target_block = self.stack.get(index).ok_or_else(|| ())?;
+				target_block.is_loop
+			};
+			if target_is_loop {
+				continue;
+			}
+
+			let control_block = self.stack.last_mut().ok_or_else(|| ())?;
+			control_block.lowest_forward_br = min(control_block.lowest_forward_br, index);
+		}
+
+		Ok(())
+	}
+
+	/// Returns the height of the control stack.
+	fn stack_height(&self) -> usize {
+		self.stack.len()
+	}
+
+	/// Get a reference to the currently active metered block.
+	fn active_block(&mut self) -> Result<&mut MeteredBlock, ()> {
+		let top_block = self.stack.last_mut().ok_or_else(|| ())?;
+		Ok(&mut top_block.active_metered_block)
+	}
+
 	/// Increment the cost of the current block by the specified value.
 	fn increment(&mut self, val: u32) -> Result<(), ()> {
-		let stack_top = self.stack.last_mut().ok_or_else(|| ())?;
-		let top_block = self.blocks.get_mut(*stack_top).ok_or_else(|| ())?;
-
+		let top_block = self.active_block()?;
 		top_block.cost = top_block.cost.checked_add(val).ok_or_else(|| ())?;
-
 		Ok(())
 	}
 }
@@ -137,53 +265,74 @@ pub fn inject_counter(
 	let mut counter = Counter::new();
 
 	// Begin an implicit function (i.e. `func...end`) block.
-	counter.begin(0);
+	counter.begin_control_block(0, false);
 
 	for cursor in 0..instructions.elements().len() {
 		let instruction = &instructions.elements()[cursor];
+		let instruction_cost = rules.process(instruction)?;
 		match *instruction {
-			Block(_) | If(_) | Loop(_) => {
-				// Increment previous block with the cost of the current opcode.
-				let instruction_cost = rules.process(instruction)?;
+			Block(_) => {
 				counter.increment(instruction_cost)?;
 
-				// Begin new block. The cost of the following opcodes until `End` or `Else` will
-				// be included into this block.
-				counter.begin(cursor + 1);
+				// Begin new block. The cost of the following opcodes until `end` or `else` will
+				// be included into this block. The start position is set to that of the previous
+				// active metered block to signal that they should be merged in order to reduce
+				// unnecessary metering instructions.
+				let top_block_start_pos = counter.active_block()
+					.expect("stack is not empty; qed")
+					.start_pos;
+				counter.begin_control_block(top_block_start_pos, false);
+			}
+			If(_) => {
+				counter.increment(instruction_cost)?;
+				counter.begin_control_block(cursor + 1, false);
+			}
+			Loop(_) => {
+				counter.increment(instruction_cost)?;
+				counter.begin_control_block(cursor + 1, true);
 			}
 			End => {
-				// Just finalize current block.
-				counter.finalize()?;
+				counter.finalize_control_block(cursor)?;
 			},
 			Else => {
-				// `Else` opcode is being encountered. So the case we are looking at:
-				//
-				// if
-				//   ...
-				// else <-- cursor
-				//   ...
-				// end
-				//
-				// Finalize the current block ('then' part of the if statement),
-				// and begin another one for the 'else' part.
-				counter.finalize()?;
-				counter.begin(cursor + 1);
+				counter.finalize_metered_block(cursor)?;
+			}
+			Br(label) | BrIf(label) => {
+				counter.increment(instruction_cost)?;
+
+				// Label is a relative index into the control stack.
+				let index = counter.stack_height().checked_sub(1 + label as usize)
+					.ok_or_else(|| ())?;
+				counter.branch(cursor, &[index])?;
+			}
+			BrTable(ref label_vec, label_default) => {
+				counter.increment(instruction_cost)?;
+
+				let stack_height = counter.stack_height();
+				let indices = [label_default].iter().chain(label_vec.iter())
+					.map(|label| stack_height.checked_sub(1 + *label as usize))
+					.collect::<Option<Vec<_>>>()
+					.ok_or_else(|| ())?;
+				counter.branch(cursor, &indices)?;
+			}
+			Return => {
+				counter.increment(instruction_cost)?;
+				counter.branch(cursor, &[0])?;
 			}
 			_ => {
-				// An ordinal non control flow instruction. Just increment the cost of the current block.
-				let instruction_cost = rules.process(instruction)?;
+				// An ordinal non control flow instruction increments the cost of the current block.
 				counter.increment(instruction_cost)?;
 			}
 		}
 	}
 
-	insert_metering_calls(instructions, counter.blocks, gas_func)
+	insert_metering_calls(instructions, counter.finalized_blocks, gas_func)
 }
 
 // Then insert metering calls into a sequence of instructions given the block locations and costs.
 fn insert_metering_calls(
 	instructions: &mut elements::Instructions,
-	blocks: Vec<BlockEntry>,
+	mut blocks: Vec<MeteredBlock>,
 	gas_func: u32,
 )
 	-> Result<(), ()>
@@ -198,8 +347,10 @@ fn insert_metering_calls(
 	);
 	let new_instrs = instructions.elements_mut();
 
-	let mut original_pos = 0;
+	blocks.sort_unstable_by_key(|block| block.start_pos);
 	let mut block_iter = blocks.into_iter().peekable();
+
+	let mut original_pos = 0;
 	for instr in original_instrs.into_iter() {
 		// If there the next block starts at this position, inject metering instructions.
 		let used_block = if let Some(ref block) = block_iter.peek() {
@@ -339,6 +490,7 @@ mod tests {
 	extern crate wabt;
 
 	use parity_wasm::{serialize, builder, elements};
+	use parity_wasm::elements::Instruction::*;
 	use super::*;
 	use rules;
 
@@ -352,8 +504,6 @@ mod tests {
 
 	#[test]
 	fn simple_grow() {
-		use parity_wasm::elements::Instruction::*;
-
 		let module = builder::module()
 			.global()
 				.value_type().i32()
@@ -377,7 +527,7 @@ mod tests {
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
 			&vec![
-				I32Const(3),
+				I32Const(2),
 				Call(0),
 				GetGlobal(0),
 				Call(2),
@@ -403,8 +553,6 @@ mod tests {
 
 	#[test]
 	fn grow_no_gas_no_track() {
-		use parity_wasm::elements::Instruction::*;
-
 		let module = builder::module()
 			.global()
 				.value_type().i32()
@@ -428,7 +576,7 @@ mod tests {
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
 			&vec![
-				I32Const(3),
+				I32Const(2),
 				Call(0),
 				GetGlobal(0),
 				GrowMemory(0),
@@ -444,8 +592,6 @@ mod tests {
 
 	#[test]
 	fn simple() {
-		use parity_wasm::elements::Instruction::*;
-
 		let module = builder::module()
 			.global()
 				.value_type().i32()
@@ -468,7 +614,7 @@ mod tests {
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
 			&vec![
-				I32Const(2),
+				I32Const(1),
 				Call(0),
 				GetGlobal(0),
 				End
@@ -478,8 +624,6 @@ mod tests {
 
 	#[test]
 	fn nested() {
-		use parity_wasm::elements::Instruction::*;
-
 		let module = builder::module()
 			.global()
 				.value_type().i32()
@@ -508,12 +652,10 @@ mod tests {
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
 			&vec![
-				I32Const(4),
+				I32Const(6),
 				Call(0),
 				GetGlobal(0),
 				Block(elements::BlockType::NoResult),
-					I32Const(4),
-					Call(0),
 					GetGlobal(0),
 					GetGlobal(0),
 					GetGlobal(0),
@@ -526,8 +668,6 @@ mod tests {
 
 	#[test]
 	fn ifelse() {
-		use parity_wasm::elements::Instruction::*;
-
 		let module = builder::module()
 			.global()
 				.value_type().i32()
@@ -559,17 +699,17 @@ mod tests {
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
 			&vec![
-				I32Const(4),
+				I32Const(3),
 				Call(0),
 				GetGlobal(0),
 				If(elements::BlockType::NoResult),
-					I32Const(4),
+					I32Const(3),
 					Call(0),
 					GetGlobal(0),
 					GetGlobal(0),
 					GetGlobal(0),
 				Else,
-					I32Const(3),
+					I32Const(2),
 					Call(0),
 					GetGlobal(0),
 					GetGlobal(0),
@@ -582,8 +722,6 @@ mod tests {
 
 	#[test]
 	fn call_index() {
-		use parity_wasm::elements::Instruction::*;
-
 		let module = builder::module()
 			.global()
 				.value_type().i32()
@@ -619,17 +757,17 @@ mod tests {
 		assert_eq!(
 			get_function_body(&injected_module, 1).unwrap(),
 			&vec![
-				I32Const(4),
+				I32Const(3),
 				Call(0),
 				Call(1),
 				If(elements::BlockType::NoResult),
-					I32Const(4),
+					I32Const(3),
 					Call(0),
 					Call(1),
 					Call(1),
 					Call(1),
 				Else,
-					I32Const(3),
+					I32Const(2),
 					Call(0),
 					Call(1),
 					Call(1),
@@ -643,8 +781,6 @@ mod tests {
 
 	#[test]
 	fn forbidden() {
-		use parity_wasm::elements::Instruction::*;
-
 		let module = builder::module()
 			.global()
 				.value_type().i32()
@@ -669,4 +805,225 @@ mod tests {
 
 	}
 
+	#[test]
+	fn branch_innermost() {
+		let module = builder::module()
+			.global()
+				.value_type().i32()
+				.build()
+			.function()
+				.signature().param().i32().build()
+				.body()
+					.with_instructions(elements::Instructions::new(
+						vec![
+							GetGlobal(0),
+							Block(elements::BlockType::NoResult),
+								GetGlobal(0),
+								Drop,
+								Br(0),
+								GetGlobal(0),
+								Drop,
+							End,
+							GetGlobal(0),
+							End
+						]
+					))
+					.build()
+				.build()
+			.build();
+
+		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
+
+		assert_eq!(
+			get_function_body(&injected_module, 0).unwrap(),
+			&vec![
+				I32Const(6),
+				Call(0),
+				GetGlobal(0),
+				Block(elements::BlockType::NoResult),
+					GetGlobal(0),
+					Drop,
+					Br(0),
+					I32Const(2),
+					Call(0),
+					GetGlobal(0),
+					Drop,
+				End,
+				GetGlobal(0),
+				End
+			][..]
+		);
+	}
+
+	#[test]
+	fn branch_outer_block() {
+		let module = builder::module()
+			.global()
+				.value_type().i32()
+				.build()
+			.function()
+				.signature().param().i32().build()
+				.body()
+					.with_instructions(elements::Instructions::new(
+						vec![
+							GetGlobal(0),
+							Block(elements::BlockType::NoResult),
+								GetGlobal(0),
+								If(elements::BlockType::NoResult),
+									GetGlobal(0),
+									GetGlobal(0),
+									Drop,
+									BrIf(1),
+								End,
+								GetGlobal(0),
+								Drop,
+							End,
+							GetGlobal(0),
+							End,
+						]
+					))
+					.build()
+				.build()
+			.build();
+
+		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
+
+		assert_eq!(
+			get_function_body(&injected_module, 0).unwrap(),
+			&vec![
+				I32Const(5),
+				Call(0),
+				GetGlobal(0),
+				Block(elements::BlockType::NoResult),
+					GetGlobal(0),
+					If(elements::BlockType::NoResult),
+						I32Const(4),
+						Call(0),
+						GetGlobal(0),
+						GetGlobal(0),
+						Drop,
+						BrIf(1),
+					End,
+					I32Const(2),
+					Call(0),
+					GetGlobal(0),
+					Drop,
+				End,
+				GetGlobal(0),
+				End,
+			][..]
+		);
+	}
+
+	#[test]
+	fn branch_outer_loop() {
+		let module = builder::module()
+			.global()
+				.value_type().i32()
+				.build()
+			.function()
+				.signature().param().i32().build()
+				.body()
+					.with_instructions(elements::Instructions::new(
+						vec![
+							GetGlobal(0),
+							Loop(elements::BlockType::NoResult),
+								GetGlobal(0),
+								If(elements::BlockType::NoResult),
+									GetGlobal(0),
+									BrIf(0),
+								Else,
+									GetGlobal(0),
+									GetGlobal(0),
+									Drop,
+									BrIf(1),
+								End,
+								GetGlobal(0),
+								Drop,
+							End,
+							GetGlobal(0),
+							End,
+						]
+					))
+					.build()
+				.build()
+			.build();
+
+		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
+
+		assert_eq!(
+			get_function_body(&injected_module, 0).unwrap(),
+			&vec![
+				I32Const(3),
+				Call(0),
+				GetGlobal(0),
+				Loop(elements::BlockType::NoResult),
+					I32Const(4),
+					Call(0),
+					GetGlobal(0),
+					If(elements::BlockType::NoResult),
+						I32Const(2),
+						Call(0),
+						GetGlobal(0),
+						BrIf(0),
+					Else,
+						I32Const(4),
+						Call(0),
+						GetGlobal(0),
+						GetGlobal(0),
+						Drop,
+						BrIf(1),
+					End,
+					GetGlobal(0),
+					Drop,
+				End,
+				GetGlobal(0),
+				End,
+			][..]
+		);
+	}
+
+	#[test]
+	fn return_from_func() {
+		let module = builder::module()
+			.global()
+				.value_type().i32()
+				.build()
+			.function()
+				.signature().param().i32().build()
+				.body()
+					.with_instructions(elements::Instructions::new(
+						vec![
+							GetGlobal(0),
+							If(elements::BlockType::NoResult),
+								Return,
+							End,
+							GetGlobal(0),
+							End,
+						]
+					))
+					.build()
+				.build()
+			.build();
+
+		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
+
+		assert_eq!(
+			get_function_body(&injected_module, 0).unwrap(),
+			&vec![
+				I32Const(2),
+				Call(0),
+				GetGlobal(0),
+				If(elements::BlockType::NoResult),
+					I32Const(1),
+					Call(0),
+					Return,
+				End,
+				I32Const(1),
+				Call(0),
+				GetGlobal(0),
+				End,
+			][..]
+		);
+	}
 }

From 0cf7daa9e529621fe75f11fe7e0fd5fb5f30509b Mon Sep 17 00:00:00 2001
From: Jim Posen <jim.posen@gmail.com>
Date: Tue, 2 Jul 2019 15:04:06 +0200
Subject: [PATCH 4/9] Update high level inject_gas_counter function
 documentation.

---
 src/gas.rs | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/gas.rs b/src/gas.rs
index a26dd69..d57c4ee 100644
--- a/src/gas.rs
+++ b/src/gas.rs
@@ -385,17 +385,18 @@ fn insert_metering_calls(
 /// 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.
 ///
-/// The calls to charge gas are inserted at the beginning of every block of code. A block is
-/// defined by `block`, `if`, `else`, `loop`, and `end` boundaries. Blocks form a nested hierarchy
-/// where `block`, `if`, `else`, and `loop` begin a new nested block, and `end` and `else` mark the
-/// end of a block. The gas cost of a block is determined statically as 1 plus the gas cost of all
-/// instructions directly in that block. Each instruction is only counted in the most deeply
-/// nested block containing it (ie. a block's cost does not include the cost of instructions in any
-/// blocks nested within it). The cost of the `begin`, `if`, and `loop` instructions is counted
-/// towards the block containing them, not the nested block that they open. There is no gas cost
-/// added for `end`/`else`, as they are pseudo-instructions. The gas cost of each instruction is
-/// determined by a `rules::Set` parameter. At the beginning of each block, this procedure injects
-/// new instructions to call the "gas" function with the gas cost of the block as an argument.
+/// The body of each function is divided into metered blocks, and the calls to charge gas are
+/// inserted at the beginning of every such block of code. A metered block is defined so that,
+/// unless there is a trap, either all of the instructions are executed or none are. These are
+/// similar to basic blocks in a control flow graph, except that in some cases multiple basic
+/// blocks can be merged into a single metered block. This is the case if any path through the
+/// control flow graph containing one basic block also contains another.
+///
+/// Charging gas is at the beginning of each metered block ensures that 1) all instructions
+/// executed are already paid for, 2) instructions that will not be executed are not charged for
+/// unless execution traps, and 3) the number of calls to "gas" is minimized. The corollary is that
+/// modules instrumented with this metering code may charge gas for instructions not executed in
+/// the event of a trap.
 ///
 /// Additionally, each `memory.grow` instruction found in the module is instrumented to first make
 /// a call to charge gas for the additional pages requested. This cannot be done as part of the
@@ -406,6 +407,8 @@ fn insert_metering_calls(
 /// function also rewrites all function indices references by code, table elements, etc., since
 /// the addition of an imported functions changes the indices of module-defined functions.
 ///
+/// This routine runs in time linear in the size of the input module.
+///
 /// The function fails if the module contains any operation forbidden by gas rule set, returning
 /// the original module as an Err.
 pub fn inject_gas_counter(module: elements::Module, rules: &rules::Set)

From b3f8f62105d8964c5a85de79f148ce5189b604fa Mon Sep 17 00:00:00 2001
From: Jim Posen <jim.posen@gmail.com>
Date: Tue, 2 Jul 2019 15:56:52 +0200
Subject: [PATCH 5/9] Update gas expectations.

---
 tests/expectations/gas/branch.wat |  6 +++---
 tests/expectations/gas/call.wat   |  4 ++--
 tests/expectations/gas/ifs.wat    |  6 +++---
 tests/expectations/gas/simple.wat | 10 ++++------
 tests/expectations/gas/start.wat  |  6 ++----
 5 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/tests/expectations/gas/branch.wat b/tests/expectations/gas/branch.wat
index 56c5ecf..4adafa9 100644
--- a/tests/expectations/gas/branch.wat
+++ b/tests/expectations/gas/branch.wat
@@ -4,11 +4,9 @@
   (import "env" "gas" (func (;0;) (type 1)))
   (func (;1;) (type 0) (result i32)
     (local i32 i32)
-    i32.const 3
+    i32.const 13
     call 0
     block  ;; label = @1
-      i32.const 17
-      call 0
       i32.const 0
       set_local 0
       i32.const 1
@@ -20,6 +18,8 @@
       set_local 1
       i32.const 1
       br_if 0 (;@1;)
+      i32.const 5
+      call 0
       get_local 0
       get_local 1
       tee_local 0
diff --git a/tests/expectations/gas/call.wat b/tests/expectations/gas/call.wat
index 878022a..fb32a0b 100644
--- a/tests/expectations/gas/call.wat
+++ b/tests/expectations/gas/call.wat
@@ -4,7 +4,7 @@
   (import "env" "gas" (func (;0;) (type 1)))
   (func (;1;) (type 0) (param i32 i32) (result i32)
     (local i32)
-    i32.const 6
+    i32.const 5
     call 0
     get_local 0
     get_local 1
@@ -12,7 +12,7 @@
     set_local 2
     get_local 2)
   (func (;2;) (type 0) (param i32 i32) (result i32)
-    i32.const 4
+    i32.const 3
     call 0
     get_local 0
     get_local 1
diff --git a/tests/expectations/gas/ifs.wat b/tests/expectations/gas/ifs.wat
index 89c72e7..7575504 100644
--- a/tests/expectations/gas/ifs.wat
+++ b/tests/expectations/gas/ifs.wat
@@ -3,17 +3,17 @@
   (type (;1;) (func (param i32)))
   (import "env" "gas" (func (;0;) (type 1)))
   (func (;1;) (type 0) (param i32) (result i32)
-    i32.const 3
+    i32.const 2
     call 0
     i32.const 1
     if (result i32)  ;; label = @1
-      i32.const 4
+      i32.const 3
       call 0
       get_local 0
       i32.const 1
       i32.add
     else
-      i32.const 3
+      i32.const 2
       call 0
       get_local 0
       i32.popcnt
diff --git a/tests/expectations/gas/simple.wat b/tests/expectations/gas/simple.wat
index d23ec7c..b967d21 100644
--- a/tests/expectations/gas/simple.wat
+++ b/tests/expectations/gas/simple.wat
@@ -3,24 +3,22 @@
   (type (;1;) (func (param i32)))
   (import "env" "gas" (func (;0;) (type 1)))
   (func (;1;) (type 0)
-    i32.const 3
+    i32.const 2
     call 0
     i32.const 1
     if  ;; label = @1
-      i32.const 2
+      i32.const 1
       call 0
       loop  ;; label = @2
-        i32.const 3
+        i32.const 2
         call 0
         i32.const 123
         drop
       end
     end)
   (func (;2;) (type 0)
-    i32.const 2
+    i32.const 1
     call 0
     block  ;; label = @1
-      i32.const 1
-      call 0
     end)
   (export "simple" (func 1)))
diff --git a/tests/expectations/gas/start.wat b/tests/expectations/gas/start.wat
index 89f43b5..18a447d 100644
--- a/tests/expectations/gas/start.wat
+++ b/tests/expectations/gas/start.wat
@@ -6,15 +6,13 @@
   (import "env" "memory" (memory (;0;) 1 1))
   (import "env" "gas" (func (;1;) (type 2)))
   (func (;2;) (type 1)
-    i32.const 5
+    i32.const 4
     call 1
     i32.const 8
     i32.const 4
     call 0
     unreachable)
-  (func (;3;) (type 1)
-    i32.const 1
-    call 1)
+  (func (;3;) (type 1))
   (export "call" (func 3))
   (start 2)
   (data (i32.const 8) "\01\02\03\04"))

From b5472bcd8f23af65b820398d55976719c7b9cae5 Mon Sep 17 00:00:00 2001
From: Jim Posen <jim.posen@gmail.com>
Date: Thu, 11 Jul 2019 16:45:46 +0200
Subject: [PATCH 6/9] !fixup Address review comments.

---
 src/gas.rs | 40 +++++++++++++++++-----------------------
 1 file changed, 17 insertions(+), 23 deletions(-)

diff --git a/src/gas.rs b/src/gas.rs
index d57c4ee..e1f276c 100644
--- a/src/gas.rs
+++ b/src/gas.rs
@@ -128,8 +128,6 @@ impl Counter {
 
 		// If there may have been a branch to a lower index, then also finalize the active metered
 		// block for the previous control block. Otherwise, finalize it and begin a new one.
-		//
-		// TODO: handle diverging function calls.
 		let may_br_out = closing_control_block.lowest_forward_br < closing_control_index;
 		if may_br_out {
 			self.finalize_metered_block(cursor)?;
@@ -153,9 +151,9 @@ impl Counter {
 			)
 		};
 
-		// If the block was opened with a `begin`, then its start position will be set to that of
+		// If the block was opened with a `block`, then its start position will be set to that of
 		// the active metered block in the control block one higher on the stack. This is because
-		// any instructions between a `begin` and the first branch are part of the same basic block
+		// any instructions between a `block` and the first branch are part of the same basic block
 		// as the preceding instruction. In this case, instead of finalizing the block, merge its
 		// cost into the other active metered block to avoid injecting unnecessary instructions.
 		let last_index = self.stack.len() - 1;
@@ -183,7 +181,7 @@ impl Counter {
 		self.finalize_metered_block(cursor)?;
 
 		// Update the lowest_forward_br of the current control block.
-		for index in indices.iter().cloned() {
+		for &index in indices {
 			let target_is_loop = {
 				let target_block = self.stack.get(index).ok_or_else(|| ())?;
 				target_block.is_loop
@@ -199,20 +197,20 @@ impl Counter {
 		Ok(())
 	}
 
-	/// Returns the height of the control stack.
-	fn stack_height(&self) -> usize {
-		self.stack.len()
+	/// Returns the stack index of the active control block. Returns None if stack is empty.
+	fn active_control_block_index(&self) -> Option<usize> {
+		self.stack.len().checked_sub(1)
 	}
 
 	/// Get a reference to the currently active metered block.
-	fn active_block(&mut self) -> Result<&mut MeteredBlock, ()> {
+	fn active_metered_block(&mut self) -> Result<&mut MeteredBlock, ()> {
 		let top_block = self.stack.last_mut().ok_or_else(|| ())?;
 		Ok(&mut top_block.active_metered_block)
 	}
 
 	/// Increment the cost of the current block by the specified value.
 	fn increment(&mut self, val: u32) -> Result<(), ()> {
-		let top_block = self.active_block()?;
+		let top_block = self.active_metered_block()?;
 		top_block.cost = top_block.cost.checked_add(val).ok_or_else(|| ())?;
 		Ok(())
 	}
@@ -278,9 +276,7 @@ pub fn inject_counter(
 				// be included into this block. The start position is set to that of the previous
 				// active metered block to signal that they should be merged in order to reduce
 				// unnecessary metering instructions.
-				let top_block_start_pos = counter.active_block()
-					.expect("stack is not empty; qed")
-					.start_pos;
+				let top_block_start_pos = counter.active_metered_block()?.start_pos;
 				counter.begin_control_block(top_block_start_pos, false);
 			}
 			If(_) => {
@@ -301,19 +297,19 @@ pub fn inject_counter(
 				counter.increment(instruction_cost)?;
 
 				// Label is a relative index into the control stack.
-				let index = counter.stack_height().checked_sub(1 + label as usize)
-					.ok_or_else(|| ())?;
-				counter.branch(cursor, &[index])?;
+				let active_index = counter.active_control_block_index().ok_or_else(|| ())?;
+				let target_index = active_index.checked_sub(label as usize).ok_or_else(|| ())?;
+				counter.branch(cursor, &[target_index])?;
 			}
 			BrTable(ref label_vec, label_default) => {
 				counter.increment(instruction_cost)?;
 
-				let stack_height = counter.stack_height();
-				let indices = [label_default].iter().chain(label_vec.iter())
-					.map(|label| stack_height.checked_sub(1 + *label as usize))
+				let active_index = counter.active_control_block_index().ok_or_else(|| ())?;
+				let target_indices = [label_default].iter().chain(label_vec.iter())
+					.map(|label| active_index.checked_sub(*label as usize))
 					.collect::<Option<Vec<_>>>()
 					.ok_or_else(|| ())?;
-				counter.branch(cursor, &indices)?;
+				counter.branch(cursor, &target_indices)?;
 			}
 			Return => {
 				counter.increment(instruction_cost)?;
@@ -350,8 +346,7 @@ fn insert_metering_calls(
 	blocks.sort_unstable_by_key(|block| block.start_pos);
 	let mut block_iter = blocks.into_iter().peekable();
 
-	let mut original_pos = 0;
-	for instr in original_instrs.into_iter() {
+	for (original_pos, instr) in original_instrs.into_iter().enumerate() {
 		// If there the next block starts at this position, inject metering instructions.
 		let used_block = if let Some(ref block) = block_iter.peek() {
 			if block.start_pos == original_pos {
@@ -367,7 +362,6 @@ fn insert_metering_calls(
 
 		// Copy over the original instruction.
 		new_instrs.push(instr);
-		original_pos += 1;
 	}
 
 	if block_iter.next().is_some() {

From ed7f31ec20ab443390c8f11557e6d44ce5c41993 Mon Sep 17 00:00:00 2001
From: Jim Posen <jim.posen@gmail.com>
Date: Fri, 12 Jul 2019 10:15:37 +0200
Subject: [PATCH 7/9] Use macro + WAT to make gas injection unit tests more
 readable.

---
 src/gas.rs | 573 ++++++++++++++++++++++-------------------------------
 1 file changed, 232 insertions(+), 341 deletions(-)

diff --git a/src/gas.rs b/src/gas.rs
index e1f276c..d068443 100644
--- a/src/gas.rs
+++ b/src/gas.rs
@@ -587,136 +587,6 @@ mod tests {
 		self::wabt::wasm2wat(&binary).unwrap();
 	}
 
-	#[test]
-	fn simple() {
-		let module = builder::module()
-			.global()
-				.value_type().i32()
-				.build()
-			.function()
-				.signature().param().i32().build()
-				.body()
-					.with_instructions(elements::Instructions::new(
-						vec![
-							GetGlobal(0),
-							End
-						]
-					))
-					.build()
-				.build()
-			.build();
-
-		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
-
-		assert_eq!(
-			get_function_body(&injected_module, 0).unwrap(),
-			&vec![
-				I32Const(1),
-				Call(0),
-				GetGlobal(0),
-				End
-			][..]
-		);
-	}
-
-	#[test]
-	fn nested() {
-		let module = builder::module()
-			.global()
-				.value_type().i32()
-				.build()
-			.function()
-				.signature().param().i32().build()
-				.body()
-					.with_instructions(elements::Instructions::new(
-						vec![
-							GetGlobal(0),
-							Block(elements::BlockType::NoResult),
-								GetGlobal(0),
-								GetGlobal(0),
-								GetGlobal(0),
-							End,
-							GetGlobal(0),
-							End
-						]
-					))
-					.build()
-				.build()
-			.build();
-
-		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
-
-		assert_eq!(
-			get_function_body(&injected_module, 0).unwrap(),
-			&vec![
-				I32Const(6),
-				Call(0),
-				GetGlobal(0),
-				Block(elements::BlockType::NoResult),
-					GetGlobal(0),
-					GetGlobal(0),
-					GetGlobal(0),
-				End,
-				GetGlobal(0),
-				End
-			][..]
-		);
-	}
-
-	#[test]
-	fn ifelse() {
-		let module = builder::module()
-			.global()
-				.value_type().i32()
-				.build()
-			.function()
-				.signature().param().i32().build()
-				.body()
-					.with_instructions(elements::Instructions::new(
-						vec![
-							GetGlobal(0),
-							If(elements::BlockType::NoResult),
-								GetGlobal(0),
-								GetGlobal(0),
-								GetGlobal(0),
-							Else,
-								GetGlobal(0),
-								GetGlobal(0),
-							End,
-							GetGlobal(0),
-							End
-						]
-					))
-					.build()
-				.build()
-			.build();
-
-		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
-
-		assert_eq!(
-			get_function_body(&injected_module, 0).unwrap(),
-			&vec![
-				I32Const(3),
-				Call(0),
-				GetGlobal(0),
-				If(elements::BlockType::NoResult),
-					I32Const(3),
-					Call(0),
-					GetGlobal(0),
-					GetGlobal(0),
-					GetGlobal(0),
-				Else,
-					I32Const(2),
-					Call(0),
-					GetGlobal(0),
-					GetGlobal(0),
-				End,
-				GetGlobal(0),
-				End
-			][..]
-		);
-	}
-
 	#[test]
 	fn call_index() {
 		let module = builder::module()
@@ -775,7 +645,6 @@ mod tests {
 		);
 	}
 
-
 	#[test]
 	fn forbidden() {
 		let module = builder::module()
@@ -802,225 +671,247 @@ mod tests {
 
 	}
 
-	#[test]
-	fn branch_innermost() {
-		let module = builder::module()
-			.global()
-				.value_type().i32()
-				.build()
-			.function()
-				.signature().param().i32().build()
-				.body()
-					.with_instructions(elements::Instructions::new(
-						vec![
-							GetGlobal(0),
-							Block(elements::BlockType::NoResult),
-								GetGlobal(0),
-								Drop,
-								Br(0),
-								GetGlobal(0),
-								Drop,
-							End,
-							GetGlobal(0),
-							End
-						]
-					))
-					.build()
-				.build()
-			.build();
-
-		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
-
-		assert_eq!(
-			get_function_body(&injected_module, 0).unwrap(),
-			&vec![
-				I32Const(6),
-				Call(0),
-				GetGlobal(0),
-				Block(elements::BlockType::NoResult),
-					GetGlobal(0),
-					Drop,
-					Br(0),
-					I32Const(2),
-					Call(0),
-					GetGlobal(0),
-					Drop,
-				End,
-				GetGlobal(0),
-				End
-			][..]
-		);
+	fn parse_wat(source: &str) -> elements::Module {
+		let module_bytes = wabt::Wat2Wasm::new()
+			.validate(false)
+			.convert(source)
+			.expect("failed to parse module");
+		elements::deserialize_buffer(module_bytes.as_ref())
+			.expect("failed to parse module")
 	}
 
-	#[test]
-	fn branch_outer_block() {
-		let module = builder::module()
-			.global()
-				.value_type().i32()
-				.build()
-			.function()
-				.signature().param().i32().build()
-				.body()
-					.with_instructions(elements::Instructions::new(
-						vec![
-							GetGlobal(0),
-							Block(elements::BlockType::NoResult),
-								GetGlobal(0),
-								If(elements::BlockType::NoResult),
-									GetGlobal(0),
-									GetGlobal(0),
-									Drop,
-									BrIf(1),
-								End,
-								GetGlobal(0),
-								Drop,
-							End,
-							GetGlobal(0),
-							End,
-						]
-					))
-					.build()
-				.build()
-			.build();
+	macro_rules! test_gas_counter_injection {
+		(name = $name:ident; input = $input:expr; expected = $expected:expr) => {
+			#[test]
+			fn $name() {
+				let input_module = parse_wat($input);
+				let expected_module = parse_wat($expected);
 
-		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
+				let injected_module = inject_gas_counter(input_module, &Default::default())
+					.expect("inject_gas_counter call failed");
 
-		assert_eq!(
-			get_function_body(&injected_module, 0).unwrap(),
-			&vec![
-				I32Const(5),
-				Call(0),
-				GetGlobal(0),
-				Block(elements::BlockType::NoResult),
-					GetGlobal(0),
-					If(elements::BlockType::NoResult),
-						I32Const(4),
-						Call(0),
-						GetGlobal(0),
-						GetGlobal(0),
-						Drop,
-						BrIf(1),
-					End,
-					I32Const(2),
-					Call(0),
-					GetGlobal(0),
-					Drop,
-				End,
-				GetGlobal(0),
-				End,
-			][..]
-		);
+				let actual_func_body = get_function_body(&injected_module, 0)
+					.expect("injected module must have a function body");
+				let expected_func_body = get_function_body(&expected_module, 0)
+					.expect("post-module must have a function body");
+
+				assert_eq!(actual_func_body, expected_func_body);
+			}
+		}
 	}
 
-	#[test]
-	fn branch_outer_loop() {
-		let module = builder::module()
-			.global()
-				.value_type().i32()
-				.build()
-			.function()
-				.signature().param().i32().build()
-				.body()
-					.with_instructions(elements::Instructions::new(
-						vec![
-							GetGlobal(0),
-							Loop(elements::BlockType::NoResult),
-								GetGlobal(0),
-								If(elements::BlockType::NoResult),
-									GetGlobal(0),
-									BrIf(0),
-								Else,
-									GetGlobal(0),
-									GetGlobal(0),
-									Drop,
-									BrIf(1),
-								End,
-								GetGlobal(0),
-								Drop,
-							End,
-							GetGlobal(0),
-							End,
-						]
-					))
-					.build()
-				.build()
-			.build();
-
-		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
-
-		assert_eq!(
-			get_function_body(&injected_module, 0).unwrap(),
-			&vec![
-				I32Const(3),
-				Call(0),
-				GetGlobal(0),
-				Loop(elements::BlockType::NoResult),
-					I32Const(4),
-					Call(0),
-					GetGlobal(0),
-					If(elements::BlockType::NoResult),
-						I32Const(2),
-						Call(0),
-						GetGlobal(0),
-						BrIf(0),
-					Else,
-						I32Const(4),
-						Call(0),
-						GetGlobal(0),
-						GetGlobal(0),
-						Drop,
-						BrIf(1),
-					End,
-					GetGlobal(0),
-					Drop,
-				End,
-				GetGlobal(0),
-				End,
-			][..]
-		);
+	test_gas_counter_injection! {
+		name = simple;
+		input = r#"
+		(module
+			(func (result i32)
+				(get_global 0)))
+		"#;
+		expected = r#"
+		(module
+			(func (result i32)
+				(call 0 (i32.const 1))
+				(get_global 0)))
+		"#
 	}
 
-	#[test]
-	fn return_from_func() {
-		let module = builder::module()
-			.global()
-				.value_type().i32()
-				.build()
-			.function()
-				.signature().param().i32().build()
-				.body()
-					.with_instructions(elements::Instructions::new(
-						vec![
-							GetGlobal(0),
-							If(elements::BlockType::NoResult),
-								Return,
-							End,
-							GetGlobal(0),
-							End,
-						]
-					))
-					.build()
-				.build()
-			.build();
+	test_gas_counter_injection! {
+		name = nested;
+		input = r#"
+		(module
+			(func (result i32)
+				(get_global 0)
+				(block
+					(get_global 0)
+					(get_global 0)
+					(get_global 0))
+				(get_global 0)))
+		"#;
+		expected = r#"
+		(module
+			(func (result i32)
+				(call 0 (i32.const 6))
+				(get_global 0)
+				(block
+					(get_global 0)
+					(get_global 0)
+					(get_global 0))
+				(get_global 0)))
+		"#
+	}
 
-		let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
+	test_gas_counter_injection! {
+		name = ifelse;
+		input = r#"
+		(module
+			(func (result i32)
+				(get_global 0)
+				(if
+					(then
+						(get_global 0)
+						(get_global 0)
+						(get_global 0))
+					(else
+						(get_global 0)
+						(get_global 0)))
+				(get_global 0)))
+		"#;
+		expected = r#"
+		(module
+			(func (result i32)
+				(call 0 (i32.const 3))
+				(get_global 0)
+				(if
+					(then
+						(call 0 (i32.const 3))
+						(get_global 0)
+						(get_global 0)
+						(get_global 0))
+					(else
+						(call 0 (i32.const 2))
+						(get_global 0)
+						(get_global 0)))
+				(get_global 0)))
+		"#
+	}
 
-		assert_eq!(
-			get_function_body(&injected_module, 0).unwrap(),
-			&vec![
-				I32Const(2),
-				Call(0),
-				GetGlobal(0),
-				If(elements::BlockType::NoResult),
-					I32Const(1),
-					Call(0),
-					Return,
-				End,
-				I32Const(1),
-				Call(0),
-				GetGlobal(0),
-				End,
-			][..]
-		);
+	test_gas_counter_injection! {
+		name = branch_innermost;
+		input = r#"
+		(module
+			(func (result i32)
+				(get_global 0)
+				(block
+					(get_global 0)
+					(drop)
+					(br 0)
+					(get_global 0)
+					(drop))
+				(get_global 0)))
+		"#;
+		expected = r#"
+		(module
+			(func (result i32)
+				(call 0 (i32.const 6))
+				(get_global 0)
+				(block
+					(get_global 0)
+					(drop)
+					(br 0)
+					(call 0 (i32.const 2))
+					(get_global 0)
+					(drop))
+				(get_global 0)))
+		"#
+	}
+
+	test_gas_counter_injection! {
+		name = branch_outer_block;
+		input = r#"
+		(module
+			(func (result i32)
+				(get_global 0)
+				(block
+					(get_global 0)
+					(if
+						(then
+							(get_global 0)
+							(get_global 0)
+							(drop)
+							(br_if 1)))
+					(get_global 0)
+					(drop))
+				(get_global 0)))
+		"#;
+		expected = r#"
+		(module
+			(func (result i32)
+				(call 0 (i32.const 5))
+				(get_global 0)
+				(block
+					(get_global 0)
+					(if
+						(then
+							(call 0 (i32.const 4))
+							(get_global 0)
+							(get_global 0)
+							(drop)
+							(br_if 1)))
+					(call 0 (i32.const 2))
+					(get_global 0)
+					(drop))
+				(get_global 0)))
+		"#
+	}
+
+	test_gas_counter_injection! {
+		name = branch_outer_loop;
+		input = r#"
+		(module
+			(func (result i32)
+				(get_global 0)
+				(loop
+					(get_global 0)
+					(if
+						(then
+							(get_global 0)
+							(br_if 0))
+						(else
+							(get_global 0)
+							(get_global 0)
+							(drop)
+							(br_if 1)))
+					(get_global 0)
+					(drop))
+				(get_global 0)))
+		"#;
+		expected = r#"
+		(module
+			(func (result i32)
+				(call 0 (i32.const 3))
+				(get_global 0)
+				(loop
+					(call 0 (i32.const 4))
+					(get_global 0)
+					(if
+						(then
+							(call 0 (i32.const 2))
+							(get_global 0)
+							(br_if 0))
+						(else
+							(call 0 (i32.const 4))
+							(get_global 0)
+							(get_global 0)
+							(drop)
+							(br_if 1)))
+					(get_global 0)
+					(drop))
+				(get_global 0)))
+		"#
+	}
+
+	test_gas_counter_injection! {
+		name = return_from_func;
+		input = r#"
+		(module
+			(func (result i32)
+				(get_global 0)
+				(if
+					(then
+						(return)))
+				(get_global 0)))
+		"#;
+		expected = r#"
+		(module
+			(func (result i32)
+				(call 0 (i32.const 2))
+				(get_global 0)
+				(if
+					(then
+						(call 0 (i32.const 1))
+						(return)))
+				(call 0 (i32.const 1))
+				(get_global 0)))
+		"#
 	}
 }

From 93abbcfe564824cdb08bd44cb87bd56f93c1153e Mon Sep 17 00:00:00 2001
From: Jim Posen <jim.posen@gmail.com>
Date: Fri, 12 Jul 2019 10:30:11 +0200
Subject: [PATCH 8/9] Add gas injection unit test for correct else behavior.

---
 src/gas.rs | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/src/gas.rs b/src/gas.rs
index d068443..27aba65 100644
--- a/src/gas.rs
+++ b/src/gas.rs
@@ -914,4 +914,40 @@ mod tests {
 				(get_global 0)))
 		"#
 	}
+
+	test_gas_counter_injection! {
+		name = branch_from_if_not_else;
+		input = r#"
+		(module
+			(func (result i32)
+				(get_global 0)
+				(block
+					(get_global 0)
+					(if
+						(then (br 1))
+						(else (br 0)))
+					(get_global 0)
+					(drop))
+				(get_global 0)))
+		"#;
+		expected = r#"
+		(module
+			(func (result i32)
+				(call 0 (i32.const 5))
+				(get_global 0)
+				(block
+					(get_global 0)
+					(if
+						(then
+							(call 0 (i32.const 1))
+							(br 1))
+						(else
+							(call 0 (i32.const 1))
+							(br 0)))
+					(call 0 (i32.const 2))
+					(get_global 0)
+					(drop))
+				(get_global 0)))
+		"#
+	}
 }

From 82bd972333eba943eb50bdfb80751bbf85d70d3c Mon Sep 17 00:00:00 2001
From: Jim Posen <jim.posen@gmail.com>
Date: Fri, 12 Jul 2019 10:39:57 +0200
Subject: [PATCH 9/9] !fixup lowest_forward_br rename and field docs.

---
 src/gas.rs | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/gas.rs b/src/gas.rs
index 27aba65..2efd82b 100644
--- a/src/gas.rs
+++ b/src/gas.rs
@@ -42,12 +42,13 @@ pub fn update_call_index(instructions: &mut elements::Instructions, inserted_ind
 struct ControlBlock {
 	/// The lowest control stack index corresponding to a forward jump targeted by a br, br_if, or
 	/// br_table instruction within this control block. The index must refer to a control block
-	/// that is not a loop, meaning it is a forward jump.
+	/// that is not a loop, meaning it is a forward jump. Given the way Wasm control flow is
+	/// structured, the lowest index on the stack represents the furthest forward branch target.
 	///
 	/// This value will always be at most the index of the block itself, even if there is no
 	/// explicit br instruction targeting this control block. This does not affect how the value is
 	/// used in the metering algorithm.
-	lowest_forward_br: usize,
+	lowest_forward_br_target: usize,
 
 	/// The active metering block that new instructions contribute a gas cost towards.
 	active_metered_block: MeteredBlock,
@@ -94,7 +95,7 @@ impl Counter {
 	fn begin_control_block(&mut self, cursor: usize, is_loop: bool) {
 		let index = self.stack.len();
 		self.stack.push(ControlBlock {
-			lowest_forward_br: index,
+			lowest_forward_br_target: index,
 			active_metered_block: MeteredBlock {
 				start_pos: cursor,
 				cost: 0,
@@ -118,17 +119,18 @@ impl Counter {
 			return Ok(())
 		}
 
-		// Update the lowest_forward_br for the control block now on top of the stack.
+		// Update the lowest_forward_br_target for the control block now on top of the stack.
 		{
 			let control_block = self.stack.last_mut().ok_or_else(|| ())?;
-			control_block.lowest_forward_br = min(
-				control_block.lowest_forward_br, closing_control_block.lowest_forward_br
+			control_block.lowest_forward_br_target = min(
+				control_block.lowest_forward_br_target,
+				closing_control_block.lowest_forward_br_target
 			);
 		}
 
 		// If there may have been a branch to a lower index, then also finalize the active metered
 		// block for the previous control block. Otherwise, finalize it and begin a new one.
-		let may_br_out = closing_control_block.lowest_forward_br < closing_control_index;
+		let may_br_out = closing_control_block.lowest_forward_br_target < closing_control_index;
 		if may_br_out {
 			self.finalize_metered_block(cursor)?;
 		}
@@ -180,7 +182,7 @@ impl Counter {
 	fn branch(&mut self, cursor: usize, indices: &[usize]) -> Result<(), ()> {
 		self.finalize_metered_block(cursor)?;
 
-		// Update the lowest_forward_br of the current control block.
+		// Update the lowest_forward_br_target of the current control block.
 		for &index in indices {
 			let target_is_loop = {
 				let target_block = self.stack.get(index).ok_or_else(|| ())?;
@@ -191,7 +193,8 @@ impl Counter {
 			}
 
 			let control_block = self.stack.last_mut().ok_or_else(|| ())?;
-			control_block.lowest_forward_br = min(control_block.lowest_forward_br, index);
+			control_block.lowest_forward_br_target =
+				min(control_block.lowest_forward_br_target, index);
 		}
 
 		Ok(())