casper-wasm-utils 4.0.0

use crate::std::collections::{BTreeMap as Map, BTreeSet as Set};
use crate::std::vec::Vec;

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

use super::{instrument_call, 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();

        let set_of_func_idx: Set<_> = exported_func_indices
            .chain(table_func_indices)
            .chain(start_func_idx)
            .collect();

        for func_idx in set_of_func_idx {
            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,
            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_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 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)
}