winch-codegen 0.22.2

//! Function call emission.  For more details around the ABI and
//! calling convention, see [ABI].
//!
//! This module exposes a single function [`FnCall::emit`], which is responsible
//! of orchestrating the emission of calls. In general such orchestration
//! takes place in 6 steps:
//!
//! 1. [`Callee`] resolution.
//! 2. Mapping of the [`Callee`] to the [`CalleeKind`].
//! 3. Spilling the value stack.
//! 4. Calculate the return area, for 1+ results.
//! 5. Emission.
//! 6. Stack space cleanup.
//!
//! The stack space consumed by the function call is the amount
//! of space used by any memory entries in the value stack present
//! at the callsite (after spilling the value stack), that will be
//! used as arguments for the function call. Any memory values in the
//! value stack that are needed as part of the function
//! arguments will be consumed by the function call (either by
//! assigning those values to a register or by storing those
//! values in a memory location if the callee argument is on
//! the stack).
//! This could also be done when assigning arguments every time a
//! memory entry needs to be assigned to a particular location,
//! but doing so will emit more instructions (e.g. a pop per
//! argument that needs to be assigned); it's more efficient to
//! calculate the space used by those memory values and reclaim it
//! at once when cleaning up the stack after the call has been
//! emitted.
//!
//! The machine stack throughout the function call is as follows:
//! ┌──────────────────────────────────────────────────┐
//! │                                                  │
//! │                  1                               │
//! │  Stack space created by any previous spills      │
//! │  from the value stack; and which memory values   │
//! │  are used as function arguments.                 │
//! │                                                  │
//! ├──────────────────────────────────────────────────┤ ---> The Wasm value stack at this point in time would look like:
//! │                                                  │      [ Mem(offset) | Mem(offset) | Local(index) | Local(index) ]
//! │                   2                              │
//! │   Stack space created by spilling locals and     |
//! │   registers at the callsite.                     │
//! │                                                  │
//! │                                                  │
//! ├─────────────────────────────────────────────────┬┤ ---> The Wasm value stack at this point in time would look like:
//! │                                                  │      [ Mem(offset) | Mem(offset) | Mem(offset) | Mem(offset) ]
//! │                                                  │      Assuming that the callee takes 4 arguments, we calculate
//! │                                                  │      4 memory values; all of which will be used as arguments to
//! │   Stack space allocated for                      │      the call via `assign_args`, thus the sum of the size of the
//! │   the callee function arguments in the stack;    │      memory they represent is considered to be consumed by the call.
//! │   represented by `arg_stack_space`               │
//! │                                                  │
//! │                                                  │
//! │                                                  │
//! └──────────────────────────────────────────────────┘ ------> Stack pointer when emitting the call

use crate::{
    abi::{scratch, vmctx, ABIOperand, ABISig, RetArea},
    codegen::{BuiltinFunction, BuiltinType, Callee, CodeGenContext},
    masm::{
        CalleeKind, ContextArgs, MacroAssembler, MemMoveDirection, OperandSize, SPOffset,
        VMContextLoc,
    },
    reg::Reg,
    stack::Val,
    FuncEnv,
};
use wasmtime_environ::{FuncIndex, PtrSize, VMOffsets};

/// All the information needed to emit a function call.
#[derive(Copy, Clone)]
pub(crate) struct FnCall {}

impl FnCall {
    /// Orchestrates the emission of a function call:
    /// 1. Resolves the [`Callee`] through the given callback.
    /// 2. Lowers the resolved [`Callee`] to a ([`CalleeKind`], [ContextArgs])
    /// 3. Spills the value stack.
    /// 4. Creates the stack space needed for the return area.
    /// 5. Emits the call.
    /// 6. Cleans up the stack space.
    pub fn emit<M: MacroAssembler>(
        env: &mut FuncEnv<M::Ptr>,
        masm: &mut M,
        context: &mut CodeGenContext,
        callee: Callee,
    ) {
        let (kind, callee_context) = Self::lower(env, context.vmoffsets, &callee, context, masm);

        let sig = env.callee_sig::<M::ABI>(&callee);
        context.spill(masm);
        let ret_area = Self::make_ret_area(&sig, masm);
        let arg_stack_space = sig.params_stack_size();
        let reserved_stack = masm.call(arg_stack_space, |masm| {
            Self::assign(sig, &callee_context, ret_area.as_ref(), context, masm);
            kind
        });

        Self::cleanup(
            sig,
            &callee_context,
            &kind,
            reserved_stack,
            ret_area,
            masm,
            context,
        );
    }

