wasm-instrument/src/stack_height/thunk.rs

#[cfg(features = "std")]
use crate::std::collections::{HashMap as Map};
#[cfg(not(features = "std"))]
use crate::std::collections::{BTreeMap as Map};
use crate::std::vec::Vec;

use parity_wasm::elements::{self, FunctionType, Internal};
use parity_wasm::builder;

use super::{resolve_func_type, Context, Error};

struct Thunk {
	signature: FunctionType,
	// Index in function space of this thunk.
	idx: Option<u32>,
	callee_stack_cost: u32,
}

pub(crate) fn generate_thunks(
	ctx: &mut Context,
	module: elements::Module,
) -> Result<elements::Module, Error> {
	// First, we need to collect all function indices that should be replaced by thunks

	let mut replacement_map: Map<u32, Thunk> = {
		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();

		let exported_func_indices = exports.iter().filter_map(|entry| match entry.internal() {
			Internal::Function(function_idx) => Some(*function_idx),
			_ => None,
		});
		let table_func_indices = elem_segments
			.iter()
			.flat_map(|segment| segment.members())
			.cloned();

		// Replacement map is at least export section size.
		let mut replacement_map: Map<u32, Thunk> = Map::new();

		for func_idx in exported_func_indices.chain(table_func_indices).chain(start_func_idx.into_iter()) {
			let callee_stack_cost = ctx.stack_cost(func_idx).ok_or_else(|| {
				Error(format!("function with idx {} isn't found", func_idx))
			})?;

			// 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,
				});
			}
		}

		replacement_map
	};

	// 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,
			thunk.callee_stack_cost as i32,
			ctx.stack_height_global_idx(),
			ctx.stack_limit()
		);
		// Thunk body consist of:
		//  - argument pushing
		//  - instrumented call
		//  - end
		let mut thunk_body: Vec<elements::Instruction> = Vec::with_capacity(
			thunk.signature.params().len() +
			instrumented_call.len() +
			1
		);

		for (arg_idx, _) in thunk.signature.params().iter().enumerate() {
			thunk_body.push(elements::Instruction::GetLocal(arg_idx as u32));
		}
		thunk_body.extend(instrumented_call.iter().cloned());
		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_return_type(thunk.signature.return_type())
					.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.

	// Fixup original function index to a index of a thunk generated earlier.
	let fixup = |function_idx: &mut u32| {
		// Check whether this function is in replacement_map, since
		// we can skip thunk generation (e.g. if stack_cost of function is 0).
		if let Some(thunk) = replacement_map.get(function_idx) {
			*function_idx = thunk
				.idx
				.expect("At this point an index must be assigned to each thunk");
		}
	};

	for section in module.sections_mut() {
		match 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) => {
				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)
			}
			_ => {}
		}
	}

	Ok(module)
}