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.

src/stack_limiter/thunk.rs

@@ -1,17 +1,17 @@
 #[cfg(not(features = "std"))]
 use alloc::collections::BTreeMap as Map;
 use alloc::vec::Vec;
-use parity_wasm::{
-	builder,
-	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,
@@ -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(
-					func_idx,
-					Thunk {
-						signature: resolve_func_type(func_idx, &module)?.clone(),
-						idx: None,
-						callee_stack_cost,
-					},
-				);
+				let type_idx = ctx.func_type(func_idx).ok_or("type idx for thunk not found")?;
+				let elements::Type::Function(func_ty) =
+					types.get(type_idx as usize).ok_or("sig for thunk is not found")?;
+				let param_num = func_ty.params().len() as u32;
+				replacement_map
+					.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() {
-			thunk_body.push(elements::Instruction::GetLocal(arg_idx as u32));
+		for arg_idx in 0..thunk.param_num {
+			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.
-
-		mbuilder = mbuilder
-			.function()
-			// Signature of the thunk should match the original function signature.
-			.signature()
-			.with_params(thunk.signature.params().to_vec())
-			.with_results(thunk.signature.results().to_vec())
-			.build()
-			.body()
-			.with_instructions(elements::Instructions::new(thunk_body))
-			.build()
-			.build();
-
-		thunk.idx = Some(next_func_idx);
-		next_func_idx += 1;
+		let func_idx = insert_function(
+			ctx,
+			&mut module,
+			thunk.type_idx,
+			Vec::new(), // No declared local variables.
+			elements::Instructions::new(thunk_body),
+		)?;
+		thunk.idx = Some(func_idx);
 	}
-	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)
+}