    /// Calculates the return area for the callee, if any.
    fn make_ret_area<M: MacroAssembler>(callee_sig: &ABISig, masm: &mut M) -> Option<RetArea> {
        callee_sig.has_stack_results().then(|| {
            let base = masm.sp_offset().as_u32();
            let end = base + callee_sig.results_stack_size();
            if end > base {
                masm.reserve_stack(end - base);
            }
            RetArea::sp(SPOffset::from_u32(end))
        })
    }

    /// Lowers the high-level [`Callee`] to a [`CalleeKind`] and
    /// [ContextArgs] pair which contains all the metadata needed for
    /// emission.
    fn lower<M: MacroAssembler>(
        env: &mut FuncEnv<M::Ptr>,
        vmoffsets: &VMOffsets<u8>,
        callee: &Callee,
        context: &mut CodeGenContext,
        masm: &mut M,
    ) -> (CalleeKind, ContextArgs) {
        let ptr = vmoffsets.ptr.size();
        match callee {
            Callee::Builtin(b) => Self::lower_builtin(env, b),
            Callee::FuncRef(_) => {
                Self::lower_funcref(env.callee_sig::<M::ABI>(callee), ptr, context, masm)
            }
            Callee::Local(i) => Self::lower_local(env, *i),
            Callee::Import(i) => {
                let sig = env.callee_sig::<M::ABI>(callee);
                Self::lower_import(*i, sig, context, masm, vmoffsets)
            }
        }
    }

    /// Lowers a builtin function by loading its address to the next available
    /// register.
    fn lower_builtin<P: PtrSize>(
        env: &mut FuncEnv<P>,
        builtin: &BuiltinFunction,
    ) -> (CalleeKind, ContextArgs) {
        match builtin.ty() {
            BuiltinType::Builtin(idx) => (
                CalleeKind::direct(env.name_builtin(idx)),
                ContextArgs::pinned_vmctx(),
            ),
            BuiltinType::LibCall(c) => (CalleeKind::libcall(c), ContextArgs::none()),
        }
    }

    /// Lower  a local function to a [`CalleeKind`] and [ContextArgs] pair.
    fn lower_local<P: PtrSize>(
        env: &mut FuncEnv<P>,
        index: FuncIndex,
    ) -> (CalleeKind, ContextArgs) {
        (
            CalleeKind::direct(env.name_wasm(index)),
            ContextArgs::pinned_callee_and_caller_vmctx(),
        )
    }

    /// Lowers a function import by loading its address to the next available
    /// register.
    fn lower_import<M: MacroAssembler, P: PtrSize>(
        index: FuncIndex,
        sig: &ABISig,
        context: &mut CodeGenContext,
        masm: &mut M,
        vmoffsets: &VMOffsets<P>,
    ) -> (CalleeKind, ContextArgs) {
        let (callee, callee_vmctx) =
            context.without::<(Reg, Reg), M, _>(&sig.regs, masm, |context, masm| {
                (context.any_gpr(masm), context.any_gpr(masm))
            });
        let callee_vmctx_offset = vmoffsets.vmctx_vmfunction_import_vmctx(index);
        let callee_vmctx_addr = masm.address_at_vmctx(callee_vmctx_offset);
        masm.load_ptr(callee_vmctx_addr, callee_vmctx);

        let callee_body_offset = vmoffsets.vmctx_vmfunction_import_wasm_call(index);
        let callee_addr = masm.address_at_vmctx(callee_body_offset);
        masm.load_ptr(callee_addr, callee);

        (
            CalleeKind::indirect(callee),
            ContextArgs::with_callee_and_pinned_caller(callee_vmctx),
        )
    }

