Improve performance of stack height injection

Specifically, - avoid using `import_count` and `function_space` since those are slow. - Avoid using the `builder` module since it requires recreation of the module essentially and is not efficient. - That also allows us to avoid creation of redundant signatures for thunks.
2026-06-13 07:01:11 +00:00 · 2022-10-19 17:56:29 +02:00
49 changed files with 440 additions and 2019 deletions
@@ -4,4 +4,3 @@ target
 .DS_Store
 .idea
 .vscode
 *~
@@ -1,91 +0,0 @@
 # Benchmarks
 ## Table of Contents
 - [Benchmark Results](#benchmark-results)
    - [coremark, instrumented](#coremark,-instrumented)
    - [recursive_ok, instrumented](#recursive_ok,-instrumented)
    - [fibonacci_recursive, instrumented](#fibonacci_recursive,-instrumented)
    - [factorial_recursive, instrumented](#factorial_recursive,-instrumented)
    - [count_until, instrumented](#count_until,-instrumented)
    - [memory_vec_add, instrumented](#memory_vec_add,-instrumented)
    - [wasm_kernel::tiny_keccak, instrumented](#wasm_kernel::tiny_keccak,-instrumented)
    - [global_bump, instrumented](#global_bump,-instrumented)
 ## Instrumented Modules sizes
 | fixture                      |  original size   | gas metered/host fn | gas metered/mut global | size diff |
 |------------------------------|------------------|---------------------|------------------------|-----------|
 | recursive_ok.wat             |             0 kb |         0 kb (137%) |            0 kb (177%) |      +29% |
 | count_until.wat              |             0 kb |         0 kb (125%) |            0 kb (153%) |      +21% |
 | global_bump.wat              |             0 kb |         0 kb (123%) |            0 kb (145%) |      +18% |
 | memory-vec-add.wat           |             0 kb |         0 kb (116%) |            0 kb (134%) |      +15% |
 | factorial.wat                |             0 kb |         0 kb (125%) |            0 kb (145%) |      +15% |
 | fibonacci.wat                |             0 kb |         0 kb (121%) |            0 kb (134%) |      +10% |
 | contract_terminate.wasm      |             1 kb |         1 kb (110%) |            1 kb (112%) |       +2% |
 | coremark_minimal.wasm        |             7 kb |         8 kb (114%) |            8 kb (115%) |       +0% |
 | trait_erc20.wasm             |            10 kb |        11 kb (108%) |           11 kb (108%) |       +0% |
 | rand_extension.wasm          |             4 kb |         5 kb (109%) |            5 kb (109%) |       +0% |
 | multisig.wasm                |            27 kb |        30 kb (110%) |           30 kb (110%) |       +0% |
 | wasm_kernel.wasm             |           779 kb |       787 kb (100%) |          795 kb (101%) |       +0% |
 | many_blocks.wasm             |          1023 kb |      2389 kb (233%) |         2389 kb (233%) |       +0% |
 | contract_transfer.wasm       |             7 kb |         8 kb (113%) |            8 kb (113%) |       +0% |
 | erc1155.wasm                 |            26 kb |        29 kb (111%) |           29 kb (111%) |       +0% |
 | erc20.wasm                   |             9 kb |        10 kb (108%) |           10 kb (109%) |       +0% |
 | dns.wasm                     |            10 kb |        11 kb (108%) |           11 kb (108%) |       +0% |
 | proxy.wasm                   |             3 kb |         4 kb (108%) |            4 kb (109%) |       +0% |
 | erc721.wasm                  |            13 kb |        14 kb (108%) |           14 kb (108%) |       +0% |
 ## Benchmark Results
 ### coremark, instrumented
 |        | `with host_function::Injector`          | `with mutable_global::Injector`           |
 |:-------|:----------------------------------------|:----------------------------------------- |
 |        | `20.81 s` (✅ **1.00x**)                 | `20.20 s` (✅ **1.03x faster**)            |
 ### recursive_ok, instrumented
 |        | `with host_function::Injector`          | `with mutable_global::Injector`           |
 |:-------|:----------------------------------------|:----------------------------------------- |
 |        | `367.11 us` (✅ **1.00x**)               | `585.39 us` (❌ *1.59x slower*)            |
 ### fibonacci_recursive, instrumented
 |        | `with host_function::Injector`          | `with mutable_global::Injector`           |
 |:-------|:----------------------------------------|:----------------------------------------- |
 |        | `9.15 us` (✅ **1.00x**)                 | `13.56 us` (❌ *1.48x slower*)             |
 ### factorial_recursive, instrumented
 |        | `with host_function::Injector`          | `with mutable_global::Injector`           |
 |:-------|:----------------------------------------|:----------------------------------------- |
 |        | `1.50 us` (✅ **1.00x**)                 | `1.98 us` (❌ *1.32x slower*)              |
 ### count_until, instrumented
 |        | `with host_function::Injector`          | `with mutable_global::Injector`           |
 |:-------|:----------------------------------------|:----------------------------------------- |
 |        | `5.03 ms` (✅ **1.00x**)                 | `8.13 ms` (❌ *1.62x slower*)              |
 ### memory_vec_add, instrumented
 |        | `with host_function::Injector`          | `with mutable_global::Injector`           |
 |:-------|:----------------------------------------|:----------------------------------------- |
 |        | `6.21 ms` (✅ **1.00x**)                 | `8.45 ms` (❌ *1.36x slower*)              |
 ### wasm_kernel::tiny_keccak, instrumented
 |        | `with host_function::Injector`          | `with mutable_global::Injector`           |
 |:-------|:----------------------------------------|:----------------------------------------- |
 |        | `925.22 us` (✅ **1.00x**)               | `1.08 ms` (❌ *1.17x slower*)              |
 ### global_bump, instrumented
 |        | `with host_function::Injector`          | `with mutable_global::Injector`           |
 |:-------|:----------------------------------------|:----------------------------------------- |
 |        | `3.79 ms` (✅ **1.00x**)                 | `7.03 ms` (❌ *1.86x slower*)              |
 ---
 Made with [criterion-table](https://github.com/nu11ptr/criterion-table)
@@ -16,15 +16,6 @@ The interface provided to smart contracts will adhere to semver with one excepti
 major version bumps will be backwards compatible with regard to already deployed contracts.
 In other words: Upgrading this pallet will not break pre-existing contracts.
 ## [Unreleased]
 ### New
 - 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]
 ### Changed
@@ -1,6 +1,6 @@
 [package]
 name = "wasm-instrument"
-version = "0.4.0"
+version = "0.3.0"
 edition = "2021"
 rust-version = "1.56.1"
 authors = ["Parity Technologies <admin@parity.io>"]
@@ -30,14 +30,10 @@ diff = "0.1"
 pretty_assertions = "1"
 rand = "0.8"
 wat = "1"
-wasmparser = "0.95"
+wasmparser = "0.92"
 wasmprinter = "0.2"
 wasmi = "0.20"
 [features]
 default = ["std"]
 std = ["parity-wasm/std"]
 sign_ext = ["parity-wasm/sign_ext"]
 [lib]
 bench = false
@@ -1,23 +1,20 @@
 use criterion::{
-	criterion_group, criterion_main, measurement::Measurement, Bencher, BenchmarkGroup, Criterion,
+	criterion_group, criterion_main, measurement::Measurement, BenchmarkGroup, Criterion,
 	Throughput,
 };
 use std::{
 	fs::{read, read_dir},
 	path::PathBuf,
 	slice,
 };
 use wasm_instrument::{
-	gas_metering::{self, host_function, mutable_global, Backend, ConstantCostRules},
+	gas_metering, inject_stack_limiter,
-	inject_stack_limiter,
+	parity_wasm::{deserialize_buffer, elements::Module},
 	parity_wasm::{deserialize_buffer, elements::Module, serialize},
 };
 fn fixture_dir() -> PathBuf {
 	let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
 	path.push("benches");
 	path.push("fixtures");
 	path.push("wasm");
 	path
 }
@@ -39,12 +36,7 @@ where
 fn gas_metering(c: &mut Criterion) {
 	let mut group = c.benchmark_group("Gas Metering");
 	any_fixture(&mut group, |module| {
-		gas_metering::inject(
+		gas_metering::inject(module, &gas_metering::ConstantCostRules::default(), "env").unwrap();
 			module,
 			host_function::Injector::new("env", "gas"),
 			&ConstantCostRules::default(),
 		)
 		.unwrap();
 	});
 }
@@ -55,520 +47,5 @@ fn stack_height_limiter(c: &mut Criterion) {
 	});
 }
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use wasmi::{
 	self,
 	core::{Pages, Value, F32},
 	Caller, Config, Engine, Extern, Func, Instance, Linker, Memory, StackLimits, Store,
 };
 fn prepare_module<P: Backend>(backend: P, input: &[u8]) -> (wasmi::Module, Store<u64>) {
 	let module = deserialize_buffer(input).unwrap();
 	let instrumented_module =
 		gas_metering::inject(module, backend, &ConstantCostRules::default()).unwrap();
 	let input = serialize(instrumented_module).unwrap();
 	// Prepare wasmi
 	let engine = Engine::new(&bench_config());
 	let module = wasmi::Module::new(&engine, &mut &input[..]).unwrap();
 	// Init host state with maximum gas_left
 	let store = Store::new(&engine, u64::MAX);
 	(module, store)
 }
 fn add_gas_host_func(linker: &mut Linker<u64>, store: &mut Store<u64>) {
 	// Create gas host function
 	let host_gas = Func::wrap(store, |mut caller: Caller<'_, u64>, param: u64| {
 		*caller.host_data_mut() -= param;
 	});
 	// Link the gas host function
 	linker.define("env", "gas", host_gas).unwrap();
 }
 fn add_gas_left_global(instance: &Instance, mut store: Store<u64>) -> Store<u64> {
 	instance
 		.get_export(&mut store, "gas_left")
 		.and_then(Extern::into_global)
 		.unwrap()
 		.set(&mut store, Value::I64(-1i64)) // the same as u64::MAX
 		.unwrap();
 	store
 }
 fn gas_metered_coremark(c: &mut Criterion) {
 	let mut group = c.benchmark_group("coremark, instrumented");
 	// Benchmark host_function::Injector
 	let wasm_filename = "coremark_minimal.wasm";
 	let bytes = read(fixture_dir().join(wasm_filename)).unwrap();
 	group.bench_function("with host_function::Injector", |bench| {
 		let backend = host_function::Injector::new("env", "gas");
 		let (module, mut store) = prepare_module(backend, &bytes);
 		// Link the host functions with the imported ones
 		let mut linker = <Linker<u64>>::new();
 		add_gas_host_func(&mut linker, &mut store);
 		// Create clock_ms host function.
 		let host_clock_ms = Func::wrap(&mut store, || {
 			SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as u64
 		});
 		// Link the time measurer for the coremark wasm
 		linker.define("env", "clock_ms", host_clock_ms).unwrap();
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		bench.iter(|| {
 			let run = instance
 				.get_export(&mut store, "run")
 				.and_then(Extern::into_func)
 				.unwrap()
 				.typed::<(), F32>(&mut store)
 				.unwrap();
 			// Call the wasm!
 			run.call(&mut store, ()).unwrap();
 		})
 	});
 	group.bench_function("with mutable_global::Injector", |bench| {
 		let backend = mutable_global::Injector::new("gas_left");
 		let (module, mut store) = prepare_module(backend, &bytes);
 		// Add the gas_left mutable global
 		let mut linker = <Linker<u64>>::new();
 		// Create clock_ms host function.
 		let host_clock_ms = Func::wrap(&mut store, || {
 			SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as u64
 		});
 		// Link the time measurer for the coremark wasm
 		linker.define("env", "clock_ms", host_clock_ms).unwrap();
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let mut store = add_gas_left_global(&instance, store);
 		bench.iter(|| {
 			let run = instance
 				.get_export(&mut store, "run")
 				.and_then(Extern::into_func)
 				.unwrap()
 				.typed::<(), F32>(&mut store)
 				.unwrap();
 			// Call the wasm!
 			run.call(&mut store, ()).unwrap();
 		})
 	});
 }
 /// Converts the `.wat` encoded `bytes` into `.wasm` encoded bytes.
 pub fn wat2wasm(bytes: &[u8]) -> Vec<u8> {
 	wat::parse_bytes(bytes).unwrap().into_owned()
 }
 /// Returns a [`Config`] useful for benchmarking.
 fn bench_config() -> Config {
 	let mut config = Config::default();
 	config.set_stack_limits(StackLimits::new(1024, 1024 * 1024, 64 * 1024).unwrap());
 	config
 }
 fn gas_metered_recursive_ok(c: &mut Criterion) {
 	let mut group = c.benchmark_group("recursive_ok, instrumented");
 	const RECURSIVE_DEPTH: i32 = 8000;
 	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/recursive_ok.wat"));
 	group.bench_function("with host_function::Injector", |bench| {
 		let backend = host_function::Injector::new("env", "gas");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Link the host function with the imported one
 		let mut linker = <Linker<u64>>::new();
 		add_gas_host_func(&mut linker, &mut store);
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let bench_call = instance.get_export(&store, "call").and_then(Extern::into_func).unwrap();
 		let mut result = [Value::I32(0)];
 		bench.iter(|| {
 			bench_call
 				.call(&mut store, &[Value::I32(RECURSIVE_DEPTH)], &mut result)
 				.unwrap();
 			assert_eq!(result, [Value::I32(0)]);
 		})
 	});
 	group.bench_function("with mutable_global::Injector", |bench| {
 		let backend = mutable_global::Injector::new("gas_left");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Add the gas_left mutable global
 		let linker = <Linker<u64>>::new();
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let mut store = add_gas_left_global(&instance, store);
 		let bench_call = instance.get_export(&store, "call").and_then(Extern::into_func).unwrap();
 		let mut result = [Value::I32(0)];
 		bench.iter(|| {
 			bench_call
 				.call(&mut store, &[Value::I32(RECURSIVE_DEPTH)], &mut result)
 				.unwrap();
 			assert_eq!(result, [Value::I32(0)]);
 		})
 	});
 }
 fn gas_metered_fibonacci_recursive(c: &mut Criterion) {
 	let mut group = c.benchmark_group("fibonacci_recursive, instrumented");
 	const FIBONACCI_REC_N: i64 = 10;
 	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/fibonacci.wat"));
 	group.bench_function("with host_function::Injector", |bench| {
 		let backend = host_function::Injector::new("env", "gas");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Link the host function with the imported one
 		let mut linker = <Linker<u64>>::new();
 		add_gas_host_func(&mut linker, &mut store);
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let bench_call = instance
 			.get_export(&store, "fib_recursive")
 			.and_then(Extern::into_func)
 			.unwrap();
 		let mut result = [Value::I32(0)];
 		bench.iter(|| {
 			bench_call
 				.call(&mut store, &[Value::I64(FIBONACCI_REC_N)], &mut result)
 				.unwrap();
 		});
 	});
 	group.bench_function("with mutable_global::Injector", |bench| {
 		let backend = mutable_global::Injector::new("gas_left");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Add the gas_left mutable global
 		let linker = <Linker<u64>>::new();
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let mut store = add_gas_left_global(&instance, store);
 		let bench_call = instance
 			.get_export(&store, "fib_recursive")
 			.and_then(Extern::into_func)
 			.unwrap();
 		let mut result = [Value::I32(0)];
 		bench.iter(|| {
 			bench_call
 				.call(&mut store, &[Value::I64(FIBONACCI_REC_N)], &mut result)
 				.unwrap();
 		});
 	});
 }
 fn gas_metered_fac_recursive(c: &mut Criterion) {
 	let mut group = c.benchmark_group("factorial_recursive, instrumented");
 	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/factorial.wat"));
 	group.bench_function("with host_function::Injector", |b| {
 		let backend = host_function::Injector::new("env", "gas");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Link the host function with the imported one
 		let mut linker = <Linker<u64>>::new();
 		add_gas_host_func(&mut linker, &mut store);
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let fac = instance
 			.get_export(&store, "recursive_factorial")
 			.and_then(Extern::into_func)
 			.unwrap();
 		let mut result = [Value::I64(0)];
 		b.iter(|| {
 			fac.call(&mut store, &[Value::I64(25)], &mut result).unwrap();
 			assert_eq!(result, [Value::I64(7034535277573963776)]);
 		})
 	});
 	group.bench_function("with mutable_global::Injector", |b| {
 		let backend = mutable_global::Injector::new("gas_left");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Add the gas_left mutable global
 		let linker = <Linker<u64>>::new();
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let mut store = add_gas_left_global(&instance, store);
 		let fac = instance
 			.get_export(&store, "recursive_factorial")
 			.and_then(Extern::into_func)
 			.unwrap();
 		let mut result = [Value::I64(0)];
 		b.iter(|| {
 			fac.call(&mut store, &[Value::I64(25)], &mut result).unwrap();
 			assert_eq!(result, [Value::I64(7034535277573963776)]);
 		})
 	});
 }
 fn gas_metered_count_until(c: &mut Criterion) {
 	const COUNT_UNTIL: i32 = 100_000;
 	let mut group = c.benchmark_group("count_until, instrumented");
 	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/count_until.wat"));
 	group.bench_function("with host_function::Injector", |b| {
 		let backend = host_function::Injector::new("env", "gas");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Link the host function with the imported one
 		let mut linker = <Linker<u64>>::new();
 		add_gas_host_func(&mut linker, &mut store);
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let count_until =
 			instance.get_export(&store, "count_until").and_then(Extern::into_func).unwrap();
 		let mut result = [Value::I32(0)];
 		b.iter(|| {
 			count_until.call(&mut store, &[Value::I32(COUNT_UNTIL)], &mut result).unwrap();
 			assert_eq!(result, [Value::I32(COUNT_UNTIL)]);
 		})
 	});
 	group.bench_function("with mutable_global::Injector", |b| {
 		let backend = mutable_global::Injector::new("gas_left");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Add the gas_left mutable global
 		let linker = <Linker<u64>>::new();
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let mut store = add_gas_left_global(&instance, store);
 		let count_until =
 			instance.get_export(&store, "count_until").and_then(Extern::into_func).unwrap();
 		let mut result = [Value::I32(0)];
 		b.iter(|| {
 			count_until.call(&mut store, &[Value::I32(COUNT_UNTIL)], &mut result).unwrap();
 			assert_eq!(result, [Value::I32(COUNT_UNTIL)]);
 		})
 	});
 }
 fn gas_metered_vec_add(c: &mut Criterion) {
 	fn test_for<A, B>(
 		b: &mut Bencher,
 		vec_add: Func,
 		mut store: &mut Store<u64>,
 		mem: Memory,
 		len: usize,
 		vec_a: A,
 		vec_b: B,
 	) where
 		A: IntoIterator<Item = i32>,
 		B: IntoIterator<Item = i32>,
 	{
 		use core::mem::size_of;
 		let ptr_result = 10;
 		let len_result = len * size_of::<i64>();
 		let ptr_a = ptr_result + len_result;
 		let len_a = len * size_of::<i32>();
 		let ptr_b = ptr_a + len_a;
 		// Reset `result` buffer to zeros:
 		mem.data_mut(&mut store)[ptr_result..ptr_result + (len * size_of::<i32>())].fill(0);
 		// Initialize `a` buffer:
 		for (n, a) in vec_a.into_iter().take(len).enumerate() {
 			mem.write(&mut store, ptr_a + (n * size_of::<i32>()), &a.to_le_bytes()).unwrap();
 		}
 		// Initialize `b` buffer:
 		for (n, b) in vec_b.into_iter().take(len).enumerate() {
 			mem.write(&mut store, ptr_b + (n * size_of::<i32>()), &b.to_le_bytes()).unwrap();
 		}
 		// Prepare parameters and all Wasm `vec_add`:
 		let params = [
 			Value::I32(ptr_result as i32),
 			Value::I32(ptr_a as i32),
 			Value::I32(ptr_b as i32),
 			Value::I32(len as i32),
 		];
 		b.iter(|| {
 			vec_add.call(&mut store, &params, &mut []).unwrap();
 		});
 		// Validate the result buffer:
 		for n in 0..len {
 			let mut buffer4 = [0x00; 4];
 			let mut buffer8 = [0x00; 8];
 			let a = {
 				mem.read(&store, ptr_a + (n * size_of::<i32>()), &mut buffer4).unwrap();
 				i32::from_le_bytes(buffer4)
 			};
 			let b = {
 				mem.read(&store, ptr_b + (n * size_of::<i32>()), &mut buffer4).unwrap();
 				i32::from_le_bytes(buffer4)
 			};
 			let actual_result = {
 				mem.read(&store, ptr_result + (n * size_of::<i64>()), &mut buffer8).unwrap();
 				i64::from_le_bytes(buffer8)
 			};
 			let expected_result = (a as i64) + (b as i64);
 			assert_eq!(
 				expected_result, actual_result,
 				"given a = {a} and b = {b}, results diverge at index {n}"
 			);
 		}
 	}
 	let mut group = c.benchmark_group("memory_vec_add, instrumented");
 	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/memory-vec-add.wat"));
 	const LEN: usize = 100_000;
 	group.bench_function("with host_function::Injector", |b| {
 		let backend = host_function::Injector::new("env", "gas");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Link the host function with the imported one
 		let mut linker = <Linker<u64>>::new();
 		add_gas_host_func(&mut linker, &mut store);
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let vec_add = instance.get_export(&store, "vec_add").and_then(Extern::into_func).unwrap();
 		let mem = instance.get_export(&store, "mem").and_then(Extern::into_memory).unwrap();
 		mem.grow(&mut store, Pages::new(25).unwrap()).unwrap();
 		test_for(
 			b,
 			vec_add,
 			&mut store,
 			mem,
 			LEN,
 			(0..LEN).map(|i| (i * i) as i32),
 			(0..LEN).map(|i| (i * 10) as i32),
 		)
 	});
 	group.bench_function("with mutable_global::Injector", |b| {
 		let backend = mutable_global::Injector::new("gas_left");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Add the gas_left mutable global
 		let linker = <Linker<u64>>::new();
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let mut store = add_gas_left_global(&instance, store);
 		let vec_add = instance.get_export(&store, "vec_add").and_then(Extern::into_func).unwrap();
 		let mem = instance.get_export(&store, "mem").and_then(Extern::into_memory).unwrap();
 		mem.grow(&mut store, Pages::new(25).unwrap()).unwrap();
 		test_for(
 			b,
 			vec_add,
 			&mut store,
 			mem,
 			LEN,
 			(0..LEN).map(|i| (i * i) as i32),
 			(0..LEN).map(|i| (i * 10) as i32),
 		)
 	});
 }
 fn gas_metered_tiny_keccak(c: &mut Criterion) {
 	let mut group = c.benchmark_group("wasm_kernel::tiny_keccak, instrumented");
 	let wasm_filename = "wasm_kernel.wasm";
 	let wasm_bytes = read(fixture_dir().join(wasm_filename)).unwrap();
 	group.bench_function("with host_function::Injector", |b| {
 		let backend = host_function::Injector::new("env", "gas");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Link the host function with the imported one
 		let mut linker = <Linker<u64>>::new();
 		add_gas_host_func(&mut linker, &mut store);
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let prepare = instance
 			.get_export(&store, "prepare_tiny_keccak")
 			.and_then(Extern::into_func)
 			.unwrap();
 		let keccak = instance
 			.get_export(&store, "bench_tiny_keccak")
 			.and_then(Extern::into_func)
 			.unwrap();
 		let mut test_data_ptr = Value::I32(0);
 		prepare.call(&mut store, &[], slice::from_mut(&mut test_data_ptr)).unwrap();
 		b.iter(|| {
 			keccak.call(&mut store, slice::from_ref(&test_data_ptr), &mut []).unwrap();
 		})
 	});
 	group.bench_function("with mutable_global::Injector", |b| {
 		let backend = mutable_global::Injector::new("gas_left");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Add the gas_left mutable global
 		let linker = <Linker<u64>>::new();
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let mut store = add_gas_left_global(&instance, store);
 		let prepare = instance
 			.get_export(&store, "prepare_tiny_keccak")
 			.and_then(Extern::into_func)
 			.unwrap();
 		let keccak = instance
 			.get_export(&store, "bench_tiny_keccak")
 			.and_then(Extern::into_func)
 			.unwrap();
 		let mut test_data_ptr = Value::I32(0);
 		prepare.call(&mut store, &[], slice::from_mut(&mut test_data_ptr)).unwrap();
 		b.iter(|| {
 			keccak.call(&mut store, slice::from_ref(&test_data_ptr), &mut []).unwrap();
 		})
 	});
 }
 fn gas_metered_global_bump(c: &mut Criterion) {
 	const BUMP_AMOUNT: i32 = 100_000;
 	let mut group = c.benchmark_group("global_bump, instrumented");
 	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/global_bump.wat"));
 	group.bench_function("with host_function::Injector", |b| {
 		let backend = host_function::Injector::new("env", "gas");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Link the host function with the imported one
 		let mut linker = <Linker<u64>>::new();
 		add_gas_host_func(&mut linker, &mut store);
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let bump = instance.get_export(&store, "bump").and_then(Extern::into_func).unwrap();
 		let mut result = [Value::I32(0)];
 		b.iter(|| {
 			bump.call(&mut store, &[Value::I32(BUMP_AMOUNT)], &mut result).unwrap();
 			assert_eq!(result, [Value::I32(BUMP_AMOUNT)]);
 		})
 	});
 	group.bench_function("with mutable_global::Injector", |b| {
 		let backend = mutable_global::Injector::new("gas_left");
 		let (module, mut store) = prepare_module(backend, &wasm_bytes);
 		// Add the gas_left mutable global
 		let linker = <Linker<u64>>::new();
 		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
 		let mut store = add_gas_left_global(&instance, store);
 		let bump = instance.get_export(&store, "bump").and_then(Extern::into_func).unwrap();
 		let mut result = [Value::I32(0)];
 		b.iter(|| {
 			bump.call(&mut store, &[Value::I32(BUMP_AMOUNT)], &mut result).unwrap();
 			assert_eq!(result, [Value::I32(BUMP_AMOUNT)]);
 		})
 	});
 }
 criterion_group!(benches, gas_metering, stack_height_limiter);
-criterion_group!(
+criterion_main!(benches);
 	name = coremark;
 	config = Criterion::default()
 		.sample_size(10)
 		.measurement_time(Duration::from_millis(275000))
 		.warm_up_time(Duration::from_millis(1000));
 	targets =
 		 gas_metered_coremark,
 );
 criterion_group!(
 	name = wasmi_fixtures;
 	config = Criterion::default()
 		.sample_size(10)
 		.measurement_time(Duration::from_millis(250000))
 	.warm_up_time(Duration::from_millis(1000));
 	targets =
 		gas_metered_recursive_ok,
 		gas_metered_fibonacci_recursive,
 		gas_metered_fac_recursive,
 		gas_metered_count_until,
 		gas_metered_vec_add,
 		gas_metered_tiny_keccak,
 		gas_metered_global_bump,
 );
 criterion_main!(coremark, wasmi_fixtures);
@@ -1,25 +0,0 @@
 ;; Exports a function `count_until` that takes an input `n`.
 ;; The exported function counts an integer `n` times and then returns `n`.
 (module
  (func $count_until (export "count_until") (param $limit i32) (result i32)
    (local $counter i32)
    (block
        (loop
            (br_if
                1
                (i32.eq
                    (local.tee $counter
                        (i32.add
                            (local.get $counter)
                            (i32.const 1)
                        )
                    )
                    (local.get $limit)
                )
            )
            (br 0)
        )
    )
    (return (local.get $counter))
  )
 )
@@ -1,35 +0,0 @@
 (module
    ;; Iterative factorial function, does not use recursion.
    (func (export "iterative_factorial") (param i64) (result i64)
        (local i64)
        (local.set 1 (i64.const 1))
        (block
            (br_if 0 (i64.lt_s (local.get 0) (i64.const 2)))
            (loop
                (local.set 1 (i64.mul (local.get 1) (local.get 0)))
                (local.set 0 (i64.add (local.get 0) (i64.const -1)))
                (br_if 0 (i64.gt_s (local.get 0) (i64.const 1)))
            )
        )
        (local.get 1)
    )
    ;; Recursive trivial factorial function.
    (func $rec_fac (export "recursive_factorial") (param i64) (result i64)
        (if (result i64)
            (i64.eq (local.get 0) (i64.const 0))
            (then (i64.const 1))
            (else
                (i64.mul
                    (local.get 0)
                    (call $rec_fac
                        (i64.sub
                            (local.get 0)
                            (i64.const 1)
                        )
                    )
                )
            )
        )
    )
 )
@@ -1,47 +0,0 @@
 (module
    (func $fib_recursive (export "fib_recursive") (param $N i64) (result i64)
        (if
            (i64.le_s (local.get $N) (i64.const 1))
            (then (return (local.get $N)))
        )
        (return
            (i64.add
                (call $fib_recursive
                  (i64.sub (local.get $N) (i64.const 1))
                )
                (call $fib_recursive
                  (i64.sub (local.get $N) (i64.const 2))
                )
            )
        )
    )
    (func $fib_iterative (export "fib_iterative") (param $N i64) (result i64)
        (local $n1 i64)
        (local $n2 i64)
        (local $tmp i64)
        (local $i i64)
        ;; return $N for N <= 1
        (if
            (i64.le_s (local.get $N) (i64.const 1))
            (then (return (local.get $N)))
        )
        (local.set $n1 (i64.const 1))
        (local.set $n2 (i64.const 1))
        (local.set $i (i64.const 2))
        ;;since we normally return n2, handle n=1 case specially
        (loop $again
            (if
                (i64.lt_s (local.get $i) (local.get $N))
                (then
                    (local.set $tmp (i64.add (local.get $n1) (local.get $n2)))
                    (local.set $n1 (local.get $n2))
                    (local.set $n2 (local.get $tmp))
                    (local.set $i (i64.add (local.get $i) (i64.const 1)))
                    (br $again)
                )
            )
        )
        (local.get $n2)
    )
 )
@@ -1,27 +0,0 @@
 ;; Exports a function `bump` that takes an input `n`.
 ;; The exported function bumps a global variable `n` times and then returns it.
 (module
    (global $g (mut i32) (i32.const 0))
    (func $bump (export "bump") (param $n i32) (result i32)
        (global.set $g (i32.const 0))
        (block $break
            (loop $continue
                (br_if ;; if $g == $n then break
                    $break
                    (i32.eq
                        (global.get $g)
                        (local.get $n)
                    )
                )
                (global.set $g ;; $g += 1
                    (i32.add
                        (global.get $g)
                        (i32.const 1)
                    )
                )
                (br $continue)
            )
        )
        (return (global.get $g))
    )
 )
@@ -1,58 +0,0 @@
 ;; Exports a function `vec_add` that computes the addition of 2 vectors
 ;; of length `len` starting at `ptr_a` and `ptr_b` and stores the result
 ;; into a buffer of the same length starting at `ptr_result`.
 (module
    (memory (export "mem") 1)
    (func (export "vec_add")
        (param $ptr_result i32)
        (param $ptr_a i32)
        (param $ptr_b i32)
        (param $len i32)
        (local $n i32)
        (block $exit
            (loop $loop
                (br_if ;; exit loop if $n == $len
                    $exit
                    (i32.eq
                        (local.get $n)
                        (local.get $len)
                    )
                )
                (i64.store offset=0 ;; ptr_result[n] = ptr_a[n] + ptr_b[n]
                    (i32.add
                        (local.get $ptr_result)
                        (i32.mul
                            (local.get $n)
                            (i32.const 8)
                        )
                    )
                    (i64.add
                        (i64.load32_s offset=0 ;; load ptr_a[n]
                            (i32.add
                                (local.get $ptr_a)
                                (i32.mul
                                    (local.get $n)
                                    (i32.const 4)
                                )
                            )
                        )
                        (i64.load32_s offset=0 ;; load ptr_b[n]
                            (i32.add
                                (local.get $ptr_b)
                                (i32.mul
                                    (local.get $n)
                                    (i32.const 4)
                                )
                            )
                        )
                    )
                )
                (local.set $n ;; increment n
                    (i32.add (local.get $n) (i32.const 1))
                 )
                (br $loop) ;; continue loop
            )
        )
        (return)
    )
 )
@@ -1,22 +0,0 @@
 ;; Exports a function `call` that takes an input `n`.
 ;; The exported function calls itself `n` times.
 (module
    (func $call (export "call") (param $n i32) (result i32)
        (if (result i32)
            (local.get $n)
            (then
                (return
                    (call $call
                        (i32.sub
                            (local.get $n)
                            (i32.const 1)
                        )
                    )
                )
            )
            (else
                (return (local.get $n))
            )
        )
    )
 )
@@ -1,138 +0,0 @@
 //! Provides backends for the gas metering instrumentation
 use parity_wasm::elements;
 /// Implementation details of the specific method of the gas metering.
 #[derive(Clone)]
 pub enum GasMeter {
 	/// Gas metering with an external function.
 	External {
 		/// Name of the module to import the gas function from.
 		module: &'static str,
 		/// Name of the external gas function to be imported.
 		function: &'static str,
 	},
 	/// Gas metering with a local function and a mutable global.
 	Internal {
 		/// Name of the mutable global to be exported.
 		global: &'static str,
 		/// Body of the local gas counting function to be injected.
 		func_instructions: elements::Instructions,
 		/// Cost of the gas function execution.
 		cost: u64,
 	},
 }
 use super::Rules;
 /// Under the hood part of the gas metering mechanics.
 pub trait Backend {
 	/// Provides the gas metering implementation details.  
 	fn gas_meter<R: Rules>(self, module: &elements::Module, rules: &R) -> GasMeter;
 }
 /// Gas metering with an external host function.
 pub mod host_function {
 	use super::{Backend, GasMeter, Rules};
 	use parity_wasm::elements::Module;
 	/// Injects invocations of the gas charging host function into each metering block.
 	pub struct Injector {
 		/// The name of the module to import the gas function from.
 		module: &'static str,
 		/// The name of the gas function to import.
 		name: &'static str,
 	}
 	impl Injector {
 		pub fn new(module: &'static str, name: &'static str) -> Self {
 			Self { module, name }
 		}
 	}
 	impl Backend for Injector {
 		fn gas_meter<R: Rules>(self, _module: &Module, _rules: &R) -> GasMeter {
 			GasMeter::External { module: self.module, function: self.name }
 		}
 	}
 }
 /// Gas metering with a mutable global.
 ///
 /// # Note
 ///
 /// Not for all execution engines this method gives performance wins compared to using an [external
 /// host function](host_function). See benchmarks and size overhead tests for examples of how to
 /// make measurements needed to decide which gas metering method is better for your particular case.
 ///
 /// # Warning
 ///
 /// It is not recommended to apply [stack limiter](crate::inject_stack_limiter) instrumentation to a
 /// module instrumented with this type of gas metering. This could lead to a massive module size
 /// bloat. This is a known issue to be fixed in upcoming versions.
 pub mod mutable_global {
 	use super::{Backend, GasMeter, Rules};
 	use alloc::vec;
 	use parity_wasm::elements::{self, Instruction, Module};
 	/// Injects a mutable global variable and a local function to the module to track
 	/// current gas left.
 	///
 	/// The function is called in every metering block. In case of falling out of gas, the global is
 	/// set to the sentinel value `U64::MAX` and `unreachable` instruction is called. The execution
 	/// engine should take care of getting the current global value and setting it back in order to
 	/// sync the gas left value during an execution.
 	pub struct Injector {
 		/// The export name of the gas tracking global.
 		pub global_name: &'static str,
 	}
 	impl Injector {
 		pub fn new(global_name: &'static str) -> Self {
 			Self { global_name }
 		}
 	}
 	impl Backend for Injector {
 		fn gas_meter<R: Rules>(self, module: &Module, rules: &R) -> GasMeter {
 			let gas_global_idx = module.globals_space() as u32;
 			let func_instructions = vec![
 				Instruction::GetGlobal(gas_global_idx),
 				Instruction::GetLocal(0),
 				Instruction::I64GeU,
 				Instruction::If(elements::BlockType::NoResult),
 				Instruction::GetGlobal(gas_global_idx),
 				Instruction::GetLocal(0),
 				Instruction::I64Sub,
 				Instruction::SetGlobal(gas_global_idx),
 				Instruction::Else,
 				// sentinel val u64::MAX
 				Instruction::I64Const(-1i64),           // non-charged instruction
 				Instruction::SetGlobal(gas_global_idx), // non-charged instruction
 				Instruction::Unreachable,               // non-charged instruction
 				Instruction::End,
 				Instruction::End,
 			];
 			// calculate gas used for the gas charging func execution itself
 			let mut gas_fn_cost = func_instructions.iter().fold(0, |cost, instruction| {
 				cost + (rules.instruction_cost(instruction).unwrap_or(0) as u64)
 			});
 			// don't charge for the instructions used to fail when out of gas
 			let fail_cost = vec![
 				Instruction::I64Const(-1i64),           // non-charged instruction
 				Instruction::SetGlobal(gas_global_idx), // non-charged instruction
 				Instruction::Unreachable,               // non-charged instruction
 			]
 			.iter()
 			.fold(0, |cost, instruction| {
 				cost + (rules.instruction_cost(instruction).unwrap_or(0) as u64)
 			});
 			gas_fn_cost -= fail_cost;
 			GasMeter::Internal {
 				global: self.global_name,
 				func_instructions: elements::Instructions::new(func_instructions),
 				cost: gas_fn_cost,
 			}
 		}
 	}
 }
@@ -1,13 +1,9 @@
-//! This module is used to instrument a Wasm module with the gas metering code.
+//! This module is used to instrument a Wasm module with gas metering code.
 //!
 //! The primary public interface is the [`inject`] function which transforms a given
 //! module into one that charges gas for code to be executed. See function documentation for usage
 //! and details.
 mod backend;
 pub use backend::{host_function, mutable_global, Backend, GasMeter};
 #[cfg(test)]
 mod validation;
@@ -36,9 +32,6 @@ 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.
@@ -74,12 +67,11 @@ impl MemoryGrowCost {
 /// # Note
 ///
 /// In a production environment it usually makes no sense to assign every instruction
-/// the same cost. A proper implemention of [`Rules`] should be provided that is probably
+/// the same cost. A proper implemention of [`Rules`] should be prived that is probably
 /// created by benchmarking.
 pub struct ConstantCostRules {
 	instruction_cost: u32,
 	memory_grow_cost: u32,
 	call_per_local_cost: u32,
 }
 impl ConstantCostRules {
@@ -87,15 +79,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, call_per_local_cost: u32) -> Self {
+	pub fn new(instruction_cost: u32, memory_grow_cost: u32) -> Self {
-		Self { instruction_cost, memory_grow_cost, call_per_local_cost }
+		Self { instruction_cost, memory_grow_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, call_per_local_cost: 1 }
+		Self { instruction_cost: 1, memory_grow_cost: 0 }
 	}
 }
@@ -107,167 +99,86 @@ 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.
+/// Transforms a given module into one that charges gas for code to be executed by proxy of an
 /// imported gas metering function.
 ///
-/// The output module uses the `gas` function to track the gas spent. The function could be either
+/// The output module imports a function "gas" from the specified module with type signature
-/// an imported or a local one modifying a mutable global. The argument is the amount of gas
+/// [i64] -> []. The argument is the amount of gas required to continue execution. The external
-/// required to continue execution. The execution engine is meant to keep track of the total amount
+/// function is meant to keep track of the total amount of gas used and trap or otherwise halt
-/// of gas used and trap or otherwise halt execution of the runtime if the gas usage exceeds some
+/// execution of the runtime if the gas usage exceeds some allowed limit.
 /// allowed limit.
 ///
-/// The body of each function of the original module is divided into metered blocks, and the calls
+/// The body of each function is divided into metered blocks, and the calls to charge gas are
-/// to charge gas are inserted at the beginning of every such block of code. A metered block is
+/// inserted at the beginning of every such block of code. A metered block is defined so that,
-/// defined so that, unless there is a trap, either all of the instructions are executed or none
+/// unless there is a trap, either all of the instructions are executed or none are. These are
-/// are. These are similar to basic blocks in a control flow graph, except that in some cases
+/// similar to basic blocks in a control flow graph, except that in some cases multiple basic
-/// multiple basic blocks can be merged into a single metered block. This is the case if any path
+/// blocks can be merged into a single metered block. This is the case if any path through the
-/// through the control flow graph containing one basic block also contains another.
+/// control flow graph containing one basic block also contains another.
 ///
-/// Charging gas at the beginning of each metered block ensures that 1) all instructions
+/// 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
+/// unless execution traps, and 3) the number of calls to "gas" is minimized. The corollary is that
-/// that modules instrumented with this metering code may charge gas for instructions not
+/// modules instrumented with this metering code may charge gas for instructions not executed in
-/// executed in the event of a trap.
+/// the event of a trap.
 ///
-/// Additionally, each `memory.grow` instruction found in the module is instrumented to first
+/// Additionally, each `memory.grow` instruction found in the module is instrumented to first make
-/// make a call to charge gas for the additional pages requested. This cannot be done as part of
+/// a call to charge gas for the additional pages requested. This cannot be done as part of the
-/// the block level gas charges as the gas cost is not static and depends on the stack argument
+/// block level gas charges as the gas cost is not static and depends on the stack argument to
-/// to `memory.grow`.
+/// `memory.grow`.
 ///
 /// The above transformations are performed for every function body defined in the module. This
 /// 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. If
+/// the addition of an imported functions changes the indices of module-defined functions. If the
-/// the module has a `NameSection`, added by calling `parse_names`, the indices will also be
+/// the module has a NameSection, added by calling `parse_names`, the indices will also be updated.
 /// updated.
 ///
 /// Syncronizing the amount of gas charged with the execution engine can be done in two ways. The
 /// first way is by calling the imported `gas` host function, see [`host_function`] for details. The
 /// second way is by using a local `gas` function together with a mutable global, see
 /// [`mutable_global`] for details.
 ///
 /// 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`.
+/// the original module as an Err.
-pub fn inject<R: Rules, B: Backend>(
+pub fn inject<R: Rules>(
 	module: elements::Module,
 	backend: B,
 	rules: &R,
 	gas_module_name: &str,
 ) -> Result<elements::Module, elements::Module> {
-	// Prepare module and return the gas function
+	// Injecting gas counting external
 	let gas_meter = backend.gas_meter(&module, rules);
 	let import_count = module.import_count(elements::ImportCountType::Function) as u32;
 	let functions_space = module.functions_space() as u32;
 	let gas_global_idx = module.globals_space() as u32;
 	let mut mbuilder = builder::from_module(module);
 	let import_sig =
 		mbuilder.push_signature(builder::signature().with_param(ValueType::I64).build_sig());
-	// Calculate the indexes and gas function cost,
+	mbuilder.push_import(
-	// for external gas function the cost is counted on the host side
+		builder::import()
-	let (gas_func_idx, total_func, gas_fn_cost) = match gas_meter {
+			.module(gas_module_name)
-		GasMeter::External { module: gas_module, function } => {
+			.field("gas")
-			// Inject the import of the gas function
+			.external()
-			let import_sig = mbuilder
+			.func(import_sig)
-				.push_signature(builder::signature().with_param(ValueType::I64).build_sig());
+			.build(),
-			mbuilder.push_import(
+	);
 				builder::import()
 					.module(gas_module)
 					.field(function)
 					.external()
 					.func(import_sig)
 					.build(),
 			);
-			(import_count, functions_space + 1, 0)
+	// back to plain module
 		},
 		GasMeter::Internal { global, ref func_instructions, cost } => {
 			// Inject the gas counting global
 			mbuilder.push_global(
 				builder::global()
 					.with_type(ValueType::I64)
 					.mutable()
 					.init_expr(Instruction::I64Const(0))
 					.build(),
 			);
 			// Inject the export entry for the gas counting global
 			let ebuilder = builder::ExportBuilder::new();
 			let global_export = ebuilder
 				.field(global)
 				.with_internal(elements::Internal::Global(gas_global_idx))
 				.build();
 			mbuilder.push_export(global_export);
 			let func_idx = functions_space;
 			// Build local gas function
 			let gas_func_sig =
 				builder::SignatureBuilder::new().with_param(ValueType::I64).build_sig();
 			let function = builder::FunctionBuilder::new()
 				.with_signature(gas_func_sig)
 				.body()
 				.with_instructions(func_instructions.clone())
 				.build()
 				.build();
 			// Inject local gas function
 			mbuilder.push_function(function);
 			(func_idx, func_idx + 1, cost)
 		},
 	};
 	// We need the built the module for making injections to its blocks
 	let mut module = mbuilder.build();
 	// calculate actual function index of the imported definition
 	//    (subtract all imports that are NOT functions)
 	let gas_func = module.import_count(elements::ImportCountType::Function) as u32 - 1;
 	let total_func = module.functions_space() as u32;
 	let mut need_grow_counter = false;
 	let mut error = false;
-	// Iterate over module sections and perform needed transformations.
+	// Updating calling addresses (all calls to function index >= `gas_func` should be incremented)
 	// Indexes are needed to be fixed up in `GasMeter::External` case, as it adds an imported
 	// function, which goes to the beginning of the module's functions space.
 	for section in module.sections_mut() {
 		match section {
-			elements::Section::Code(code_section) => {
+			elements::Section::Code(code_section) =>
-				let injection_targets = match gas_meter {
+				for func_body in code_section.bodies_mut() {
-					GasMeter::External { .. } => code_section.bodies_mut().as_mut_slice(),
+					for instruction in func_body.code_mut().elements_mut().iter_mut() {
-					// Don't inject counters to the local gas function, which is the last one as
+						if let Instruction::Call(call_index) = instruction {
-					// it's just added. Cost for its execution is added statically before each
+							if *call_index >= gas_func {
-					// invocation (see `inject_counter()`).
+								*call_index += 1
 					GasMeter::Internal { .. } => {
 						let len = code_section.bodies().len();
 						&mut code_section.bodies_mut()[..len - 1]
 					},
 				};
 				for func_body in injection_targets {
 					// Increment calling addresses if needed
 					if let GasMeter::External { .. } = gas_meter {
 						for instruction in func_body.code_mut().elements_mut().iter_mut() {
 							if let Instruction::Call(call_index) = instruction {
 								if *call_index >= gas_func_idx {
 									*call_index += 1
 								}
 							}
 						}
 					}
-					let locals_count =
+					if inject_counter(func_body.code_mut(), rules, gas_func).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
 					}
@@ -276,50 +187,42 @@ pub fn inject<R: Rules, B: Backend>(
 					{
 						need_grow_counter = true;
 					}
-				}
+				},
-			},
+			elements::Section::Export(export_section) => {
-			elements::Section::Export(export_section) =>
+				for export in export_section.entries_mut() {
-				if let GasMeter::External { module: _, function: _ } = gas_meter {
+					if let elements::Internal::Function(func_index) = export.internal_mut() {
-					for export in export_section.entries_mut() {
+						if *func_index >= gas_func {
-						if let elements::Internal::Function(func_index) = export.internal_mut() {
+							*func_index += 1
 							if *func_index >= gas_func_idx {
 								*func_index += 1
 							}
 						}
 					}
-				},
+				}
 			},
 			elements::Section::Element(elements_section) => {
 				// Note that we do not need to check the element type referenced because in the
 				// WebAssembly 1.0 spec, the only allowed element type is funcref.
-				if let GasMeter::External { .. } = gas_meter {
+				for segment in elements_section.entries_mut() {
-					for segment in elements_section.entries_mut() {
+					// update all indirect call addresses initial values
-						// update all indirect call addresses initial values
+					for func_index in segment.members_mut() {
-						for func_index in segment.members_mut() {
+						if *func_index >= gas_func {
-							if *func_index >= gas_func_idx {
+							*func_index += 1
 								*func_index += 1
 							}
 						}
 					}
 				}
 			},
 			elements::Section::Start(start_idx) =>
-				if let GasMeter::External { .. } = gas_meter {
+				if *start_idx >= gas_func {
-					if *start_idx >= gas_func_idx {
+					*start_idx += 1
 						*start_idx += 1
 					}
 				},
 			elements::Section::Name(s) =>
-				if let GasMeter::External { .. } = gas_meter {
+				for functions in s.functions_mut() {
-					for functions in s.functions_mut() {
+					*functions.names_mut() =
-						*functions.names_mut() =
+						IndexMap::from_iter(functions.names().iter().map(|(mut idx, name)| {
-							IndexMap::from_iter(functions.names().iter().map(|(mut idx, name)| {
+							if idx >= gas_func {
-								if idx >= gas_func_idx {
+								idx += 1;
-									idx += 1;
+							}
 								}
-								(idx, name.clone())
+							(idx, name.clone())
-							}));
+						}));
 					}
 				},
 			_ => {},
 		}
@@ -330,7 +233,7 @@ pub fn inject<R: Rules, B: Backend>(
 	}
 	if need_grow_counter {
-		Ok(add_grow_counter(module, rules, gas_func_idx))
+		Ok(add_grow_counter(module, rules, gas_func))
 	} else {
 		Ok(module)
 	}
@@ -582,7 +485,6 @@ 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::*;
@@ -590,9 +492,6 @@ 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];
@@ -659,19 +558,16 @@ 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, locals_count)?;
+	let blocks = determine_metered_blocks(instructions, rules)?;
-	insert_metering_calls(instructions, gas_function_cost, blocks, gas_func)
+	insert_metering_calls(instructions, 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,
 	gas_function_cost: u64,
 	blocks: Vec<MeteredBlock>,
 	gas_func: u32,
 ) -> Result<(), ()> {
@@ -689,8 +585,7 @@ 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
+				new_instrs.push(I64Const(block.cost as i64));
 					.push(I64Const((block.cost.checked_add(gas_function_cost).ok_or(())?) as i64));
 				new_instrs.push(Call(gas_func));
 				true
 			} else {
@@ -732,7 +627,7 @@ mod tests {
 	}
 	#[test]
-	fn simple_grow_host_fn() {
+	fn simple_grow() {
 		let module = parse_wat(
 			r#"(module
 			(func (result i32)
@@ -742,9 +637,8 @@ mod tests {
 			(memory 0 1)
 			)"#,
 		);
-		let backend = host_function::Injector::new("env", "gas");
+
-		let injected_module =
+		let injected_module = inject(module, &ConstantCostRules::new(1, 10_000), "env").unwrap();
 			super::inject(module, backend, &ConstantCostRules::new(1, 10_000, 1)).unwrap();
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
@@ -769,64 +663,7 @@ mod tests {
 	}
 	#[test]
-	fn simple_grow_mut_global() {
+	fn grow_no_gas_no_track() {
 		let module = parse_wat(
 			r#"(module
 			(func (result i32)
 			  global.get 0
 			  memory.grow)
 			(global i32 (i32.const 42))
 			(memory 0 1)
 			)"#,
 		);
 		let backend = mutable_global::Injector::new("gas_left");
 		let injected_module =
 			super::inject(module, backend, &ConstantCostRules::new(1, 10_000, 1)).unwrap();
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
 			&vec![I64Const(13), Call(1), GetGlobal(0), Call(2), End][..]
 		);
 		assert_eq!(
 			get_function_body(&injected_module, 1).unwrap(),
 			&vec![
 				Instruction::GetGlobal(1),
 				Instruction::GetLocal(0),
 				Instruction::I64GeU,
 				Instruction::If(elements::BlockType::NoResult),
 				Instruction::GetGlobal(1),
 				Instruction::GetLocal(0),
 				Instruction::I64Sub,
 				Instruction::SetGlobal(1),
 				Instruction::Else,
 				// sentinel val u64::MAX
 				Instruction::I64Const(-1i64), // non-charged instruction
 				Instruction::SetGlobal(1),    // non-charged instruction
 				Instruction::Unreachable,     // non-charged instruction
 				Instruction::End,
 				Instruction::End,
 			][..]
 		);
 		assert_eq!(
 			get_function_body(&injected_module, 2).unwrap(),
 			&vec![
 				GetLocal(0),
 				GetLocal(0),
 				I64ExtendUI32,
 				I64Const(10000),
 				I64Mul,
 				Call(1),
 				GrowMemory(0),
 				End,
 			][..]
 		);
 		let binary = serialize(injected_module).expect("serialization failed");
 		wasmparser::validate(&binary).unwrap();
 	}
 	#[test]
 	fn grow_no_gas_no_track_host_fn() {
 		let module = parse_wat(
 			r"(module
 			(func (result i32)
@@ -836,9 +673,8 @@ mod tests {
 			(memory 0 1)
 			)",
 		);
-		let backend = host_function::Injector::new("env", "gas");
+
-		let injected_module =
+		let injected_module = inject(module, &ConstantCostRules::default(), "env").unwrap();
 			super::inject(module, backend, &ConstantCostRules::default()).unwrap();
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
@@ -852,33 +688,7 @@ mod tests {
 	}
 	#[test]
-	fn grow_no_gas_no_track_mut_global() {
+	fn call_index() {
 		let module = parse_wat(
 			r"(module
 			(func (result i32)
 			  global.get 0
 			  memory.grow)
 			(global i32 (i32.const 42))
 			(memory 0 1)
 			)",
 		);
 		let backend = mutable_global::Injector::new("gas_left");
 		let injected_module =
 			super::inject(module, backend, &ConstantCostRules::default()).unwrap();
 		assert_eq!(
 			get_function_body(&injected_module, 0).unwrap(),
 			&vec![I64Const(13), Call(1), GetGlobal(0), GrowMemory(0), End][..]
 		);
 		assert_eq!(injected_module.functions_space(), 2);
 		let binary = serialize(injected_module).expect("serialization failed");
 		wasmparser::validate(&binary).unwrap();
 	}
 	#[test]
 	fn call_index_host_fn() {
 		let module = builder::module()
 			.global()
 			.value_type()
@@ -915,9 +725,7 @@ mod tests {
 			.build()
 			.build();
-		let backend = host_function::Injector::new("env", "gas");
+		let injected_module = inject(module, &ConstantCostRules::default(), "env").unwrap();
 		let injected_module =
 			super::inject(module, backend, &ConstantCostRules::default()).unwrap();
 		assert_eq!(
 			get_function_body(&injected_module, 1).unwrap(),
@@ -943,89 +751,20 @@ mod tests {
 		);
 	}
 	#[test]
 	fn call_index_mut_global() {
 		let module = builder::module()
 			.global()
 			.value_type()
 			.i32()
 			.build()
 			.function()
 			.signature()
 			.param()
 			.i32()
 			.build()
 			.body()
 			.build()
 			.build()
 			.function()
 			.signature()
 			.param()
 			.i32()
 			.build()
 			.body()
 			.with_instructions(elements::Instructions::new(vec![
 				Call(0),
 				If(elements::BlockType::NoResult),
 				Call(0),
 				Call(0),
 				Call(0),
 				Else,
 				Call(0),
 				Call(0),
 				End,
 				Call(0),
 				End,
 			]))
 			.build()
 			.build()
 			.build();
 		let backend = mutable_global::Injector::new("gas_left");
 		let injected_module =
 			super::inject(module, backend, &ConstantCostRules::default()).unwrap();
 		assert_eq!(
 			get_function_body(&injected_module, 1).unwrap(),
 			&vec![
 				I64Const(14),
 				Call(2),
 				Call(0),
 				If(elements::BlockType::NoResult),
 				I64Const(14),
 				Call(2),
 				Call(0),
 				Call(0),
 				Call(0),
 				Else,
 				I64Const(13),
 				Call(2),
 				Call(0),
 				Call(0),
 				End,
 				Call(0),
 				End
 			][..]
 		);
 	}
 	fn parse_wat(source: &str) -> elements::Module {
 		let module_bytes = wat::parse_str(source).unwrap();
 		elements::deserialize_buffer(module_bytes.as_ref()).unwrap()
 	}
 	macro_rules! test_gas_counter_injection {
-		(names = ($name1:ident, $name2:ident); input = $input:expr; expected = $expected:expr) => {
+		(name = $name:ident; input = $input:expr; expected = $expected:expr) => {
 			#[test]
-			fn $name1() {
+			fn $name() {
 				let input_module = parse_wat($input);
 				let expected_module = parse_wat($expected);
-				let injected_module = super::inject(
+
-					input_module,
+				let injected_module = inject(input_module, &ConstantCostRules::default(), "env")
-					host_function::Injector::new("env", "gas"),
+					.expect("inject_gas_counter call failed");
 					&ConstantCostRules::default(),
 				)
 				.expect("inject_gas_counter call failed");
 				let actual_func_body = get_function_body(&injected_module, 0)
 					.expect("injected module must have a function body");
@@ -1034,51 +773,11 @@ mod tests {
 				assert_eq!(actual_func_body, expected_func_body);
 			}
 			#[test]
 			fn $name2() {
 				let input_module = parse_wat($input);
 				let draft_module = parse_wat($expected);
 				let gas_fun_cost = match mutable_global::Injector::new("gas_left")
 					.gas_meter(&input_module, &ConstantCostRules::default())
 				{
 					GasMeter::Internal { cost, .. } => cost as i64,
 					_ => 0i64,
 				};
 				let injected_module = super::inject(
 					input_module,
 					mutable_global::Injector::new("gas_left"),
 					&ConstantCostRules::default(),
 				)
 				.expect("inject_gas_counter call failed");
 				let actual_func_body = get_function_body(&injected_module, 0)
 					.expect("injected module must have a function body");
 				let mut expected_func_body = get_function_body(&draft_module, 0)
 					.expect("post-module must have a function body")
 					.to_vec();
 				// modify expected instructions set for gas_metering::mutable_global
 				let mut iter = expected_func_body.iter_mut();
 				while let Some(ins) = iter.next() {
 					if let I64Const(cost) = ins {
 						if let Some(ins_next) = iter.next() {
 							if let Call(0) = ins_next {
 								*cost += gas_fun_cost;
 								*ins_next = Call(1);
 							}
 						}
 					}
 				}
 				assert_eq!(actual_func_body, &expected_func_body);
 			}
 		};
 	}
 	test_gas_counter_injection! {
-		names = (simple_host_fn, simple_mut_global);
+		name = simple;
 		input = r#"
 		(module
 			(func (result i32)
@@ -1093,7 +792,7 @@ mod tests {
 	}
 	test_gas_counter_injection! {
-		names = (nested_host_fn, nested_mut_global);
+		name = nested;
 		input = r#"
 		(module
 			(func (result i32)
@@ -1118,7 +817,7 @@ mod tests {
 	}
 	test_gas_counter_injection! {
-		names = (ifelse_host_fn, ifelse_mut_global);
+		name = ifelse;
 		input = r#"
 		(module
 			(func (result i32)
@@ -1153,7 +852,7 @@ mod tests {
 	}
 	test_gas_counter_injection! {
-		names = (branch_innermost_host_fn, branch_innermost_mut_global);
+		name = branch_innermost;
 		input = r#"
 		(module
 			(func (result i32)
@@ -1183,7 +882,7 @@ mod tests {
 	}
 	test_gas_counter_injection! {
-		names = (branch_outer_block_host_fn, branch_outer_block_mut_global);
+		name = branch_outer_block;
 		input = r#"
 		(module
 			(func (result i32)
@@ -1222,7 +921,7 @@ mod tests {
 	}
 	test_gas_counter_injection! {
-		names = (branch_outer_loop_host_fn, branch_outer_loop_mut_global);
+		name = branch_outer_loop;
 		input = r#"
 		(module
 			(func (result i32)
@@ -1268,7 +967,7 @@ mod tests {
 	}
 	test_gas_counter_injection! {
-		names = (return_from_func_host_fn, return_from_func_mut_global);
+		name = return_from_func;
 		input = r#"
 		(module
 			(func (result i32)
@@ -1293,7 +992,7 @@ mod tests {
 	}
 	test_gas_counter_injection! {
-		names = (branch_from_if_not_else_host_fn, branch_from_if_not_else_mut_global);
+		name = branch_from_if_not_else;
 		input = r#"
 		(module
 			(func (result i32)
@@ -1329,7 +1028,7 @@ mod tests {
 	}
 	test_gas_counter_injection! {
-		names = (empty_loop_host_fn, empty_loop_mut_global);
+		name = empty_loop;
 		input = r#"
 		(module
 			(func
@@ -134,10 +134,6 @@ 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()
@@ -153,10 +149,6 @@ 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(_) => {
@@ -350,11 +342,8 @@ 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 =
+				let metered_blocks = determine_metered_blocks(func_body.code(), &rules).unwrap();
 					determine_metered_blocks(func_body.code(), &rules, locals_count).unwrap();
 				let success =
 					validate_metering_injections(func_body, &rules, &metered_blocks).unwrap();
 				assert!(success);
@@ -1,10 +1,11 @@
 use super::resolve_func_type;
 use alloc::vec::Vec;
 use parity_wasm::elements::{self, BlockType, Type};
 #[cfg(feature = "sign_ext")]
 use parity_wasm::elements::SignExtInstruction;
 use super::Context;
 // 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
@@ -122,7 +123,11 @@ 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,
 	ctx: &Context,
 	module: &elements::Module,
 ) -> Result<u32, &'static str> {
 	use parity_wasm::elements::Instruction::*;
 	let func_section = module.function_section().ok_or("No function section")?;
@@ -246,8 +251,10 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 				stack.pop_values(func_arity)?;
 				stack.mark_unreachable()?;
 			},
-			Call(idx) => {
+			Call(fn_idx) => {
-				let ty = resolve_func_type(*idx, module)?;
+				let ty_idx = ctx.func_type(*fn_idx).ok_or("function idx is not found in the func types list")?;
 				let Type::Function(ty) =
 					type_section.types().get(ty_idx as usize).ok_or("Type not found")?;
 				// Pop values for arguments of the function.
 				stack.pop_values(ty.params().len() as u32)?;
@@ -411,7 +418,8 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 #[cfg(test)]
 mod tests {
-	use super::*;
+	use super::ACTIVATION_FRAME_COST;
 	use crate::stack_limiter::prepare_context;
 	use parity_wasm::elements;
 	fn parse_wat(source: &str) -> elements::Module {
@@ -419,10 +427,15 @@ mod tests {
 			.expect("Failed to deserialize the module")
 	}
 	fn test_compute(func_idx: u32, source: &str) -> u32 {
 		let module = parse_wat(source);
 		let ctx = prepare_context(&module, 0).unwrap();
 		ctx.stack_cost(func_idx).unwrap()
 	}
 	#[test]
 	fn simple_test() {
-		let module = parse_wat(
+		let module = r#"
 			r#"
 (module
 	(func
 		i32.const 1
@@ -433,34 +446,30 @@ mod tests {
 		drop
 	)
 )
-"#,
+"#;
 		);
-		let height = compute(0, &module).unwrap();
+		let height = test_compute(0, module);
 		assert_eq!(height, 3 + ACTIVATION_FRAME_COST);
 	}
 	#[test]
 	fn implicit_and_explicit_return() {
-		let module = parse_wat(
+		let module = r#"
 			r#"
 (module
 	(func (result i32)
 		i32.const 0
 		return
 	)
 )
-"#,
+"#;
 		);
-		let height = compute(0, &module).unwrap();
+		let height = test_compute(0, module);
 		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
 	}
 	#[test]
 	fn dont_count_in_unreachable() {
-		let module = parse_wat(
+		let module = r#"
 			r#"
 (module
  (memory 0)
  (func (result i32)
@@ -468,17 +477,15 @@ mod tests {
 	grow_memory
  )
 )
-"#,
+"#;
 		);
-		let height = compute(0, &module).unwrap();
+		let height = test_compute(0, module);
 		assert_eq!(height, ACTIVATION_FRAME_COST);
 	}
 	#[test]
 	fn yet_another_test() {
-		let module = parse_wat(
+		let module = r#"
 			r#"
 (module
  (memory 0)
  (func
@@ -497,17 +504,15 @@ mod tests {
 	i32.const 2
  )
 )
-"#,
+"#;
 		);
-		let height = compute(0, &module).unwrap();
+		let height = test_compute(0, module);
 		assert_eq!(height, 2 + ACTIVATION_FRAME_COST);
 	}
 	#[test]
 	fn call_indirect() {
-		let module = parse_wat(
+		let module = r#"
 			r#"
 (module
 	(table $ptr 1 1 funcref)
 	(elem $ptr (i32.const 0) func 1)
@@ -521,17 +526,15 @@ mod tests {
 		drop
 	)
 )
-"#,
+"#;
 		);
-		let height = compute(0, &module).unwrap();
+		let height = test_compute(0, module);
 		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
 	}
 	#[test]
 	fn breaks() {
-		let module = parse_wat(
+		let module = r#"
 			r#"
 (module
 	(func $main
 		block (result i32)
@@ -543,17 +546,15 @@ mod tests {
 		drop
 	)
 )
-"#,
+"#;
 		);
-		let height = compute(0, &module).unwrap();
+		let height = test_compute(0, module);
 		assert_eq!(height, 1 + ACTIVATION_FRAME_COST);
 	}
 	#[test]
 	fn if_else_works() {
-		let module = parse_wat(
+		let module = r#"
 			r#"
 (module
 	(func $main
 		i32.const 7
@@ -569,10 +570,9 @@ mod tests {
 		drop
 	)
 )
-"#,
+"#;
 		);
-		let height = compute(0, &module).unwrap();
+		let height = test_compute(0, module);
 		assert_eq!(height, 3 + ACTIVATION_FRAME_COST);
 	}
 }
@@ -39,7 +39,15 @@ mod max_height;
 mod thunk;
 pub struct Context {
 	/// Number of functions that the module imports. Required to convert defined functions indicies
 	/// into the global function index space.
 	func_imports: u32,
 	/// For each function in the function space this vector stores the respective type index.
 	func_types: Vec<u32>,
 	/// The index of the global variable that contains the current stack height.
 	stack_height_global_idx: u32,
 	/// Logical stack costs for each function in the function space. Imported functions have cost
 	/// of 0.
 	func_stack_costs: Vec<u32>,
 	stack_limit: u32,
 }
@@ -55,6 +63,11 @@ impl Context {
 		self.func_stack_costs.get(func_idx as usize).cloned()
 	}
 	/// Returns a reference to the function type index given by the index into the function space.
 	fn func_type(&self, func_idx: u32) -> Option<u32> {
 		self.func_types.get(func_idx as usize).copied()
 	}
 	/// Returns stack limit specified by the rules.
 	fn stack_limit(&self) -> u32 {
 		self.stack_limit
@@ -115,20 +128,106 @@ pub fn inject(
 	mut module: elements::Module,
 	stack_limit: u32,
 ) -> Result<elements::Module, &'static str> {
-	let mut ctx = Context {
+	let mut ctx = prepare_context(&module, stack_limit)?;
 		stack_height_global_idx: generate_stack_height_global(&mut module),
 		func_stack_costs: compute_stack_costs(&module)?,
 		stack_limit,
 	};
-	instrument_functions(&mut ctx, &mut module)?;
+	generate_stack_height_global(&mut ctx.stack_height_global_idx, &mut module)?;
 	instrument_functions(&ctx, &mut module)?;
 	let module = thunk::generate_thunks(&mut ctx, module)?;
 	Ok(module)
 }
 fn prepare_context(module: &elements::Module, stack_limit: u32) -> Result<Context, &'static str> {
 	let mut ctx = Context {
 		func_imports: module.import_count(elements::ImportCountType::Function) as u32,
 		func_types: Vec::new(),
 		stack_height_global_idx: 0,
 		func_stack_costs: Vec::new(),
 		stack_limit,
 	};
 	collect_func_types(&mut ctx, &module)?;
 	compute_stack_costs(&mut ctx, &module)?;
 	Ok(ctx)
 }
 fn collect_func_types(ctx: &mut Context, module: &elements::Module) -> Result<(), &'static str> {
 	let types = module.type_section().map(|ts| ts.types()).unwrap_or(&[]);
 	let functions = module.function_section().map(|fs| fs.entries()).unwrap_or(&[]);
 	let imports = module.import_section().map(|is| is.entries()).unwrap_or(&[]);
 	let ensure_ty = |sig_idx: u32| -> Result<(), &'static str> {
 		let Type::Function(_) = types
 			.get(sig_idx as usize)
 			.ok_or("The signature as specified by a function isn't defined")?;
 		Ok(())
 	};
 	for import in imports {
 		if let elements::External::Function(sig_idx) = import.external() {
 			ensure_ty(*sig_idx)?;
 			ctx.func_types.push(*sig_idx);
 		}
 	}
 	for def_func_idx in functions {
 		ensure_ty(def_func_idx.type_ref())?;
 		ctx.func_types.push(def_func_idx.type_ref());
 	}
 	Ok(())
 }
 /// Calculate stack costs for all functions in the function space.
 ///
 /// The function space consists of the imported functions followed by defined functions.
 /// All imported functions assumed to have the cost of 0.
 fn compute_stack_costs(ctx: &mut Context, module: &elements::Module) -> Result<(), &'static str> {
 	for _ in 0..ctx.func_imports {
 		ctx.func_stack_costs.push(0);
 	}
 	let def_func_n = module.function_section().map(|fs| fs.entries().len()).unwrap_or(0) as u32;
 	for def_func_idx in 0..def_func_n {
 		let cost = compute_stack_cost(def_func_idx, ctx, module)?;
 		ctx.func_stack_costs.push(cost);
 	}
 	Ok(())
 }
 /// Computes the stack cost of a given function. The function is specified by its index in the
 /// declared function space.
 ///
 /// Stack cost of a given 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(
 	def_func_idx: u32,
 	ctx: &Context,
 	module: &elements::Module,
 ) -> Result<u32, &'static str> {
 	let code_section =
 		module.code_section().ok_or("Due to validation code section should exists")?;
 	let body = &code_section
 		.bodies()
 		.get(def_func_idx as usize)
 		.ok_or("Function body is out of bounds")?;
 	let mut locals_count: u32 = 0;
 	for local_group in body.locals() {
 		locals_count =
 			locals_count.checked_add(local_group.count()).ok_or("Overflow in local count")?;
 	}
 	let max_stack_height = max_height::compute(def_func_idx, ctx, module)?;
 	locals_count
 		.checked_add(max_stack_height)
 		.ok_or("Overflow in adding locals_count and max_stack_height")
 }
 /// Generate a new global that will be used for tracking current stack height.
-fn generate_stack_height_global(module: &mut elements::Module) -> u32 {
+fn generate_stack_height_global(
 	stack_height_global_idx: &mut u32,
 	module: &mut elements::Module,
 ) -> Result<(), &'static str> {
 	let global_entry = builder::global()
 		.value_type()
 		.i32()
@@ -140,71 +239,41 @@ fn generate_stack_height_global(module: &mut elements::Module) -> u32 {
 	for section in module.sections_mut() {
 		if let elements::Section::Global(gs) = section {
 			gs.entries_mut().push(global_entry);
-			return (gs.entries().len() as u32) - 1
+			*stack_height_global_idx = (gs.entries().len() as u32) - 1;
 			return Ok(());
 		}
 	}
 	// Existing section not found, create one!
-	module
+	//
-		.sections_mut()
+	// It's a bit tricky since the sections have a strict prescribed order.
-		.push(elements::Section::Global(elements::GlobalSection::with_entries(vec![global_entry])));
+	let global_section = elements::GlobalSection::with_entries(vec![global_entry]);
-	0
+	let prec_index = module
-}
+		.sections()
-
+		.iter()
-/// Calculate stack costs for all functions.
+		.rposition(|section| {
-///
+			use elements::Section::*;
-/// Returns a vector with a stack cost for each function, including imports.
+			match section {
-fn compute_stack_costs(module: &elements::Module) -> Result<Vec<u32>, &'static str> {
+				Type(_) | Import(_) | Function(_) | Table(_) | Memory(_) => true,
-	let func_imports = module.import_count(elements::ImportCountType::Function);
+				_ => false,
 	// TODO: optimize!
 	(0..module.functions_space())
 		.map(|func_idx| {
 			if func_idx < func_imports {
 				// We can't calculate stack_cost of the import functions.
 				Ok(0)
 			} else {
 				compute_stack_cost(func_idx as u32, module)
 			}
 		})
-		.collect()
+		.ok_or("generate stack height global hasn't found any preceding sections")?;
 	// now `prec_index` points to the last section preceding the `global_section`. It's guaranteed that at least
 	// one of those functions is present. Therefore, the candidate position for the global section is the following
 	// one. However, technically, custom sections could occupy any place between the well-known sections.
 	//
 	// Now, regarding `+1` here. `insert` panics iff `index > len`. `prec_index + 1` can only be equal to `len`.
 	module
 		.sections_mut()
 		.insert(prec_index + 1, elements::Section::Global(global_section));
 	// First entry in the brand new globals section.
 	*stack_height_global_idx = 0;
 	Ok(())
 }
-/// Stack cost of the given *defined* function is the sum of it's locals count (that is,
+fn instrument_functions(ctx: &Context, module: &mut elements::Module) -> Result<(), &'static str> {
 /// 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> {
 	// 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;
 	let defined_func_idx = func_idx
 		.checked_sub(func_imports)
 		.ok_or("This should be a index of a defined function")?;
 	let code_section =
 		module.code_section().ok_or("Due to validation code section should exists")?;
 	let body = &code_section
 		.bodies()
 		.get(defined_func_idx as usize)
 		.ok_or("Function body is out of bounds")?;
 	let mut locals_count: u32 = 0;
 	for local_group in body.locals() {
 		locals_count =
 			locals_count.checked_add(local_group.count()).ok_or("Overflow in local count")?;
 	}
 	let max_stack_height = 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")
 }
 fn instrument_functions(
 	ctx: &mut Context,
 	module: &mut elements::Module,
 ) -> Result<(), &'static str> {
 	for section in module.sections_mut() {
 		if let elements::Section::Code(code_section) = section {
 			for func_body in code_section.bodies_mut() {
@@ -242,7 +311,7 @@ fn instrument_functions(
 ///
 /// drop
 /// ```
-fn instrument_function(ctx: &mut Context, func: &mut Instructions) -> Result<(), &'static str> {
+fn instrument_function(ctx: &Context, func: &mut Instructions) -> Result<(), &'static str> {
 	use Instruction::*;
 	struct InstrumentCall {
@@ -309,42 +378,6 @@ fn instrument_function(ctx: &mut Context, func: &mut Instructions) -> Result<(),
 	Ok(())
 }
 fn resolve_func_type(
 	func_idx: u32,
 	module: &elements::Module,
 ) -> Result<&elements::FunctionType, &'static str> {
 	let types = module.type_section().map(|ts| ts.types()).unwrap_or(&[]);
 	let functions = module.function_section().map(|fs| fs.entries()).unwrap_or(&[]);
 	let func_imports = module.import_count(elements::ImportCountType::Function);
 	let sig_idx = if func_idx < func_imports as u32 {
 		module
 			.import_section()
 			.expect("function import count is not zero; import section must exists; qed")
 			.entries()
 			.iter()
 			.filter_map(|entry| match entry.external() {
 				elements::External::Function(idx) => Some(*idx),
 				_ => None,
 			})
 			.nth(func_idx as usize)
 			.expect(
 				"func_idx is less than function imports count;
 				nth function import must be `Some`;
 				qed",
 			)
 	} else {
 		functions
 			.get(func_idx as usize - func_imports)
 			.ok_or("Function at the specified index is not defined")?
 			.type_ref()
 	};
 	let Type::Function(ty) = types
 		.get(sig_idx as usize)
 		.ok_or("The signature as specified by a function isn't defined")?;
 	Ok(ty)
 }
 #[cfg(test)]
 mod tests {
 	use super::*;
@@ -1,17 +1,17 @@
 #[cfg(not(features = "std"))]
 use alloc::collections::BTreeMap as Map;
 use alloc::vec::Vec;
-use parity_wasm::{
+use parity_wasm::elements::{self, Internal};
 	builder,
 	elements::{self, FunctionType, Internal},
 };
 #[cfg(features = "std")]
 use std::collections::HashMap as Map;
-use super::{resolve_func_type, Context};
+use super::Context;
 struct Thunk {
-	signature: FunctionType,
+	/// The index of the signature in the type section.
 	type_idx: u32,
 	/// The number of parameters the function has.
 	param_num: u32,
 	// Index in function space of this thunk.
 	idx: Option<u32>,
 	callee_stack_cost: u32,
@@ -19,10 +19,11 @@ struct Thunk {
 pub fn generate_thunks(
 	ctx: &mut Context,
-	module: elements::Module,
+	mut module: elements::Module,
 ) -> Result<elements::Module, &'static str> {
 	// First, we need to collect all function indices that should be replaced by thunks
 	let mut replacement_map: Map<u32, Thunk> = {
 		let types = module.type_section().map(|ts| ts.types()).unwrap_or(&[]);
 		let exports = module.export_section().map(|es| es.entries()).unwrap_or(&[]);
 		let elem_segments = module.elements_section().map(|es| es.entries()).unwrap_or(&[]);
 		let start_func_idx = module.start_section();
@@ -45,14 +46,12 @@ pub fn generate_thunks(
 			// Don't generate a thunk if stack_cost of a callee is zero.
 			if callee_stack_cost != 0 {
-				replacement_map.insert(
+				let type_idx = ctx.func_type(func_idx).ok_or("type idx for thunk not found")?;
-					func_idx,
+				let elements::Type::Function(func_ty) =
-					Thunk {
+					types.get(type_idx as usize).ok_or("sig for thunk is not found")?;
-						signature: resolve_func_type(func_idx, &module)?.clone(),
+				let param_num = func_ty.params().len() as u32;
-						idx: None,
+				replacement_map
-						callee_stack_cost,
+					.insert(func_idx, Thunk { type_idx, param_num, idx: None, callee_stack_cost });
 					},
 				);
 			}
 		}
@@ -61,10 +60,6 @@ pub fn generate_thunks(
 	// Then, we generate a thunk for each original function.
 	// Save current func_idx
 	let mut next_func_idx = module.functions_space() as u32;
 	let mut mbuilder = builder::from_module(module);
 	for (func_idx, thunk) in replacement_map.iter_mut() {
 		let instrumented_call = instrument_call!(
 			*func_idx,
@@ -77,32 +72,23 @@ pub fn generate_thunks(
 		//  - instrumented call
 		//  - end
 		let mut thunk_body: Vec<elements::Instruction> =
-			Vec::with_capacity(thunk.signature.params().len() + instrumented_call.len() + 1);
+			Vec::with_capacity(thunk.param_num as usize + instrumented_call.len() + 1);
-		for (arg_idx, _) in thunk.signature.params().iter().enumerate() {
+		for arg_idx in 0..thunk.param_num {
-			thunk_body.push(elements::Instruction::GetLocal(arg_idx as u32));
+			thunk_body.push(elements::Instruction::GetLocal(arg_idx));
 		}
 		thunk_body.extend_from_slice(&instrumented_call);
 		thunk_body.push(elements::Instruction::End);
-		// TODO: Don't generate a signature, but find an existing one.
+		let func_idx = insert_function(
-
+			ctx,
-		mbuilder = mbuilder
+			&mut module,
-			.function()
+			thunk.type_idx,
-			// Signature of the thunk should match the original function signature.
+			Vec::new(), // No declared local variables.
-			.signature()
+			elements::Instructions::new(thunk_body),
-			.with_params(thunk.signature.params().to_vec())
+		)?;
-			.with_results(thunk.signature.results().to_vec())
+		thunk.idx = Some(func_idx);
 			.build()
 			.body()
 			.with_instructions(elements::Instructions::new(thunk_body))
 			.build()
 			.build();
 		thunk.idx = Some(next_func_idx);
 		next_func_idx += 1;
 	}
 	let mut module = mbuilder.build();
 	// And finally, fixup thunks in export and table sections.
@@ -118,18 +104,20 @@ pub fn generate_thunks(
 	for section in module.sections_mut() {
 		match section {
-			elements::Section::Export(export_section) =>
+			elements::Section::Export(export_section) => {
 				for entry in export_section.entries_mut() {
 					if let Internal::Function(function_idx) = entry.internal_mut() {
 						fixup(function_idx)
 					}
-				},
+				}
-			elements::Section::Element(elem_section) =>
+			},
 			elements::Section::Element(elem_section) => {
 				for segment in elem_section.entries_mut() {
 					for function_idx in segment.members_mut() {
 						fixup(function_idx)
 					}
-				},
+				}
 			},
 			elements::Section::Start(start_idx) => fixup(start_idx),
 			_ => {},
 		}
@@ -137,3 +125,30 @@ pub fn generate_thunks(
 	Ok(module)
 }
 /// Inserts a new function into the module and returns it's index in the function space.
 ///
 /// Specifically, inserts entires into the function section and the code section.
 fn insert_function(
 	ctx: &Context,
 	module: &mut elements::Module,
 	type_idx: u32,
 	locals: Vec<elements::Local>,
 	insns: elements::Instructions,
 ) -> Result<u32, &'static str> {
 	let funcs = module
 		.function_section_mut()
 		.ok_or("insert function no function section")?
 		.entries_mut();
 	let new_func_idx = ctx
 		.func_imports
 		.checked_add(funcs.len() as u32)
 		.ok_or("insert function func idx overflow")?;
 	funcs.push(elements::Func::new(type_idx));
 	let func_bodies =
 		module.code_section_mut().ok_or("insert function no code section")?.bodies_mut();
 	func_bodies.push(elements::FuncBody::new(locals, insns));
 	Ok(new_func_idx)
 }
@@ -1,9 +1,10 @@
 use parity_wasm::elements::Module;
 use std::{
 	fs,
 	io::{self, Read, Write},
 	path::{Path, PathBuf},
 };
-use wasm_instrument::{self as instrument, gas_metering, parity_wasm::elements};
+use wasm_instrument::{self as instrument, parity_wasm::elements};
 use wasmparser::validate;
 fn slurp<P: AsRef<Path>>(path: P) -> io::Result<Vec<u8>> {
@@ -19,23 +20,18 @@ fn dump<P: AsRef<Path>>(path: P, buf: &[u8]) -> io::Result<()> {
 	Ok(())
 }
-fn run_diff_test<F: FnOnce(&[u8]) -> Vec<u8>>(
+fn run_diff_test<F: FnOnce(&[u8]) -> Vec<u8>>(test_dir: &str, name: &str, test: F) {
 	test_dir: &str,
 	in_name: &str,
 	out_name: &str,
 	test: F,
 ) {
 	let mut fixture_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
 	fixture_path.push("tests");
 	fixture_path.push("fixtures");
 	fixture_path.push(test_dir);
-	fixture_path.push(in_name);
+	fixture_path.push(name);
 	let mut expected_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
 	expected_path.push("tests");
 	expected_path.push("expectations");
 	expected_path.push(test_dir);
-	expected_path.push(out_name);
+	expected_path.push(name);
 	let fixture_wasm = wat::parse_file(&fixture_path).expect("Failed to read fixture");
 	validate(&fixture_wasm).expect("Fixture is invalid");
@@ -52,7 +48,7 @@ fn run_diff_test<F: FnOnce(&[u8]) -> Vec<u8>>(
 	if actual_wat != expected_wat {
 		println!("difference!");
 		println!("--- {}", expected_path.display());
-		println!("+++ {} test {}", test_dir, out_name);
+		println!("+++ {} test {}", test_dir, name);
 		for diff in diff::lines(expected_wat, &actual_wat) {
 			match diff {
 				diff::Result::Left(l) => println!("-{}", l),
@@ -76,18 +72,13 @@ mod stack_height {
 		( $name:ident ) => {
 			#[test]
 			fn $name() {
-				run_diff_test(
+				run_diff_test("stack-height", concat!(stringify!($name), ".wat"), |input| {
-					"stack-height",
+					let module =
-					concat!(stringify!($name), ".wat"),
+						elements::deserialize_buffer(input).expect("Failed to deserialize");
-					concat!(stringify!($name), ".wat"),
+					let instrumented = instrument::inject_stack_limiter(module, 1024)
-					|input| {
+						.expect("Failed to instrument with stack counter");
-						let module =
+					elements::serialize(instrumented).expect("Failed to serialize")
-							elements::deserialize_buffer(input).expect("Failed to deserialize");
+				});
 						let instrumented = instrument::inject_stack_limiter(module, 1024)
 							.expect("Failed to instrument with stack counter");
 						elements::serialize(instrumented).expect("Failed to serialize")
 					},
 				);
 			}
 		};
 	}
@@ -105,53 +96,27 @@ mod gas {
 	use super::*;
 	macro_rules! def_gas_test {
-		( ($input:ident, $name1:ident, $name2:ident) ) => {
+		( $name:ident ) => {
 			#[test]
-			fn $name1() {
+			fn $name() {
-				run_diff_test(
+				run_diff_test("gas", concat!(stringify!($name), ".wat"), |input| {
-					"gas",
+					let rules = instrument::gas_metering::ConstantCostRules::default();
 					concat!(stringify!($input), ".wat"),
 					concat!(stringify!($name1), ".wat"),
 					|input| {
 						let rules = gas_metering::ConstantCostRules::default();
-						let module: elements::Module =
+					let module: Module =
-							elements::deserialize_buffer(input).expect("Failed to deserialize");
+						elements::deserialize_buffer(input).expect("Failed to deserialize");
-						let module = module.parse_names().expect("Failed to parse names");
+					let module = module.parse_names().expect("Failed to parse names");
 						let backend = gas_metering::host_function::Injector::new("env", "gas");
-						let instrumented = gas_metering::inject(module, backend, &rules)
+					let instrumented = instrument::gas_metering::inject(module, &rules, "env")
-							.expect("Failed to instrument with gas metering");
+						.expect("Failed to instrument with gas metering");
-						elements::serialize(instrumented).expect("Failed to serialize")
+					elements::serialize(instrumented).expect("Failed to serialize")
-					},
+				});
 				);
 			}
 			#[test]
 			fn $name2() {
 				run_diff_test(
 					"gas",
 					concat!(stringify!($input), ".wat"),
 					concat!(stringify!($name2), ".wat"),
 					|input| {
 						let rules = gas_metering::ConstantCostRules::default();
 						let module: elements::Module =
 							elements::deserialize_buffer(input).expect("Failed to deserialize");
 						let module = module.parse_names().expect("Failed to parse names");
 						let backend = gas_metering::mutable_global::Injector::new("gas_left");
 						let instrumented = gas_metering::inject(module, backend, &rules)
 							.expect("Failed to instrument with gas metering");
 						elements::serialize(instrumented).expect("Failed to serialize")
 					},
 				);
 			}
 		};
 	}
-	def_gas_test!((ifs, ifs_host_fn, ifs_mut_global));
+	def_gas_test!(ifs);
-	def_gas_test!((simple, simple_host_fn, simple_mut_global));
+	def_gas_test!(simple);
-	def_gas_test!((start, start_host_fn, start_mut_global));
+	def_gas_test!(start);
-	def_gas_test!((call, call_host_fn, call_mut_global));
+	def_gas_test!(call);
-	def_gas_test!((branch, branch_host_fn, branch_mut_global));
+	def_gas_test!(branch);
 }
@@ -4,9 +4,9 @@
  (import "env" "gas" (func (;0;) (type 1)))
  (func $fibonacci_with_break (;1;) (type 0) (result i32)
    (local i32 i32)
-    i64.const 15
+    i64.const 13
    call 0
-    block ;; label = @1
+    block  ;; label = @1
      i32.const 0
      local.set 0
      i32.const 1
@@ -1,47 +0,0 @@
 (module
  (type (;0;) (func (result i32)))
  (type (;1;) (func (param i64)))
  (func $fibonacci_with_break (;0;) (type 0) (result i32)
    (local $x i32) (local $y i32)
    i64.const 26
    call 1
    block ;; label = @1
      i32.const 0
      local.set $x
      i32.const 1
      local.set $y
      local.get $x
      local.get $y
      local.tee $x
      i32.add
      local.set $y
      i32.const 1
      br_if 0 (;@1;)
      i64.const 16
      call 1
      local.get $x
      local.get $y
      local.tee $x
      i32.add
      local.set $y
    end
    local.get $y
  )
  (func (;1;) (type 1) (param i64)
    global.get 0
    local.get 0
    i64.ge_u
    if ;; label = @1
      global.get 0
      local.get 0
      i64.sub
      global.set 0
    else
      i64.const -1
      global.set 0
      unreachable
    end
  )
  (global (;0;) (mut i64) i64.const 0)
  (export "gas_left" (global 0))
 )
@@ -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 6
+    i64.const 5
    call 0
    local.get $x
    local.get $y
@@ -1,38 +0,0 @@
 (module
  (type (;0;) (func (param i32 i32) (result i32)))
  (type (;1;) (func (param i64)))
  (func $add_locals (;0;) (type 0) (param $x i32) (param $y i32) (result i32)
    (local $t i32)
    i64.const 17
    call 2
    local.get $x
    local.get $y
    call $add
    local.set $t
    local.get $t
  )
  (func $add (;1;) (type 0) (param $x i32) (param $y i32) (result i32)
    i64.const 14
    call 2
    local.get $x
    local.get $y
    i32.add
  )
  (func (;2;) (type 1) (param i64)
    global.get 0
    local.get 0
    i64.ge_u
    if ;; label = @1
      global.get 0
      local.get 0
      i64.sub
      global.set 0
    else
      i64.const -1
      global.set 0
      unreachable
    end
  )
  (global (;0;) (mut i64) i64.const 0)
  (export "gas_left" (global 0))
 )
@@ -6,7 +6,7 @@
    i64.const 2
    call 0
    i32.const 1
-    if (result i32) ;; label = @1
+    if (result i32)  ;; label = @1
      i64.const 3
      call 0
      local.get 0
@@ -1,38 +0,0 @@
 (module
  (type (;0;) (func (param i32) (result i32)))
  (type (;1;) (func (param i64)))
  (func (;0;) (type 0) (param $x i32) (result i32)
    i64.const 13
    call 1
    i32.const 1
    if (result i32) ;; label = @1
      i64.const 14
      call 1
      local.get $x
      i32.const 1
      i32.add
    else
      i64.const 13
      call 1
      local.get $x
      i32.popcnt
    end
  )
  (func (;1;) (type 1) (param i64)
    global.get 0
    local.get 0
    i64.ge_u
    if ;; label = @1
      global.get 0
      local.get 0
      i64.sub
      global.set 0
    else
      i64.const -1
      global.set 0
      unreachable
    end
  )
  (global (;0;) (mut i64) i64.const 0)
  (export "gas_left" (global 0))
 )
@@ -6,10 +6,10 @@
    i64.const 2
    call 0
    i32.const 1
-    if ;; label = @1
+    if  ;; label = @1
      i64.const 1
      call 0
-      loop ;; label = @2
+      loop  ;; label = @2
        i64.const 2
        call 0
        i32.const 123
@@ -20,7 +20,7 @@
  (func (;2;) (type 0)
    i64.const 1
    call 0
-    block ;; label = @1
+    block  ;; label = @1
    end
  )
  (export "simple" (func 1))
@@ -1,43 +0,0 @@
 (module
  (type (;0;) (func))
  (type (;1;) (func (param i64)))
  (func (;0;) (type 0)
    i64.const 13
    call 2
    i32.const 1
    if ;; label = @1
      i64.const 12
      call 2
      loop ;; label = @2
        i64.const 13
        call 2
        i32.const 123
        drop
      end
    end
  )
  (func (;1;) (type 0)
    i64.const 12
    call 2
    block ;; label = @1
    end
  )
  (func (;2;) (type 1) (param i64)
    global.get 0
    local.get 0
    i64.ge_u
    if ;; label = @1
      global.get 0
      local.get 0
      i64.sub
      global.set 0
    else
      i64.const -1
      global.set 0
      unreachable
    end
  )
  (global (;0;) (mut i64) i64.const 0)
  (export "simple" (func 0))
  (export "gas_left" (global 0))
 )
@@ -1,36 +0,0 @@
 (module
  (type (;0;) (func (param i32 i32)))
  (type (;1;) (func))
  (type (;2;) (func (param i64)))
  (import "env" "ext_return" (func $ext_return (;0;) (type 0)))
  (import "env" "memory" (memory (;0;) 1 1))
  (func $start (;1;) (type 1)
    i64.const 15
    call 3
    i32.const 8
    i32.const 4
    call $ext_return
    unreachable
  )
  (func (;2;) (type 1))
  (func (;3;) (type 2) (param i64)
    global.get 0
    local.get 0
    i64.ge_u
    if ;; label = @1
      global.get 0
      local.get 0
      i64.sub
      global.set 0
    else
      i64.const -1
      global.set 0
      unreachable
    end
  )
  (global (;0;) (mut i64) i64.const 0)
  (export "call" (func 2))
  (export "gas_left" (global 0))
  (start $start)
  (data (;0;) (i32.const 8) "\01\02\03\04")
 )
@@ -8,7 +8,7 @@
    global.get 0
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call 0
@@ -25,7 +25,7 @@
    global.get 0
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call 0
@@ -42,7 +42,7 @@
    global.get 0
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call 1
@@ -24,7 +24,7 @@
    global.get 1
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call $i32.add
@@ -44,7 +44,7 @@
    global.get 1
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call $i32.add
@@ -20,7 +20,7 @@
    global.get 0
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call 2
@@ -11,7 +11,7 @@
    global.get 0
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call $one-group-many-locals
@@ -12,7 +12,7 @@
    global.get 0
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call 0
@@ -15,7 +15,7 @@
    global.get 0
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call $start
@@ -32,7 +32,7 @@
    global.get 0
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call 2
@@ -13,7 +13,7 @@
    global.get 0
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call $i32.add
@@ -37,7 +37,7 @@
    global.get 0
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call 1
@@ -56,7 +56,7 @@
    global.get 0
    i32.const 1024
    i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
      unreachable
    end
    call $i32.add
@@ -1,10 +1,9 @@
 use std::{
-	fs::{read, read_dir, ReadDir},
+	fs::{read, read_dir},
 	path::PathBuf,
 };
 use wasm_instrument::{
-	gas_metering::{self, host_function, mutable_global, ConstantCostRules},
+	gas_metering, inject_stack_limiter,
 	inject_stack_limiter,
 	parity_wasm::{deserialize_buffer, elements::Module, serialize},
 };
@@ -15,157 +14,60 @@ fn fixture_dir() -> PathBuf {
 	path
 }
-use gas_metering::Backend;
+/// Print the overhead of applying gas metering, stack height limiting or both.
-fn gas_metered_mod_len<B: Backend>(orig_module: Module, backend: B) -> (Module, usize) {
+///
-	let module = gas_metering::inject(orig_module, backend, &ConstantCostRules::default()).unwrap();
+/// Use `cargo test print_overhead -- --nocapture`.
-	let bytes = serialize(module.clone()).unwrap();
+#[test]
-	let len = bytes.len();
+fn print_size_overhead() {
-	(module, len)
+	let mut results: Vec<_> = read_dir(fixture_dir())
-}
+		.unwrap()
 fn stack_limited_mod_len(module: Module) -> (Module, usize) {
 	let module = inject_stack_limiter(module, 128).unwrap();
 	let bytes = serialize(module.clone()).unwrap();
 	let len = bytes.len();
 	(module, len)
 }
 struct InstrumentedWasmResults {
 	filename: String,
 	original_module_len: usize,
 	stack_limited_len: usize,
 	gas_metered_host_fn_len: usize,
 	gas_metered_mut_glob_len: usize,
 	gas_metered_host_fn_then_stack_limited_len: usize,
 	gas_metered_mut_glob_then_stack_limited_len: usize,
 }
 fn size_overheads_all(files: ReadDir) -> Vec<InstrumentedWasmResults> {
 	files
 		.map(|entry| {
 			let entry = entry.unwrap();
-			let filename = entry.file_name().into_string().unwrap();
+			let (orig_len, orig_module) = {
-
+				let bytes = read(&entry.path()).unwrap();
 			let (original_module_len, orig_module) = {
 				let bytes = match entry.path().extension().unwrap().to_str() {
 					Some("wasm") => read(&entry.path()).unwrap(),
 					Some("wat") =>
 						wat::parse_bytes(&read(&entry.path()).unwrap()).unwrap().into_owned(),
 					_ => panic!("expected fixture_dir containing .wasm or .wat files only"),
 				};
 				let len = bytes.len();
 				let module: Module = deserialize_buffer(&bytes).unwrap();
 				(len, module)
 			};
 			let (gas_metering_len, gas_module) = {
 				let module = gas_metering::inject(
 					orig_module.clone(),
 					&gas_metering::ConstantCostRules::default(),
 					"env",
 				)
 				.unwrap();
 				let bytes = serialize(module.clone()).unwrap();
 				let len = bytes.len();
 				(len, module)
 			};
 			let stack_height_len = {
 				let module = inject_stack_limiter(orig_module, 128).unwrap();
 				let bytes = serialize(module).unwrap();
 				bytes.len()
 			};
 			let both_len = {
 				let module = inject_stack_limiter(gas_module, 128).unwrap();
 				let bytes = serialize(module).unwrap();
 				bytes.len()
 			};
-			let (gm_host_fn_module, gas_metered_host_fn_len) = gas_metered_mod_len(
+			let overhead = both_len * 100 / orig_len;
 				orig_module.clone(),
 				host_function::Injector::new("env", "gas"),
 			);
-			let (gm_mut_global_module, gas_metered_mut_glob_len) =
+			(
-				gas_metered_mod_len(orig_module.clone(), mutable_global::Injector::new("gas_left"));
+				overhead,
-
+				format!(
-			let stack_limited_len = stack_limited_mod_len(orig_module).1;
+					"{:30}: orig = {:4} kb, gas_metering = {} %, stack_limiter = {} %, both = {} %",
-
+					entry.file_name().to_str().unwrap(),
-			let (_gm_hf_sl_mod, gas_metered_host_fn_then_stack_limited_len) =
+					orig_len / 1024,
-				stack_limited_mod_len(gm_host_fn_module);
+					gas_metering_len * 100 / orig_len,
-
+					stack_height_len * 100 / orig_len,
-			let (_gm_mg_sl_module, gas_metered_mut_glob_then_stack_limited_len) =
+					overhead,
-				stack_limited_mod_len(gm_mut_global_module);
+				),
-
+			)
 			InstrumentedWasmResults {
 				filename,
 				original_module_len,
 				stack_limited_len,
 				gas_metered_host_fn_len,
 				gas_metered_mut_glob_len,
 				gas_metered_host_fn_then_stack_limited_len,
 				gas_metered_mut_glob_then_stack_limited_len,
 			}
 		})
-		.collect()
+		.collect();
 }
 fn calc_size_overheads() -> Vec<InstrumentedWasmResults> {
 	let mut wasm_path = fixture_dir();
 	wasm_path.push("wasm");
 	let mut wat_path = fixture_dir();
 	wat_path.push("wat");
 	let mut results = size_overheads_all(read_dir(wasm_path).unwrap());
 	let results_wat = size_overheads_all(read_dir(wat_path).unwrap());
 	results.extend(results_wat);
 	results
 }
 /// Print the overhead of applying gas metering, stack
 /// height limiting or both.
 ///
 /// Use `cargo test print_size_overhead -- --nocapture`.
 #[test]
 fn print_size_overhead() {
 	let mut results = calc_size_overheads();
 	results.sort_unstable_by(|a, b| {
 		b.gas_metered_mut_glob_then_stack_limited_len
 			.cmp(&a.gas_metered_mut_glob_then_stack_limited_len)
 	});
 	for r in results {
 		let filename = r.filename;
 		let original_size = r.original_module_len / 1024;
 		let stack_limit = r.stack_limited_len * 100 / r.original_module_len;
 		let host_fn = r.gas_metered_host_fn_len * 100 / r.original_module_len;
 		let mut_glob = r.gas_metered_mut_glob_len * 100 / r.original_module_len;
 		let host_fn_sl = r.gas_metered_host_fn_then_stack_limited_len * 100 / r.original_module_len;
 		let mut_glob_sl =
 			r.gas_metered_mut_glob_then_stack_limited_len * 100 / r.original_module_len;
 		println!(
 			"{filename:30}: orig = {original_size:4} kb, stack_limiter = {stack_limit} %, \
 			  gas_metered_host_fn =    {host_fn} %, both = {host_fn_sl} %,\n \
 			 {:69} gas_metered_mut_global = {mut_glob} %, both = {mut_glob_sl} %",
 			""
 		);
 	}
 }
 /// Compare module size overhead of applying gas metering with two methods.
 ///
 /// Use `cargo test print_gas_metered_sizes -- --nocapture`.
 #[test]
 fn print_gas_metered_sizes() {
 	let overheads = calc_size_overheads();
 	let mut results = overheads
 		.iter()
 		.map(|r| {
 			let diff = (r.gas_metered_mut_glob_len * 100 / r.gas_metered_host_fn_len) as i32 - 100;
 			(diff, r)
 		})
 		.collect::<Vec<(i32, &InstrumentedWasmResults)>>();
 	results.sort_unstable_by(|a, b| b.0.cmp(&a.0));
-
+	for entry in results {
-	println!(
+		println!("{}", entry.1);
 		"| {:28} | {:^16} | gas metered/host fn | gas metered/mut global | size diff |",
 		"fixture", "original size",
 	);
 	println!("|{:-^30}|{:-^18}|{:-^21}|{:-^24}|{:-^11}|", "", "", "", "", "",);
 	for r in results {
 		let filename = &r.1.filename;
 		let original_size = &r.1.original_module_len / 1024;
 		let host_fn = &r.1.gas_metered_host_fn_len / 1024;
 		let mut_glob = &r.1.gas_metered_mut_glob_len / 1024;
 		let host_fn_percent = &r.1.gas_metered_host_fn_len * 100 / r.1.original_module_len;
 		let mut_glob_percent = &r.1.gas_metered_mut_glob_len * 100 / r.1.original_module_len;
 		let host_fn = format!("{host_fn} kb ({host_fn_percent:}%)");
 		let mut_glob = format!("{mut_glob} kb ({mut_glob_percent:}%)");
 		let diff = &r.0;
 		println!(
 			"| {filename:28} | {original_size:13} kb | {host_fn:>19} | {mut_glob:>22} | {diff:+8}% |"
 		);
 	}
 }