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.

.github/workflows/check.yml

@@ -11,7 +11,7 @@ env:
 
 jobs:
   rustfmt:
-    runs-on: "ubuntu_20_64_core"
+    runs-on: "ubuntu-latest"
     steps:
       - name: Install Rust toolchain
         uses: actions-rs/toolchain@v1
@@ -20,9 +20,9 @@ jobs:
           toolchain: nightly
           components: rustfmt
 
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v3
 
-      - name: Fmt
+      - name: Cargo fmt
         uses: actions-rs/cargo@v1
         with:
           toolchain: nightly
@@ -30,7 +30,7 @@ jobs:
           args: --all -- --check
 
   clippy:
-    runs-on: "ubuntu_20_64_core"
+    runs-on: "ubuntu-latest"
     steps:
       - name: Install Rust toolchain
         uses: actions-rs/toolchain@v1
@@ -38,75 +38,60 @@ jobs:
           profile: minimal
           toolchain: stable
           components: clippy
-          default: true
 
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v3
 
-      - name: Clippy
+      - name: Cargo clippy
-        uses: actions-rs/clippy-check@v1
+        uses: actions-rs/cargo@v1
         with:
-          token: ${{ secrets.GITHUB_TOKEN }}
-          args: --all-targets --all-features -- -D warnings
-
-  build:
-    runs-on: "ubuntu_20_64_core"
-    steps:
-      - name: Install Rust toolchain
-        uses: actions-rs/toolchain@v1
-        with:
-          profile: minimal
-          target: wasm32-unknown-unknown
           toolchain: stable
-          default: true
+          command: clippy
-
+          args: --all-targets --all-features -- -D warnings
-      - uses: actions/checkout@v4
-
-      - name: Cargo build
-        uses: actions-rs/cargo@v1
-        with:
-          command: build
-
-      - name: Cargo build (std)
-        uses: actions-rs/cargo@v1
-        with:
-          command: build
-          args: --all-features
-
-      - name: Cargo build (no_std)
-        uses: actions-rs/cargo@v1
-        with:
-          command: build
-          args: --no-default-features
-
-      - name: Cargo build (wasm)
-        uses: actions-rs/cargo@v1
-        with:
-          command: build
-          args: --no-default-features --target wasm32-unknown-unknown
 
   test:
     strategy:
       matrix:
-        os: ["ubuntu_20_64_core", "macos-latest", "windows-latest"]
+        os: ["ubuntu-latest", "macos-latest", "windows-latest"]
+        toolchain: ["stable", "nightly"]
     runs-on: ${{ matrix.os }}
     steps:
       - name: Install Rust toolchain
         uses: actions-rs/toolchain@v1
         with:
           profile: minimal
-          toolchain: stable
+          target: wasm32-unknown-unknown
-          default: true
+          toolchain: ${{ matrix.toolchain }}
 
       - name: Set git to use LF
         run: |
           git config --global core.autocrlf false
           git config --global core.eol lf
 
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v3
+
+      - name: Cargo build
+        uses: actions-rs/cargo@v1
+        with:
+          toolchain: ${{ matrix.toolchain }}
+          command: build
+
+      - name: Cargo build (no_std)
+        uses: actions-rs/cargo@v1
+        with:
+          toolchain: ${{ matrix.toolchain }}
+          command: build
+          args: --no-default-features
+
+      - name: Cargo build (wasm)
+        uses: actions-rs/cargo@v1
+        with:
+          toolchain: ${{ matrix.toolchain }}
+          command: build
+          args: --no-default-features --target wasm32-unknown-unknown
 
       - name: Cargo test
         uses: actions-rs/cargo@v1
         with:
+          toolchain: ${{ matrix.toolchain }}
           command: test
           args: --all-features
-

.gitignore

@@ -4,4 +4,3 @@ target
 .DS_Store
 .idea
 .vscode
-*~

BENCHMARKS.md

@@ -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)
-

CHANGELOG.md

@@ -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

Cargo.toml

@@ -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>"]
@@ -12,14 +12,9 @@ repository = "https://github.com/paritytech/wasm-instrument"
 include = ["src/**/*", "LICENSE-*", "README.md"]
 
 [[bench]]