    /// Lowers a function reference by loading its address into the next
    /// available register.
    fn lower_funcref<M: MacroAssembler>(
        sig: &ABISig,
        ptr: impl PtrSize,
        context: &mut CodeGenContext,
        masm: &mut M,
    ) -> (CalleeKind, ContextArgs) {
        // Pop the funcref pointer to a register and allocate a register to hold the
        // address of the funcref. Since the callee is not addressed from a global non
        // allocatable register (like the vmctx in the case of an import), we load the
        // funcref to a register ensuring that it doesn't get assigned to a register
        // used in the callee's signature.
        let (funcref_ptr, funcref, callee_vmctx) =
            context.without::<_, M, _>(&sig.regs, masm, |cx, masm| {
                (
                    cx.pop_to_reg(masm, None).into(),
                    cx.any_gpr(masm),
                    cx.any_gpr(masm),
                )
            });

        // Load the callee VMContext, that will be passed as first argument to
        // the function call.
        masm.load_ptr(
            masm.address_at_reg(funcref_ptr, ptr.vm_func_ref_vmctx().into()),
            callee_vmctx,
        );

        // Load the function pointer to be called.
        masm.load_ptr(
            masm.address_at_reg(funcref_ptr, ptr.vm_func_ref_wasm_call().into()),
            funcref,
        );
        context.free_reg(funcref_ptr);

        (
            CalleeKind::indirect(funcref),
            ContextArgs::with_callee_and_pinned_caller(callee_vmctx),
        )
    }

    /// Materializes any [ContextArgs] as a function argument.
    fn assign_context_args<M: MacroAssembler>(sig: &ABISig, context: &ContextArgs, masm: &mut M) {
        debug_assert!(sig.params().len() >= context.len());
        for (context_arg, operand) in context
            .as_slice()
            .iter()
            .zip(sig.params_without_retptr().iter().take(context.len()))
        {
            match (context_arg, operand) {
                (VMContextLoc::Pinned, ABIOperand::Reg { ty, reg, .. }) => {
                    masm.mov(vmctx!(M).into(), *reg, (*ty).into());
                }
                (VMContextLoc::Pinned, ABIOperand::Stack { ty, offset, .. }) => {
                    let addr = masm.address_at_sp(SPOffset::from_u32(*offset));
                    masm.store(vmctx!(M).into(), addr, (*ty).into());
                }

                (VMContextLoc::Reg(src), ABIOperand::Reg { ty, reg, .. }) => {
                    masm.mov((*src).into(), *reg, (*ty).into());
                }

                (VMContextLoc::Reg(src), ABIOperand::Stack { ty, offset, .. }) => {
                    let addr = masm.address_at_sp(SPOffset::from_u32(*offset));
                    masm.store((*src).into(), addr, (*ty).into());
                }
            }
        }
    }

