Superblock metering

wasm-utils/cli/gas/main.rs

@@ -15,9 +15,12 @@ fn main() {
 
 	let memory_page_cost = 256 * 1024; // 256k gas for 1 page (64k) of memory
 
+	// let config = pwasm_utils::rules::Set::default()
+	//	.with_forbidden_floats() // Reject floating point opreations.
+	//	.with_grow_cost(memory_page_cost);
+
 	let config = pwasm_utils::rules::Set::default()
-		.with_forbidden_floats() // Reject floating point opreations.
-		.with_grow_cost(memory_page_cost);
+		.with_forbidden_floats();
 
 	// Loading module
 	let module = parity_wasm::deserialize_file(&args[1]).expect("Module deserialization to succeed");

wasm-utils/src/gas.rs

@@ -36,6 +36,7 @@ struct BlockEntry {
 	start_pos: usize,
 	/// Sum of costs of all instructions until end of the block.
 	cost: u32,
+	flow_up: bool,
 }
 
 struct Counter {
@@ -55,11 +56,12 @@ impl Counter {
 	}
 
 	/// Begin a new block.
-	fn begin(&mut self, cursor: usize) {
+	fn begin(&mut self, cursor: usize, cost: u32, flow_up: bool) {
 		let block_idx = self.blocks.len();
 		self.blocks.push(BlockEntry {
 			start_pos: cursor,
-			cost: 1,
+			cost: cost,
+			flow_up: flow_up,
 		});
 		self.stack.push(block_idx);
 	}
@@ -72,15 +74,52 @@ impl Counter {
 		Ok(())
 	}
 
-	/// 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(|| ())?;
-
-		top_block.cost = top_block.cost.checked_add(val).ok_or_else(|| ())?;
-
+	fn increment_control_flow(&mut self, val: u32) -> Result<(), ()> {
+		/*
+		;; if the current block and parent block has 0 cost, then we're seeing a sequence of nested blocks
+
+		(block $B1
+		  (block $B2
+		    (block $B3
+		      ...
+		     )))
+
+		;; instead of calling to useGas once after each block, we can sum up the 1 gas per `block` instructions
+		;; and charge for all three at the top of block $B1
+
+		;; instead of this:
+		(block $B1
+		  (call useGas (i32.const 1))
+		  (block $B2
+		    (call useGas (i32.const 1))
+		    (block $B3
+		      (call useGas (i32.const 1))
+		      ...
+		      )))
+
+		;; do this:
+		(block $B1
+		  (call useGas (i32.const 3))
+		  (block $B2
+		    (block $B3
+		      ...
+		      )))
+		*/
+
+		// find closest ancestor block (starting from top of stack and going down) with blocked flow and add 1
+
+		for (_i, stack_i) in self.stack.iter().rev().enumerate() {
+			let block_i = self.blocks.get_mut(*stack_i).ok_or_else(|| ())?;
+			if !block_i.flow_up || *stack_i == 0 {
+				block_i.cost = block_i.cost.checked_add(val).ok_or_else(|| ())?;
+				//println!("found ancestor with blocked flow or no parent. incrementing to new cost: {:?} and returning...", block_i.cost);
+				break;
+			}
+		}
 		Ok(())
+
 	}
+
 }
 
 fn inject_grow_counter(instructions: &mut elements::Instructions, grow_counter_func: u32) -> usize {
@@ -130,23 +169,78 @@ pub fn inject_counter(
 	let mut counter = Counter::new();
 
 	// Begin an implicit function (i.e. `func...end`) block.
-	counter.begin(0);
+	counter.begin(0, 0, false);
 
 	for cursor in 0..instructions.elements().len() {
 		let instruction = &instructions.elements()[cursor];
 		match *instruction {
-			Block(_) | If(_) | Loop(_) => {
+			Block(_) => {
 				// Increment previous block with the cost of the current opcode.
 				let instruction_cost = rules.process(instruction)?;
-				counter.increment(instruction_cost)?;
+				//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);
-			}
+
+				// add cost, which may flow up to ancestor block
+				counter.increment_control_flow(instruction_cost)?;
+
+				counter.begin(cursor + 1, 0, true);
+
+			},
+			If(_) => {
+				// Increment previous block with the cost of the current opcode.
+				let instruction_cost = rules.process(instruction)?;
+				//counter.increment(instruction_cost)?;
+				counter.increment_control_flow(instruction_cost)?;
+
+				// begin If with cost 1, to force new costs added to top of block
+				counter.begin(cursor + 1, 0, false);
+			},
+			BrIf(_) => {
+				// Increment previous block with the cost of the current opcode.
+				let instruction_cost = rules.process(instruction)?;
+				//counter.increment(instruction_cost)?;
+				counter.increment_control_flow(instruction_cost)?;
+
+				// on a br_if, we finalize the previous block because those instructions will always be executed.
+				// intructions after the if will be executed conditionally, so we start a new block so that gasUsed can be called after the if.
+				counter.finalize()?;
+
+				// begin If with cost 1, to force new costs added to top of block
+				counter.begin(cursor + 1, 0, false);
+			},
+			Loop(_) => {
+				let instruction_cost = rules.process(instruction)?;
+				//counter.increment(instruction_cost)?;
+				//counter.increment_control_flow(instruction_cost)?;
+
+				counter.begin(cursor + 1, 0, false);
+				// charge for the loop after the loop instruction
+				// need to do this because the loop could be executed many times (the br_if that jumps to the loop is a separate instruction and gas charge)
+				counter.increment_control_flow(instruction_cost)?;
+			},
+			Br(_) => {
+				// anything after a break is dead code.
+				// for now, we treat dead code blocks like any other (the metering will not be executed)
+				// TODO: handle properly and don't inject metering inside dead code blocks
+				let instruction_cost = rules.process(instruction)?;
+				counter.increment_control_flow(instruction_cost)?;
+				counter.finalize()?;
+				counter.begin(cursor + 1, 0, false);
+			},
+			// br_table is always followed by end (in the ecmul wasm code, at least)
+			// BrTable(_,_) => { },
+			// return is always followed by end (in the ecmul wasm code, at least)
+			// Return => { },
 			End => {
 				// Just finalize current block.
+				//counter.increment_control_flow(instruction_cost)?;
+				// wasabi doesn't count end as an instruction, so neither will we (no gas charge)
+
+
 				counter.finalize()?;
+				counter.begin(cursor + 1, 0, false);
 			},
 			Else => {
 				// `Else` opcode is being encountered. So the case we are looking at:
@@ -160,26 +254,32 @@ pub fn inject_counter(
 				// 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.begin(cursor + 1, 1, false);
+			},
+			Unreachable => {
+				// charge nothing, do nothing
+			},
 			_ => {
 				// An ordinal non control flow instruction. Just increment the cost of the current block.
 				let instruction_cost = rules.process(instruction)?;
-				counter.increment(instruction_cost)?;
+				//counter.increment(instruction_cost)?;
+				counter.increment_control_flow(instruction_cost)?;
 			}
 		}
 	}
 
 	// Then insert metering calls.
 	let mut cumulative_offset = 0;
 	for block in counter.blocks {
-		let effective_pos = block.start_pos + cumulative_offset;
+		if block.cost > 0 {
+			let effective_pos = block.start_pos + cumulative_offset;
 
-		instructions.elements_mut().insert(effective_pos, I64Const(block.cost as i64));
-		instructions.elements_mut().insert(effective_pos+1, Call(gas_func));
+			instructions.elements_mut().insert(effective_pos, I64Const(block.cost as i64));
+			instructions.elements_mut().insert(effective_pos+1, Call(gas_func));
 
-		// Take into account these two inserted instructions.
-		cumulative_offset += 2;
+			// Take into account these two inserted instructions.
+			cumulative_offset += 2;
+		}
 	}
 
 	Ok(())