-name = "instrumentation"
+name = "benches"
 harness = false
-path = "benches/instrumentation.rs"
+path = "benches/benches.rs"
-
-[[bench]]
-name = "execution"
-harness = false
-path = "benches/execution.rs"
 
 [profile.bench]
 lto = "fat"
@@ -30,19 +25,15 @@ parity-wasm = { version = "0.45", default-features = false }
 
 [dev-dependencies]
 binaryen = "0.12"
-criterion = "0.5"
+criterion = "0.4"
 diff = "0.1"
 pretty_assertions = "1"
 rand = "0.8"
 wat = "1"
-wasmparser = "0.206"
+wasmparser = "0.92"
-wasmprinter = "0.200"
+wasmprinter = "0.2"
-wasmi = "0.31"
 
 [features]
 default = ["std"]
 std = ["parity-wasm/std"]
 sign_ext = ["parity-wasm/sign_ext"]
-
-[lib]
-bench = false

benches/benches.rs

@@ -7,8 +7,7 @@ use std::{
 	path::PathBuf,
 };
 use wasm_instrument::{
-	gas_metering::{self, host_function, ConstantCostRules},
+	gas_metering, inject_stack_limiter,
-	inject_stack_limiter,
 	parity_wasm::{deserialize_buffer, elements::Module},
 };
 
@@ -16,18 +15,17 @@ fn fixture_dir() -> PathBuf {
 	let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
 	path.push("benches");
 	path.push("fixtures");
-	path.push("wasm");
 	path
 }
 