    /// Assign arguments for the function call.
    fn assign<M: MacroAssembler>(
        sig: &ABISig,
        callee_context: &ContextArgs,
        ret_area: Option<&RetArea>,
        context: &mut CodeGenContext,
        masm: &mut M,
    ) {
        let arg_count = sig.params.len_without_retptr();
        debug_assert!(arg_count >= callee_context.len());
        let stack = &context.stack;
        let stack_values = stack.peekn(arg_count - callee_context.len());

        if callee_context.len() > 0 {
            Self::assign_context_args(&sig, &callee_context, masm);
        }

        for (arg, val) in sig
            .params_without_retptr()
            .iter()
            .skip(callee_context.len())
            .zip(stack_values)
        {
            match arg {
                &ABIOperand::Reg { reg, .. } => {
                    context.move_val_to_reg(&val, reg, masm);
                }
                &ABIOperand::Stack { ty, offset, .. } => {
                    let addr = masm.address_at_sp(SPOffset::from_u32(offset));
                    let size: OperandSize = ty.into();
                    let scratch = scratch!(M, &ty);
                    context.move_val_to_reg(val, scratch, masm);
                    masm.store(scratch.into(), addr, size);
                }
            }
        }

        if sig.has_stack_results() {
            let operand = sig.params.unwrap_results_area_operand();
            let base = ret_area.unwrap().unwrap_sp();
            let addr = masm.address_from_sp(base);

            match operand {
                &ABIOperand::Reg { ty, reg, .. } => {
                    masm.load_addr(addr, reg, ty.into());
                }
                &ABIOperand::Stack { ty, offset, .. } => {
                    let slot = masm.address_at_sp(SPOffset::from_u32(offset));
                    // Don't rely on `ABI::scratch_for` as we always use
                    // an int register as the return pointer.
                    let scratch = scratch!(M);
                    masm.load_addr(addr, scratch, ty.into());
                    masm.store(scratch.into(), slot, ty.into());
                }
            }
        }
    }

    /// Cleanup stack space, handle multiple results, and free registers after
    /// emitting the call.
    fn cleanup<M: MacroAssembler>(
        sig: &ABISig,
        callee_context: &ContextArgs,
        callee_kind: &CalleeKind,
        reserved_space: u32,
        ret_area: Option<RetArea>,
        masm: &mut M,
        context: &mut CodeGenContext,
    ) {
        // Free any registers holding any function references.
        match callee_kind {
            CalleeKind::Indirect(r) => context.free_reg(*r),
            _ => {}
        }

        // Free any registers used as part of the [ContextArgs].
        for loc in callee_context.as_slice() {
            match loc {
                VMContextLoc::Reg(r) => context.free_reg(*r),
                _ => {}
            }
        }
        // Deallocate the reserved space for stack arguments and for alignment,
        // which was allocated last.
        masm.free_stack(reserved_space);

        debug_assert!(sig.params.len_without_retptr() >= callee_context.len());

        // Drop params from value stack and calculate amount of machine stack
        // space they consumed.
        let mut stack_consumed = 0;
        context.drop_last(
            sig.params.len_without_retptr() - callee_context.len(),
            |_regalloc, v| {
                debug_assert!(v.is_mem() || v.is_const());
                if let Val::Memory(mem) = v {
                    stack_consumed += mem.slot.size;
                }
            },
        );

        if let Some(ret_area) = ret_area {
            if stack_consumed > 0 {
                // Perform a memory move, by shuffling the result area to
                // higher addresses. This is needed because the result area
                // is located after any memory addresses located on the stack,
                // and after spilled values consumed by the call.
                let sp = ret_area.unwrap_sp();
                let result_bytes = sig.results_stack_size();
                debug_assert!(sp.as_u32() >= stack_consumed + result_bytes);
                let dst = SPOffset::from_u32(sp.as_u32() - stack_consumed);
                masm.memmove(sp, dst, result_bytes, MemMoveDirection::LowToHigh);
            }
        };

        // Free the bytes consumed by the call.
        masm.free_stack(stack_consumed);

        let ret_area = ret_area.map(|area| {
            if stack_consumed > 0 {
                // If there's a return area and stack space was consumed by the
                // call, adjust the return area to be to the current stack
                // pointer offset.
                RetArea::sp(masm.sp_offset())
            } else {
                // Else if no stack space was consumed by the call, simply use
                // the previously calculated area.
                debug_assert_eq!(area.unwrap_sp(), masm.sp_offset());
                area
            }
        });

        // In the case of [Callee], there's no need to set the [RetArea] of the
        // signature, as it's only used here to push abi results.
        context.push_abi_results(&sig.results, masm, |_, _, _| ret_area);
        // Reload the [VMContext] pointer into the corresponding pinned
        // register. Winch currently doesn't have any callee-saved registers in
        // the default ABI. So the callee might clobber the designated pinned
        // register.
        context.load_vmctx(masm);
    }
}