Charge gas for local variables on the callee side (#38)

Co-authored-by: Alexander Theißen <alex.theissen@me.com>

CHANGELOG.md

@@ -22,6 +22,8 @@ In other words: Upgrading this pallet will not break pre-existing contracts.
 
 - Add new gas metering method: mutable global + local gas function
 [#34](https://github.com/paritytech/wasm-instrument/pull/34)
+- Account for locals initialization costs
+[#38](https://github.com/paritytech/wasm-instrument/pull/38)
 
 ## [v0.3.0]
 

src/gas_metering/mod.rs

@@ -36,6 +36,9 @@ pub trait Rules {
 	/// code into the function calling `memory.grow`. Therefore returning anything but
 	/// [`MemoryGrowCost::Free`] introduces some overhead to the `memory.grow` instruction.
 	fn memory_grow_cost(&self) -> MemoryGrowCost;
+
+	/// A surcharge cost to calling a function that is added per local of that function.
+	fn call_per_local_cost(&self) -> u32;
 }
 
 /// Dynamic costs for memory growth.
@@ -76,6 +79,7 @@ impl MemoryGrowCost {
 pub struct ConstantCostRules {
 	instruction_cost: u32,
 	memory_grow_cost: u32,
+	call_per_local_cost: u32,
 }
 
 impl ConstantCostRules {
@@ -83,15 +87,15 @@ impl ConstantCostRules {
 	///
 	/// Uses `instruction_cost` for every instruction and `memory_grow_cost` to dynamically
 	/// meter the memory growth instruction.
-	pub fn new(instruction_cost: u32, memory_grow_cost: u32) -> Self {
+	pub fn new(instruction_cost: u32, memory_grow_cost: u32, call_per_local_cost: u32) -> Self {
-		Self { instruction_cost, memory_grow_cost }
+		Self { instruction_cost, memory_grow_cost, call_per_local_cost }
 	}
 }
 
 impl Default for ConstantCostRules {
 	/// Uses instruction cost of `1` and disables memory growth instrumentation.
 	fn default() -> Self {
-		Self { instruction_cost: 1, memory_grow_cost: 0 }
+		Self { instruction_cost: 1, memory_grow_cost: 0, call_per_local_cost: 1 }
 	}
 }
 
@@ -103,6 +107,10 @@ impl Rules for ConstantCostRules {
 	fn memory_grow_cost(&self) -> MemoryGrowCost {
 		NonZeroU32::new(self.memory_grow_cost).map_or(MemoryGrowCost::Free, MemoryGrowCost::Linear)
 	}
+
+	fn call_per_local_cost(&self) -> u32 {
+		self.call_per_local_cost
+	}
 }
 
 /// Transforms a given module into one that tracks the gas charged during its execution.
@@ -249,8 +257,16 @@ pub fn inject<R: Rules, B: Backend>(
 							}
 						}
 					}
-					if inject_counter(func_body.code_mut(), gas_fn_cost, rules, gas_func_idx)
+					let locals_count =
-						.is_err()
+						func_body.locals().iter().map(|val_type| val_type.count()).sum();
+					if inject_counter(
+						func_body.code_mut(),
+						gas_fn_cost,
+						locals_count,
+						rules,
+						gas_func_idx,
+					)
+					.is_err()
 					{
 						error = true;
 						break
@@ -566,6 +582,7 @@ fn add_grow_counter<R: Rules>(
 fn determine_metered_blocks<R: Rules>(
 	instructions: &elements::Instructions,
 	rules: &R,
+	locals_count: u32,
 ) -> Result<Vec<MeteredBlock>, ()> {
 	use parity_wasm::elements::Instruction::*;
 
@@ -573,6 +590,9 @@ fn determine_metered_blocks<R: Rules>(
 
 	// Begin an implicit function (i.e. `func...end`) block.
 	counter.begin_control_block(0, false);
+	// Add locals initialization cost to the function block.
+	let locals_init_cost = rules.call_per_local_cost().checked_mul(locals_count).ok_or(())?;
+	counter.increment(locals_init_cost)?;
 
 	for cursor in 0..instructions.elements().len() {
 		let instruction = &instructions.elements()[cursor];
@@ -640,10 +660,11 @@ fn determine_metered_blocks<R: Rules>(
 fn inject_counter<R: Rules>(
 	instructions: &mut elements::Instructions,
 	gas_function_cost: u64,
+	locals_count: u32,
 	rules: &R,
 	gas_func: u32,
 ) -> Result<(), ()> {
-	let blocks = determine_metered_blocks(instructions, rules)?;
+	let blocks = determine_metered_blocks(instructions, rules, locals_count)?;
 	insert_metering_calls(instructions, gas_function_cost, blocks, gas_func)
 }
 
@@ -668,7 +689,8 @@ fn insert_metering_calls(
 		// If there the next block starts at this position, inject metering instructions.
 		let used_block = if let Some(block) = block_iter.peek() {
 			if block.start_pos == original_pos {
-				new_instrs.push(I64Const((block.cost + gas_function_cost) as i64));
+				new_instrs
+					.push(I64Const((block.cost.checked_add(gas_function_cost).ok_or(())?) as i64));
 				new_instrs.push(Call(gas_func));
 				true
 			} else {
@@ -722,7 +744,7 @@ mod tests {
 		);
 		let backend = host_function::Injector::new("env", "gas");
 		let injected_module =
-			super::inject(module, backend, &ConstantCostRules::new(1, 10_000)).unwrap();
+			super::inject(module, backend, &ConstantCostRules::new(1, 10_000, 1)).unwrap();
 
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
@@ -759,7 +781,7 @@ mod tests {
 		);
 		let backend = mutable_global::Injector::new("gas_left");
 		let injected_module =
-			super::inject(module, backend, &ConstantCostRules::new(1, 10_000)).unwrap();
+			super::inject(module, backend, &ConstantCostRules::new(1, 10_000, 1)).unwrap();
 
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),

src/gas_metering/validation.rs

@@ -134,6 +134,10 @@ fn build_control_flow_graph(
 
 	let mut stack = vec![ControlFrame::new(entry_node_id, terminal_node_id, false)];
 	let mut metered_blocks_iter = blocks.iter().peekable();
+
+	let locals_count = body.locals().iter().fold(0, |count, val_type| count + val_type.count());
+	let locals_init_cost = (rules.call_per_local_cost()).checked_mul(locals_count).ok_or(())?;
+
 	for (cursor, instruction) in body.code().elements().iter().enumerate() {
 		let active_node_id = stack
 			.last()
@@ -149,6 +153,10 @@ fn build_control_flow_graph(
 			graph.increment_charged_cost(active_node_id, next_metered_block.cost);
 		}
 
+		// Add locals initialization cost to the function block.
+		if cursor == 0 {
+			graph.increment_actual_cost(active_node_id, locals_init_cost);
+		}
 		let instruction_cost = rules.instruction_cost(instruction).ok_or(())?;
 		match instruction {
 			Instruction::Block(_) => {
@@ -342,8 +350,11 @@ mod tests {
 
 			for func_body in module.code_section().iter().flat_map(|section| section.bodies()) {
 				let rules = ConstantCostRules::default();
+				let locals_count =
+					func_body.locals().iter().fold(0, |count, val_type| count + val_type.count());
 
-				let metered_blocks = determine_metered_blocks(func_body.code(), &rules).unwrap();
+				let metered_blocks =
+					determine_metered_blocks(func_body.code(), &rules, locals_count).unwrap();
 				let success =
 					validate_metering_injections(func_body, &rules, &metered_blocks).unwrap();
 				assert!(success);

tests/expectations/gas/branch_host_fn.wat

@@ -4,7 +4,7 @@
   (import "env" "gas" (func (;0;) (type 1)))
   (func $fibonacci_with_break (;1;) (type 0) (result i32)
     (local i32 i32)
-    i64.const 13
+    i64.const 15
     call 0
     block ;; label = @1
       i32.const 0

tests/expectations/gas/branch_mut_global.wat

@@ -3,7 +3,7 @@
   (type (;1;) (func (param i64)))
   (func $fibonacci_with_break (;0;) (type 0) (result i32)
     (local $x i32) (local $y i32)
-    i64.const 24
+    i64.const 26
     call 1
     block ;; label = @1
       i32.const 0

tests/expectations/gas/call_host_fn.wat

@@ -4,7 +4,7 @@
   (import "env" "gas" (func (;0;) (type 1)))
   (func $add_locals (;1;) (type 0) (param $x i32) (param $y i32) (result i32)
     (local i32)
-    i64.const 5
+    i64.const 6
     call 0
     local.get $x
     local.get $y

tests/expectations/gas/call_mut_global.wat

@@ -3,7 +3,7 @@
   (type (;1;) (func (param i64)))
   (func $add_locals (;0;) (type 0) (param $x i32) (param $y i32) (result i32)
     (local $t i32)
-    i64.const 16
+    i64.const 17
     call 2
     local.get $x
     local.get $y