-fn for_fixtures<F, M>(group: &mut BenchmarkGroup<M>, f: F)
+fn any_fixture<F, M>(group: &mut BenchmarkGroup<M>, f: F)
 where
 	F: Fn(Module),
 	M: Measurement,
 {
 	for entry in read_dir(fixture_dir()).unwrap() {
 		let entry = entry.unwrap();
-		let bytes = read(entry.path()).unwrap();
+		let bytes = read(&entry.path()).unwrap();
 		group.throughput(Throughput::Bytes(bytes.len().try_into().unwrap()));
 		group.bench_with_input(entry.file_name().to_str().unwrap(), &bytes, |bench, input| {
 			bench.iter(|| f(deserialize_buffer(input).unwrap()))
@@ -37,19 +35,14 @@ where
 
 fn gas_metering(c: &mut Criterion) {
 	let mut group = c.benchmark_group("Gas Metering");
-	for_fixtures(&mut group, |module| {
+	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();
 	});
 }
 
 fn stack_height_limiter(c: &mut Criterion) {
 	let mut group = c.benchmark_group("Stack Height Limiter");
-	for_fixtures(&mut group, |module| {
+	any_fixture(&mut group, |module| {
 		inject_stack_limiter(module, 128).unwrap();
 	});
 }

benches/execution.rs

@@ -1,437 +0,0 @@
-use criterion::{
-	criterion_group, criterion_main, measurement::Measurement, Bencher, BenchmarkGroup, Criterion,
-};
-use std::{
-	fs::read,
-	path::PathBuf,
-	time::{Duration, SystemTime, UNIX_EPOCH},
-};
-use wasm_instrument::{
-	gas_metering::{self, host_function, mutable_global, ConstantCostRules},
-	parity_wasm::{deserialize_buffer, elements::Module, serialize},
-};
-use wasmi::{
-	core::{Pages, TrapCode, F32},
-	Caller, Config, Engine, Instance, Linker, Memory, StackLimits, Store, TypedFunc, Value,
-};
-
-/// Describes a gas metering strategy we want to benchmark.
-///
-/// Most strategies just need a subset of these functions. Hence we added default
-/// implementations for all of them.
-trait MeteringStrategy {
-	/// The wasmi config we should be using for this strategy.
-	fn config() -> Config {
-		Config::default()
-	}
-
-	/// The strategy may or may not want to instrument the module.
-	fn instrument_module(module: Module) -> Module {
-		module
-	}
-
-	/// The strategy might need to define additional host functions.
-	fn define_host_funcs(_linker: &mut Linker<u64>) {}
-
-	/// The strategy might need to do some initialization of the wasm instance.
-	fn init_instance(_module: &mut BenchInstance) {}
-}
-
-/// Don't do any metering at all. This is helpful as a baseline.
-struct NoMetering;
-
-/// Use wasmi's builtin fuel metering.
-struct WasmiMetering;
-
-/// Instrument the module using [`host_function::Injector`].
-struct HostFunctionMetering;
-
-/// Instrument the module using [`mutable_global::Injector`].
-struct MutableGlobalMetering;
-
-impl MeteringStrategy for NoMetering {}
-
-impl MeteringStrategy for WasmiMetering {
-	fn config() -> Config {
-		let mut config = Config::default();
-		config.consume_fuel(true);
-		config
-	}
-
-	fn init_instance(module: &mut BenchInstance) {
-		module.store.add_fuel(u64::MAX).unwrap();
-	}
-}
-
-impl MeteringStrategy for HostFunctionMetering {
-	fn instrument_module(module: Module) -> Module {
-		let backend = host_function::Injector::new("env", "gas");
-		gas_metering::inject(module, backend, &ConstantCostRules::default()).unwrap()
-	}
-
-	fn define_host_funcs(linker: &mut Linker<u64>) {
-		// the instrumentation relies on the existing of this function
-		linker
-			.func_wrap("env", "gas", |mut caller: Caller<'_, u64>, amount_consumed: u64| {
-				let gas_remaining = caller.data_mut();
-				*gas_remaining =
-					gas_remaining.checked_sub(amount_consumed).ok_or(TrapCode::OutOfFuel)?;
-				Ok(())
-			})
-			.unwrap();
-	}
-}
-
-impl MeteringStrategy for MutableGlobalMetering {
-	fn instrument_module(module: Module) -> Module {
-		let backend = mutable_global::Injector::new("gas_left");
-		gas_metering::inject(module, backend, &ConstantCostRules::default()).unwrap()
-	}
-
-	fn init_instance(module: &mut BenchInstance) {
-		// the instrumentation relies on the host to initialize the global with the gas limit
-		// we just init to the maximum so it will never run out
-		module
-			.instance
-			.get_global(&mut module.store, "gas_left")
-			.unwrap()
-			.set(&mut module.store, Value::I64(-1i64)) // the same as u64::MAX
-			.unwrap();
-	}
-}
-
-/// A wasm instance ready to be benchmarked.
-struct BenchInstance {
-	store: Store<u64>,
-	instance: Instance,
-}
-
-impl BenchInstance {
-	/// Create a new instance for the supplied metering strategy.
-	///
-	/// `wasm`: The raw wasm module for the benchmark.
-	/// `define_host_func`: In here the caller can define additional host function.
-	fn new<S, H>(wasm: &[u8], define_host_funcs: &H) -> Self
-	where
-		S: MeteringStrategy,
-		H: Fn(&mut Linker<u64>),
-	{
-		let module = deserialize_buffer(wasm).unwrap();
-		let instrumented_module = S::instrument_module(module);
-		let input = serialize(instrumented_module).unwrap();
-		let mut config = S::config();
-		config.set_stack_limits(StackLimits::new(1024, 1024 * 1024, 64 * 1024).unwrap());
-		let engine = Engine::new(&config);
-		let module = wasmi::Module::new(&engine, &mut &input[..]).unwrap();
-		let mut linker = Linker::new(&engine);
-		S::define_host_funcs(&mut linker);
-		define_host_funcs(&mut linker);
-		// init host state with maximum gas_left (only used by host_function instrumentation)
-		let mut store = Store::new(&engine, u64::MAX);
-		let instance = linker.instantiate(&mut store, &module).unwrap().start(&mut store).unwrap();
-		let mut bench_module = Self { store, instance };
-		S::init_instance(&mut bench_module);
-		bench_module
-	}
-}
-
-/// Runs a benchmark for every strategy.
-///
-/// We require the closures to implement `Fn` as they are executed for every strategy and we
-/// don't want them to change in between.
-///
-/// `group`: The benchmark group within the benchmarks will be executed.
-/// `wasm`: The raw wasm module for the benchmark.
-/// `define_host_func`: In here the caller can define additional host function.
-/// `f`: In here the user should perform the benchmark. Will be executed for every strategy.
-fn for_strategies<M, H, F>(group: &mut BenchmarkGroup<M>, wasm: &[u8], define_host_funcs: H, f: F)
-where
-	M: Measurement,
-	H: Fn(&mut Linker<u64>),
-	F: Fn(&mut Bencher<M>, &mut BenchInstance),
-{
-	let mut module = BenchInstance::new::<NoMetering, _>(wasm, &define_host_funcs);
-	group.bench_function("no_metering", |bench| f(bench, &mut module));
-
-	let mut module = BenchInstance::new::<WasmiMetering, _>(wasm, &define_host_funcs);
-	group.bench_function("wasmi_builtin", |bench| f(bench, &mut module));
-
-	let mut module = BenchInstance::new::<HostFunctionMetering, _>(wasm, &define_host_funcs);
-	group.bench_function("host_function", |bench| f(bench, &mut module));
-
-	let mut module = BenchInstance::new::<MutableGlobalMetering, _>(wasm, &define_host_funcs);
-	group.bench_function("mutable_global", |bench| f(bench, &mut module));
-}
-
-/// Converts the `.wat` encoded `bytes` into `.wasm` encoded bytes.
-fn wat2wasm(bytes: &[u8]) -> Vec<u8> {
-	wat::parse_bytes(bytes).unwrap().into_owned()
-}
-
-fn fixture_dir() -> PathBuf {
-	let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
-	path.push("benches");
-	path.push("fixtures");
-	path.push("wasm");
-	path
-}
-
-fn gas_metered_coremark(c: &mut Criterion) {
-	let mut group = c.benchmark_group("coremark");
-	// Benchmark host_function::Injector
-	let wasm_filename = "coremark_minimal.wasm";
-	let bytes = read(fixture_dir().join(wasm_filename)).unwrap();
-	let define_host_funcs = |linker: &mut Linker<u64>| {
-		linker
-			.func_wrap("env", "clock_ms", || {
-				SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as u64
-			})
-			.unwrap();
-	};
-	for_strategies(&mut group, &bytes, define_host_funcs, |bench, module| {
-		bench.iter(|| {
-			let run = module.instance.get_typed_func::<(), F32>(&mut module.store, "run").unwrap();
-			// Call the wasm!
-			run.call(&mut module.store, ()).unwrap();
-		})
-	});
-}
-
-fn gas_metered_recursive_ok(c: &mut Criterion) {
-	let mut group = c.benchmark_group("recursive_ok");
-	const RECURSIVE_DEPTH: i32 = 8000;
-	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/recursive_ok.wat"));
-	for_strategies(
-		&mut group,
-		&wasm_bytes,
-		|_| {},
-		|bench, module| {
-			let bench_call =
-				module.instance.get_typed_func::<i32, i32>(&module.store, "call").unwrap();
-			bench.iter(|| {
-				let result = bench_call.call(&mut module.store, RECURSIVE_DEPTH).unwrap();
-				assert_eq!(result, 0);
-			})
-		},
-	);
-}
-
-fn gas_metered_fibonacci_recursive(c: &mut Criterion) {
-	let mut group = c.benchmark_group("fibonacci_recursive");
-	const FIBONACCI_REC_N: i64 = 10;
-	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/fibonacci.wat"));
-	for_strategies(
-		&mut group,
-		&wasm_bytes,
-		|_| {},
-		|bench, module| {
-			let bench_call = module
-				.instance
-				.get_typed_func::<i64, i64>(&module.store, "fib_recursive")
-				.unwrap();
-			bench.iter(|| {
-				bench_call.call(&mut module.store, FIBONACCI_REC_N).unwrap();
-			});
-		},
-	);
-}
-
-fn gas_metered_fac_recursive(c: &mut Criterion) {
-	let mut group = c.benchmark_group("factorial_recursive");
-	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/factorial.wat"));
-	for_strategies(
-		&mut group,
-		&wasm_bytes,
-		|_| {},
-		|bench, module| {
-			let fac = module
-				.instance
-				.get_typed_func::<i64, i64>(&module.store, "recursive_factorial")
-				.unwrap();
-			bench.iter(|| {
-				let result = fac.call(&mut module.store, 25).unwrap();
-				assert_eq!(result, 7034535277573963776);
-			})
-		},
-	);
-}
-
-fn gas_metered_count_until(c: &mut Criterion) {
-	const COUNT_UNTIL: i32 = 100_000;
-	let mut group = c.benchmark_group("count_until");
-	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/count_until.wat"));
-	for_strategies(
-		&mut group,
-		&wasm_bytes,
-		|_| {},
-		|bench, module| {
-			let count_until = module
-				.instance
-				.get_typed_func::<i32, i32>(&module.store, "count_until")
-				.unwrap();
-			bench.iter(|| {
-				let result = count_until.call(&mut module.store, COUNT_UNTIL).unwrap();
-				assert_eq!(result, COUNT_UNTIL);
-			})
-		},
-	);
-}
-
-fn gas_metered_vec_add(c: &mut Criterion) {
-	fn test_for<A, B>(
-		b: &mut Bencher,
-		vec_add: TypedFunc<(i32, i32, i32, i32), ()>,
-		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 = (ptr_result as i32, ptr_a as i32, ptr_b as i32, len as i32);
-		b.iter(|| {
-			vec_add.call(&mut store, params).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");
-	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/memory-vec-add.wat"));
-	const LEN: usize = 100_000;
-
-	for_strategies(
-		&mut group,
-		&wasm_bytes,
-		|_| {},
-		|bench, module| {
-			let vec_add = module
-				.instance
-				.get_typed_func::<(i32, i32, i32, i32), ()>(&module.store, "vec_add")
-				.unwrap();
-			let mem = module.instance.get_memory(&module.store, "mem").unwrap();
-			mem.grow(&mut module.store, Pages::new(25).unwrap()).unwrap();
-			test_for(
-				bench,
-				vec_add,
-				&mut module.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");
-	let wasm_filename = "wasm_kernel.wasm";
-	let wasm_bytes = read(fixture_dir().join(wasm_filename)).unwrap();
-	for_strategies(
-		&mut group,
-		&wasm_bytes,
-		|_| {},
-		|bench, module| {
-			let prepare = module
-				.instance
-				.get_typed_func::<(), i32>(&module.store, "prepare_tiny_keccak")
-				.unwrap();
-			let keccak = module
-				.instance
-				.get_typed_func::<i32, ()>(&module.store, "bench_tiny_keccak")
-				.unwrap();
-			let test_data_ptr = prepare.call(&mut module.store, ()).unwrap();
-			bench.iter(|| {
-				keccak.call(&mut module.store, test_data_ptr).unwrap();
-			})
-		},
-	);
-}
-
-fn gas_metered_global_bump(c: &mut Criterion) {
-	const BUMP_AMOUNT: i32 = 100_000;
-	let mut group = c.benchmark_group("global_bump");
-	let wasm_bytes = wat2wasm(include_bytes!("fixtures/wat/global_bump.wat"));
-	for_strategies(
-		&mut group,
-		&wasm_bytes,
-		|_| {},
-		|bench, module| {
-			let bump = module.instance.get_typed_func::<i32, i32>(&module.store, "bump").unwrap();
-			bench.iter(|| {
-				let result = bump.call(&mut module.store, BUMP_AMOUNT).unwrap();
-				assert_eq!(result, BUMP_AMOUNT);
-			})
-		},
-	);
-}
-
-criterion_group!(
-	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);

benches/fixtures/wat/count_until.wat

@@ -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))
-  )
-)

benches/fixtures/wat/factorial.wat

@@ -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)
-                        )
-                    )
-                )
-            )
-        )
-    )
-)

benches/fixtures/wat/fibonacci.wat

@@ -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)
-    )
-)

benches/fixtures/wat/global_bump.wat

@@ -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))
-    )
-)

benches/fixtures/wat/memory-vec-add.wat

@@ -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)
-    )
-)

benches/fixtures/wat/recursive_ok.wat

@@ -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))
-            )
-        )
-    )
-)

src/gas_metering/backend.rs

@@ -1,139 +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: u64, instruction| {
-				cost.saturating_add(rules.instruction_cost(instruction).unwrap_or(u32::MAX).into())
-			});
-			// don't charge for the instructions used to fail when out of gas
-			let fail_cost = [
-				Instruction::I64Const(-1i64),           // non-charged instruction
-				Instruction::SetGlobal(gas_global_idx), // non-charged instruction
-				Instruction::Unreachable,               // non-charged instruction
-			]
-			.iter()
-			.fold(0, |cost: u64, instruction| {
-				cost.saturating_add(rules.instruction_cost(instruction).unwrap_or(u32::MAX).into())
-			});
-
-			// the fail costs are a subset of the overall costs and hence this never underflows
-			gas_fn_cost -= fail_cost;
-
-			GasMeter::Internal {
-				global: self.global_name,
-				func_instructions: elements::Instructions::new(func_instructions),
-				cost: gas_fn_cost,
-			}
-		}
-	}
-}

src/gas_metering/validation.rs

@@ -134,14 +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()
-		.try_fold(0u32, |count, val_type| count.checked_add(val_type.count()))
-		.ok_or(())?;
-	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()
@@ -157,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(_) => {
@@ -354,10 +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().map(|val_type| val_type.count()).sum();
 
-				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);

src/stack_limiter/max_height.rs

@@ -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")?;
@@ -176,7 +181,7 @@ pub fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, &'static
 		match opcode {
 			Nop => {},
 			Block(ty) | Loop(ty) | If(ty) => {
-				let end_arity = u32::from(*ty != BlockType::NoResult);
+				let end_arity = if *ty == BlockType::NoResult { 0 } else { 1 };
 				let branch_arity = if let Loop(_) = *opcode { 0 } else { end_arity };
 				if let If(_) = *opcode {
 					stack.pop_values(1)?;
@@ -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,52 +446,46 @@ 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)
 	unreachable
-	memory.grow
+	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);
 	}
 }

src/stack_limiter/mod.rs

@@ -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() {
@@ -222,8 +291,8 @@ fn instrument_functions(
 /// Before:
 ///
 /// ```text
-/// local.get 0
+/// get_local 0
-/// local.get 1
+/// get_local 1
 /// call 228
 /// drop
 /// ```
@@ -231,8 +300,8 @@ fn instrument_functions(
 /// After:
 ///
 /// ```text
-/// local.get 0
+/// get_local 0
-/// local.get 1
+/// get_local 1
 ///
 /// < ... preamble ... >
 ///
@@ -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::*;
@@ -366,9 +399,9 @@ mod tests {
 			r#"
 (module
 	(func (export "i32.add") (param i32 i32) (result i32)
-		local.get 0
+		get_local 0
-		local.get 1
+	get_local 1
-		i32.add
+	i32.add
 	)
 )
 "#,

src/stack_limiter/thunk.rs

@@ -1,13 +1,17 @@
-use alloc::{collections::BTreeMap as Map, vec::Vec};
+#[cfg(not(features = "std"))]
-use parity_wasm::{
+use alloc::collections::BTreeMap as Map;
-	builder,
+use alloc::vec::Vec;
-	elements::{self, FunctionType, Internal},
+use parity_wasm::elements::{self, 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,
@@ -15,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();
@@ -41,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 });
-					},
-				);
 			}
 		}
 
@@ -57,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,
@@ -73,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.
 
@@ -114,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),
 			_ => {},
 		}
@@ -133,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)
+}

tests/diff.rs

@@ -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);
 }

tests/expectations/gas/branch.wat

@@ -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

tests/expectations/gas/branch_mut_global.wat

@@ -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))
-)

tests/expectations/gas/call.wat

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

tests/expectations/gas/call_mut_global.wat

@@ -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))
-)

tests/expectations/gas/ifs.wat

@@ -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

tests/expectations/gas/ifs_mut_global.wat

@@ -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))
-)

tests/expectations/gas/simple.wat

@@ -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))

tests/expectations/gas/simple_mut_global.wat

@@ -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))
-)

tests/expectations/gas/start_mut_global.wat

@@ -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")
-)

tests/expectations/stack-height/empty_functions.wat

@@ -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

tests/expectations/stack-height/global.wat

@@ -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

tests/expectations/stack-height/imports.wat

@@ -20,7 +20,7 @@
     global.get 0
     i32.const 1024
     i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
       unreachable
     end
     call 2

tests/expectations/stack-height/many_locals.wat

@@ -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

tests/expectations/stack-height/simple.wat

@@ -12,7 +12,7 @@
     global.get 0
     i32.const 1024
     i32.gt_u
-    if ;; label = @1
+    if  ;; label = @1
       unreachable
     end
     call 0

tests/expectations/stack-height/start.wat

@@ -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

tests/expectations/stack-height/table.wat

@@ -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

tests/fixtures/gas/branch.wat

@@ -3,25 +3,25 @@
 		(local $x i32) (local $y i32)
 
 		(block $unrolled_loop
-			(local.set $x (i32.const 0))
+			(set_local $x (i32.const 0))
-			(local.set $y (i32.const 1))
+			(set_local $y (i32.const 1))
 
-			local.get $x
+			get_local $x
-			local.get $y
+			get_local $y
-			local.tee $x
+			tee_local $x
 			i32.add
-			local.set $y
+			set_local $y
 
 			i32.const 1
 			br_if $unrolled_loop
 
-			local.get $x
+			get_local $x
-			local.get $y
+			get_local $y
-			local.tee $x
+			tee_local $x
 			i32.add
-			local.set $y
+			set_local $y
 		)
 
-		local.get $y
+		get_local $y
 	)
 )

tests/fixtures/gas/call.wat

@@ -2,18 +2,18 @@
 	(func $add_locals (param $x i32) (param $y i32) (result i32)
 		(local $t i32)
 
-		local.get $x
+		get_local $x
-		local.get $y
+		get_local $y
 		call $add
-		local.set $t
+		set_local $t
 
-		local.get $t
+		get_local $t
 	)
 
 	(func $add (param $x i32) (param $y i32) (result i32)
 		(i32.add
-			(local.get $x)
+			(get_local $x)
-			(local.get $y)
+			(get_local $y)
 		)
 	)
 )

tests/fixtures/gas/ifs.wat

@@ -2,8 +2,8 @@
 	(func (param $x i32) (result i32)
 		(if (result i32)
 			(i32.const 1)
-			(then (i32.add (local.get $x) (i32.const 1)))
+			(then (i32.add (get_local $x) (i32.const 1)))
-			(else (i32.popcnt (local.get $x)))
+			(else (i32.popcnt (get_local $x)))
 		)
 	)
 )

tests/fixtures/stack-height/global.wat

@@ -5,22 +5,22 @@
   (global $counter (mut i32) (i32.const 1))
 
   (func $i32.add (export "i32.add") (param i32 i32) (result i32)
-    local.get 0
+    get_local 0
-    local.get 1
+	  get_local 1
-    i32.add
+	  i32.add
   )
   (func (param $arg i32)
-    (local $tmp i32)
+     (local $tmp i32)
 
-    global.get 0
+     global.get 0
-    i32.const 1
+     i32.const 1
-    i32.add
+     i32.add
-    local.tee $tmp
+	   tee_local $tmp
-    global.set $counter
+     global.set $counter
 
-    local.get $tmp
+     get_local $tmp
-    local.get $arg
+     get_local $arg
-    call $i32.add
+     call $i32.add
-    drop
+     drop
   )
 )

tests/fixtures/stack-height/imports.wat

@@ -8,8 +8,8 @@
     call $foo
     call $boo
 
-    local.get 0
+    get_local 0
-    local.get 1
+    get_local 1
     i32.add
   )
 )

tests/fixtures/stack-height/many_locals.wat

@@ -6,5 +6,5 @@
 		(call
 			$one-group-many-locals
 		)
-	)
+  	)
 )

tests/fixtures/stack-height/table.wat

@@ -1,17 +1,17 @@
 (module
   (import "env" "foo" (func $foo))
   (func (param i32)
-    local.get 0
+    get_local 0
     i32.const 0
     call $i32.add
     drop
   )
   (func $i32.add (export "i32.add") (param i32 i32) (result i32)
-    local.get 0
+    get_local 0
-    local.get 1
+	  get_local 1
-    i32.add
+	  i32.add
   )
-  (table 10 funcref)
+  (table 10 anyfunc)
 
   ;; Refer all types of functions: imported, defined not exported and defined exported.
   (elem (i32.const 0) 0 1 2)

tests/overhead.rs

@@ -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}% |"
-		);
 	}
 }