winch_codegen/codegen/

mod.rs

use crate::{
    Result,
    abi::{ABI, ABIOperand, ABISig, LocalSlot, RetArea, vmctx},
    bail,
    codegen::BlockSig,
    ensure, format_err,
    isa::reg::{Reg, RegClass, writable},
    masm::{
        AtomicWaitKind, Extend, Imm, IntCmpKind, IntScratch, LaneSelector, LoadKind,
        MacroAssembler, OperandSize, RegImm, RmwOp, SPOffset, ShiftKind, StoreKind, TrapCode,
        UNTRUSTED_FLAGS, Zero,
    },
    stack::{TypedReg, Val},
};
use cranelift_codegen::{
    binemit::CodeOffset,
    ir::{RelSourceLoc, SourceLoc},
};
use smallvec::SmallVec;
use std::marker::PhantomData;
use wasmparser::{
    BinaryReader, FuncValidator, MemArg, Operator, OperatorsReader, ValidatorResources,
    VisitOperator, VisitSimdOperator, WasmFeatures,
};
use wasmtime_cranelift::{TRAP_BAD_SIGNATURE, TRAP_HEAP_MISALIGNED, TRAP_TABLE_OUT_OF_BOUNDS};
use wasmtime_environ::{
    DataIndex, ElemIndex, FUNCREF_INIT_BIT, FUNCREF_MASK, GlobalIndex, IndexType, MemoryIndex,
    MemoryKind, MemoryTunables, PtrSize, TableIndex, Tunables, TypeIndex, WasmHeapType,
    WasmValType, wasm_unsupported,
};

mod context;
pub(crate) use context::*;
mod env;
pub use env::*;
mod call;
pub(crate) use call::*;
mod control;
pub(crate) use control::*;
mod builtin;
pub use builtin::*;
pub(crate) mod bounds;
mod drc;
mod exceptions;
pub(crate) use exceptions::{CatchInfo, TryTableInfo};
mod gc;
use gc::GcCodegenConfig;

use bounds::{Bounds, ImmOffset, Index};

mod phase;
pub(crate) use phase::*;

mod error;
pub(crate) use error::*;

/// Branch states in the compiler, enabling the derivation of the
/// reachability state.
pub(crate) trait BranchState {
    /// Whether the compiler will enter in an unreachable state after
    /// the branch is emitted.
    fn unreachable_state_after_emission() -> bool;
}

/// A conditional branch state, with a fallthrough.
pub(crate) struct ConditionalBranch;

impl BranchState for ConditionalBranch {
    fn unreachable_state_after_emission() -> bool {
        false
    }
}

/// Unconditional branch state.
pub(crate) struct UnconditionalBranch;

impl BranchState for UnconditionalBranch {
    fn unreachable_state_after_emission() -> bool {
        true
    }
}

/// Holds metadata about the source code location and the machine code emission.
/// The fields of this struct are opaque and are not interpreted in any way.
/// They serve as a mapping between source code and machine code.
#[derive(Default)]
pub(crate) struct SourceLocation {
    /// The base source location.
    pub base: Option<SourceLoc>,
    /// The current relative source code location along with its associated
    /// machine code offset.
    pub current: (CodeOffset, RelSourceLoc),
}

/// The code generation abstraction.
pub(crate) struct CodeGen<'a, 'translation: 'a, 'data: 'translation, M, P>
where
    M: MacroAssembler,
    P: CodeGenPhase,
{
    /// The ABI-specific representation of the function signature, excluding results.
    pub sig: ABISig,

    /// The code generation context.
    pub context: CodeGenContext<'a, P>,

    /// A reference to the function compilation environment.
    pub env: FuncEnv<'a, 'translation, 'data, M::Ptr>,

    /// The MacroAssembler.
    pub masm: &'a mut M,

    /// Stack frames for control flow.
    // NB The 64 is set arbitrarily, we can adjust it as
    // we see fit.
    pub control_frames: SmallVec<[ControlStackFrame; 64]>,

    /// Information about the source code location.
    pub source_location: SourceLocation,

    /// Compilation settings for code generation.
    pub tunables: &'a Tunables,

    /// Local counter to track fuel consumption.
    pub fuel_consumed: i64,

    /// Whether this function accesses the store's GC heap.
    pub needs_gc_heap: bool,

    /// Collector-specific configuration for generating GC operations.
    gc_codegen_config: Option<GcCodegenConfig>,
    phase: PhantomData<P>,
}

impl<'a, 'translation, 'data, M> CodeGen<'a, 'translation, 'data, M, Prologue>
where
    M: MacroAssembler,
{
    pub fn new(
        tunables: &'a Tunables,
        masm: &'a mut M,
        context: CodeGenContext<'a, Prologue>,
        env: FuncEnv<'a, 'translation, 'data, M::Ptr>,
        sig: ABISig,
        wasm_features: &WasmFeatures,
    ) -> Result<CodeGen<'a, 'translation, 'data, M, Prologue>> {
        let gc_codegen_config = match tunables.collector {
            Some(collector) => Some(GcCodegenConfig::new(collector)),
            None if wasm_features.contains(WasmFeatures::EXCEPTIONS) => {
                return Err(format_err!(wasm_unsupported!(
                    "support for GC types disabled at configuration time"
                )));
            }
            None => None,
        };

        Ok(Self {
            sig,
            context,
            masm,
            env,
            tunables,
            source_location: Default::default(),
            control_frames: Default::default(),
            // Empty functions should consume at least 1 fuel unit.
            fuel_consumed: 1,
            needs_gc_heap: false,
            gc_codegen_config,
            phase: PhantomData,
        })
    }

    /// Code generation prologue.
    pub fn emit_prologue(mut self) -> Result<CodeGen<'a, 'translation, 'data, M, Emission>> {
        let vmctx = self
            .sig
            .params()
            .first()
            .ok_or_else(|| format_err!(CodeGenError::vmcontext_arg_expected()))?
            .unwrap_reg();

        self.masm.start_source_loc(Default::default())?;
        // We need to use the vmctx parameter before pinning it for stack checking.
        self.masm.prologue(vmctx)?;

        // Pin the `VMContext` pointer.
        self.masm.mov(
            writable!(vmctx!(M)),
            vmctx.into(),
            self.env.ptr_type().try_into()?,
        )?;

        self.masm.reserve_stack(self.context.frame.locals_size)?;
        self.spill_register_arguments()?;
        self.copy_stack_gc_refs_to_frame()?;

        let defined_locals_range = &self.context.frame.defined_locals_range;
        self.masm.zero_mem_range(defined_locals_range.as_range())?;

        // Save the results base parameter register into its slot.

        if self.sig.params.has_retptr() {
            match self.sig.params.unwrap_results_area_operand() {
                ABIOperand::Reg { ty, reg, .. } => {
                    let results_base_slot = self.context.frame.results_base_slot.as_ref().unwrap();
                    ensure!(
                        results_base_slot.addressed_from_sp(),
                        CodeGenError::sp_addressing_expected(),
                    );
                    let addr = self.masm.local_address(results_base_slot)?;
                    self.masm.store((*reg).into(), addr, (*ty).try_into()?)?;
                }
                // The result base parameter is a stack parameter, addressed
                // from FP.
                _ => {}
            }
        }

        self.masm.end_source_loc()?;

        Ok(CodeGen {
            sig: self.sig,
            context: self.context.for_emission(),
            masm: self.masm,
            env: self.env,
            tunables: self.tunables,
            source_location: self.source_location,
            control_frames: self.control_frames,
            fuel_consumed: self.fuel_consumed,
            needs_gc_heap: self.needs_gc_heap,
            gc_codegen_config: self.gc_codegen_config,
            phase: PhantomData,
        })
    }

    fn spill_register_arguments(&mut self) -> Result<()> {
        use WasmValType::*;
        for (operand, slot) in self
            .sig
            .params_without_retptr()
            .iter()
            .zip(self.context.frame.locals())
        {
            match (operand, slot) {
                (ABIOperand::Reg { ty, reg, .. }, slot) => {
                    let addr = self.masm.local_address(slot)?;
                    match &ty {
                        I32 | I64 | F32 | F64 | V128 => {
                            self.masm.store((*reg).into(), addr, (*ty).try_into()?)?;
                        }
                        Ref(rt) => match rt.heap_type {
                            WasmHeapType::Func => {
                                self.masm.store_ptr(*reg, addr)?;
                            }
                            WasmHeapType::Extern
                            | WasmHeapType::Exn
                            | WasmHeapType::ConcreteExn(_)
                            | WasmHeapType::NoExn => {
                                self.masm.store((*reg).into(), addr, (*ty).try_into()?)?;
                            }
                            _ => bail!(CodeGenError::unsupported_wasm_type()),
                        },
                    }
                }
                // Skip non-register arguments
                _ => {}
            }
        }
        Ok(())
    }

    /// Copy GC references passed on the stack into frame slots so that
    /// stack maps cover them: the caller's argument area is not visited by
    /// the collector, so a reference left there goes stale across a
    /// collection. Everything else stays in the caller's argument area.
    fn copy_stack_gc_refs_to_frame(&mut self) -> Result<()> {
        for (operand, slot) in self
            .sig
            .params_without_retptr()
            .iter()
            .zip(self.context.frame.locals())
        {
            match (operand, slot) {
                (ABIOperand::Stack { ty, offset, .. }, slot)
                    if ty.is_vmgcref_type_and_not_i31() =>
                {
                    ensure!(
                        slot.addressed_from_sp(),
                        CodeGenError::sp_addressing_expected(),
                    );
                    let arg_base = u32::from(<M::ABI as ABI>::arg_base_offset());
                    let src = LocalSlot::stack_arg(*ty, offset + arg_base);
                    let src_addr = self.masm.local_address(&src)?;
                    let dst_addr = self.masm.local_address(slot)?;
                    self.masm.with_scratch::<IntScratch, _>(|masm, scratch| {
                        masm.load(src_addr, scratch.writable(), (*ty).try_into()?)?;
                        masm.store(scratch.inner().into(), dst_addr, (*ty).try_into()?)
                    })?;
                }
                _ => {}
            }
        }
        Ok(())
    }
}

impl<'a, 'translation, 'data, M> CodeGen<'a, 'translation, 'data, M, Emission>
where
    M: MacroAssembler,
{
    /// Emit the function body to machine code.
    pub fn emit(
        &mut self,
        body: BinaryReader<'a>,
        validator: &mut FuncValidator<ValidatorResources>,
    ) -> Result<()> {
        self.emit_body(body, validator)
            .and_then(|_| self.emit_end())?;

        Ok(())
    }

    /// Pops a control frame from the control frame stack.
    pub fn pop_control_frame(&mut self) -> Result<ControlStackFrame> {
        let frame = self
            .control_frames
            .pop()
            .ok_or_else(|| format_err!(CodeGenError::control_frame_expected()))?;
        if let Some(info) = frame.try_table_info() {
            self.context
                .exception_handlers
                .restore_checkpoint(info.checkpoint);
        }

        Ok(frame)
    }

    /// Derives a [RelSourceLoc] from a [SourceLoc].
    pub fn source_loc_from(&mut self, loc: SourceLoc) -> RelSourceLoc {
        if self.source_location.base.is_none() && !loc.is_default() {
            self.source_location.base = Some(loc);
        }

        RelSourceLoc::from_base_offset(self.source_location.base.unwrap_or_default(), loc)
    }

    /// The following two helpers, handle else or end instructions when the
    /// compiler has entered into an unreachable code state. These instructions
    /// must be observed to determine if the reachability state should be
    /// restored.
    ///
    /// When the compiler is in an unreachable state, all the other instructions
    /// are not visited.
    pub fn handle_unreachable_else(&mut self) -> Result<()> {
        let frame = self
            .control_frames
            .last_mut()
            .ok_or_else(|| CodeGenError::control_frame_expected())?;
        ensure!(frame.is_if(), CodeGenError::if_control_frame_expected());
        if frame.is_next_sequence_reachable() {
            // We entered an unreachable state when compiling the
            // if-then branch, but if the `if` was reachable at
            // entry, the if-else branch will be reachable.
            self.context.reachable = true;
            frame.ensure_stack_state(self.masm, &mut self.context)?;
            frame.bind_else(self.masm, &mut self.context)?;
        }
        Ok(())
    }

    pub fn handle_unreachable_end(&mut self) -> Result<()> {
        let mut frame = self.pop_control_frame()?;
        if let Some(info) = frame.take_try_table_info() {
            return self.emit_try_table_end(frame, info);
        }
        // We just popped the outermost block.
        let is_outermost = self.control_frames.len() == 0;

        if frame.is_next_sequence_reachable() {
            self.context.reachable = true;
            frame.ensure_stack_state(self.masm, &mut self.context)?;
            frame.bind_end(self.masm, &mut self.context)
        } else if is_outermost {
            // If we reach the end of the function in an unreachable
            // state, perform the necessary cleanup to leave the stack
            // and SP in the expected state.  The compiler can enter
            // in this state through an infinite loop.
            frame.ensure_stack_state(self.masm, &mut self.context)
        } else {
            Ok(())
        }
    }

    fn emit_body(
        &mut self,
        body: BinaryReader<'a>,
        validator: &mut FuncValidator<ValidatorResources>,
    ) -> Result<()> {
        self.maybe_emit_fuel_check()?;

        self.maybe_emit_epoch_check()?;

        // Once we have emitted the epilogue and reserved stack space for the locals, we push the
        // base control flow block.
        self.control_frames.push(ControlStackFrame::block(
            BlockSig::from_sig(self.sig.clone()),
            self.masm,
            &mut self.context,
        )?);

        // Set the return area of the results *after* initializing the block. In
        // the function body block case, we'll treat the results as any other
        // case, addressed from the stack pointer, and when ending the function
        // the return area will be set to the return pointer.
        if self.sig.params.has_retptr() {
            self.sig
                .results
                .set_ret_area(RetArea::slot(self.context.frame.results_base_slot.unwrap()));
        }

        let mut ops = OperatorsReader::new(body);
        while !ops.eof() {
            let offset = ops.original_position();
            ops.visit_operator(&mut ValidateThenVisit(
                validator.simd_visitor(offset),
                self,
                offset,
            ))??;
        }
        ops.finish()?;
        return Ok(());

        struct ValidateThenVisit<'a, T, U>(T, &'a mut U, u64);

        macro_rules! validate_then_visit {
            ($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident $ann:tt)*) => {
                $(
                    fn $visit(&mut self $($(,$arg: $argty)*)?) -> Self::Output {
                        self.0.$visit($($($arg.clone()),*)?)?;
                        let op = Operator::$op $({ $($arg: $arg.clone()),* })?;
                        if self.1.visit(&op) {
                            self.1.before_visit_op(&op, self.2)?;
                            let res = self.1.$visit($($($arg),*)?)?;
                            self.1.after_visit_op()?;
                            Ok(res)
                        } else {
                            Ok(())
                        }
                    }
                )*
            };
        }

        fn visit_op_when_unreachable(op: &Operator) -> bool {
            use Operator::*;
            match op {
                If { .. } | Block { .. } | TryTable { .. } | Loop { .. } | Else | End => true,
                _ => false,
            }
        }

        /// Trait to handle hooks that must happen before and after visiting an
        /// operator.
        trait VisitorHooks {
            /// Hook prior to visiting an operator.
            fn before_visit_op(&mut self, operator: &Operator, offset: u64) -> Result<()>;
            /// Hook after visiting an operator.
            fn after_visit_op(&mut self) -> Result<()>;

            /// Returns `true` if the operator will be visited.
            ///
            /// Operators will be visited if the following invariants are met:
            /// * The compiler is in a reachable state.
            /// * The compiler is in an unreachable state, but the current
            ///   operator is a control flow operator. These operators need to be
            ///   visited in order to keep the control stack frames balanced and
            ///   to determine if the reachability state must be restored.
            fn visit(&self, op: &Operator) -> bool;
        }

        impl<'a, 'translation, 'data, M: MacroAssembler> VisitorHooks
            for CodeGen<'a, 'translation, 'data, M, Emission>
        {
            fn visit(&self, op: &Operator) -> bool {
                self.context.reachable || visit_op_when_unreachable(op)
            }

            fn before_visit_op(&mut self, operator: &Operator, offset: u64) -> Result<()> {
                // Handle source location mapping.
                self.source_location_before_visit_op(offset)?;

                // Handle fuel.
                if self.tunables.consume_fuel {
                    self.fuel_before_visit_op(operator)?;
                }
                Ok(())
            }

            fn after_visit_op(&mut self) -> Result<()> {
                // Handle source code location mapping.
                self.source_location_after_visit_op()
            }
        }

        impl<'a, T, U> VisitOperator<'a> for ValidateThenVisit<'_, T, U>
        where
            T: VisitSimdOperator<'a, Output = wasmparser::Result<()>>,
            U: VisitSimdOperator<'a, Output = Result<()>> + VisitorHooks,
        {
            type Output = U::Output;

            fn simd_visitor(
                &mut self,
            ) -> Option<&mut dyn VisitSimdOperator<'a, Output = Self::Output>>
            where
                T:,
            {
                Some(self)
            }

            wasmparser::for_each_visit_operator!(validate_then_visit);
        }

        impl<'a, T, U> VisitSimdOperator<'a> for ValidateThenVisit<'_, T, U>
        where
            T: VisitSimdOperator<'a, Output = wasmparser::Result<()>>,
            U: VisitSimdOperator<'a, Output = Result<()>> + VisitorHooks,
        {
            wasmparser::for_each_visit_simd_operator!(validate_then_visit);
        }
    }

    /// Whether a GC barrier must be emitted when writing or reading a
    /// reference of the given type through a collector-visible location.
    pub fn gc_barrier_needed(&self, ty: &WasmValType) -> bool {
        ty.is_vmgcref_type_and_not_i31()
            && self.tunables.collector
                == Some(wasmtime_environ::Collector::DeferredReferenceCounting)
    }

    /// Emits a a series of instructions that will type check a function reference call.
    pub fn emit_typecheck_funcref(
        &mut self,
        funcref_ptr: Reg,
        type_index: TypeIndex,
    ) -> Result<()> {
        let ptr_size: OperandSize = self.env.ptr_type().try_into()?;
        let sig_index_bytes = self.env.vmoffsets.size_of_vmshared_type_index();
        let sig_size = OperandSize::from_bytes(sig_index_bytes);
        let sig_index = self.env.translation.module.types[type_index].unwrap_module_type_index();
        let sig_offset = self.env.shared_type_index_offset(sig_index);
        let signatures_base_offset = self.env.vmoffsets.ptr.vmctx().type_ids();
        let funcref_sig_offset = self.env.vmoffsets.ptr.vm_func_ref().type_index();
        // Get the caller id.
        let caller_id = self.context.any_gpr(self.masm)?;

        self.masm.with_scratch::<IntScratch, _>(|masm, scratch| {
            // Load the signatures address into the scratch register.
            masm.load(
                masm.address_at_vmctx(signatures_base_offset.into())?,
                scratch.writable(),
                ptr_size,
            )?;

            masm.load(
                masm.address_at_reg(scratch.inner(), sig_offset)?,
                writable!(caller_id),
                sig_size,
            )
        })?;

        let callee_id = self.context.any_gpr(self.masm)?;
        self.masm.load(
            self.masm
                .address_at_reg(funcref_ptr, funcref_sig_offset.into())?,
            writable!(callee_id),
            sig_size,
        )?;

        // Typecheck.
        self.masm
            .cmp(caller_id, callee_id.into(), OperandSize::S32)?;
        self.masm.trapif(IntCmpKind::Ne, TRAP_BAD_SIGNATURE)?;
        self.context.free_reg(callee_id);
        self.context.free_reg(caller_id);
        wasmtime_environ::error::Ok(())
    }

    /// Emit the usual function end instruction sequence.
    fn emit_end(&mut self) -> Result<()> {
        // The implicit body block is treated a normal block (it pushes results
        // to the stack); so when reaching the end, we pop them taking as
        // reference the current function's signature.
        let base = SPOffset::from_u32(self.context.frame.locals_size);
        self.masm.start_source_loc(Default::default())?;
        if self.context.reachable {
            ControlStackFrame::pop_abi_results_impl(
                &mut self.sig.results,
                &mut self.context,
                self.masm,
                |results, _, _| Ok(results.ret_area().copied()),
            )?;
        } else {
            // If we reach the end of the function in an unreachable code state,
            // simply truncate to the expected values.
            // The compiler could enter this state through an infinite loop.
            self.context.truncate_stack_to(0)?;
            self.masm.reset_stack_pointer(base)?;
        }
        ensure!(
            self.context.stack.len() == 0,
            CodeGenError::unexpected_value_in_value_stack()
        );
        let stack_args_size = if self.sig.call_conv.is_default() {
            crate::abi::align_to(
                self.sig.params_stack_size(),
                u32::from(M::ABI::call_stack_align()),
            )
        } else {
            0
        };
        self.masm
            .epilogue(self.context.frame.locals_size, stack_args_size)?;
        self.masm.end_source_loc()?;
        Ok(())
    }

    /// Emits a select, using the declared result type for typed selects.
    pub fn emit_select(&mut self, ty: Option<WasmValType>) -> Result<()> {
        let cond = self.context.pop_to_reg(self.masm, None)?;
        let val2 = self.context.pop_to_reg(self.masm, None)?;
        let val1 = self.context.pop_to_reg(self.masm, None)?;
        let ty = ty.unwrap_or(val1.ty);
        self.masm.cmp(cond.reg, RegImm::i32(0), OperandSize::S32)?;
        // Conditionally move val1 to val2 if the comparison is not zero.
        self.masm.cmov(
            writable!(val2.into()),
            val1.into(),
            IntCmpKind::Ne,
            ty.try_into()?,
        )?;
        // A null reference can have an integer shadow type. The selected value
        // must retain the declared reference type so stack maps cover it.
        self.context.stack.push(TypedReg::new(ty, val2.reg).into());
        self.context.free_reg(val1.reg);
        self.context.free_reg(cond);

        Ok(())
    }

    /// Pops the value at the stack top and assigns it to the local at
    /// the given index, returning the typed register holding the
    /// source value.
    pub fn emit_set_local(&mut self, index: u32) -> Result<TypedReg> {
        // Materialize any references to the same local index that are in the
        // value stack by spilling.
        if self.context.stack.contains_latent_local(index) {
            self.context.spill(self.masm)?;
        }
        let src = self.context.pop_to_reg(self.masm, None)?;
        // Need to get address of local after `pop_to_reg` since `pop_to_reg`
        // will pop the machine stack causing an incorrect address to be
        // calculated.
        let (ty, addr) = self.context.frame.get_local_address(index, self.masm)?;
        self.masm
            .store(RegImm::reg(src.reg), addr, ty.try_into()?)?;

        Ok(src)
    }

    /// Loads the address of the given global.
    pub fn emit_get_global_addr(&mut self, index: GlobalIndex) -> Result<(WasmValType, Reg, u32)> {
        let data = self.env.resolve_global(index);

        if data.imported {
            let global_base = self.masm.address_at_reg(vmctx!(M), data.offset)?;
            let dst = self.context.any_gpr(self.masm)?;
            self.masm.load_ptr(global_base, writable!(dst))?;
            Ok((data.ty, dst, 0))
        } else {
            Ok((data.ty, vmctx!(M), data.offset))
        }
    }

    pub fn emit_table_get(&mut self, table_index: TableIndex) -> Result<()> {
        let table = self.env.table(table_index);
        let heap_type = table.ref_type.heap_type;
        ensure!(
            heap_type == WasmHeapType::Func,
            CodeGenError::unsupported_wasm_type()
        );
        ensure!(
            self.tunables.table_lazy_init,
            CodeGenError::unsupported_table_eager_init()
        );
        let table_data = self.env.resolve_table_data(table_index);
        let ptr_type = self.env.ptr_type();
        let builtin = self.env.builtins.table_get_lazy_init_func_ref::<M::ABI>()?;

        // Request the builtin's result register and use it to hold the table
        // element value. We preemptively spill and request this register to
        // avoid conflict at the control flow merge below. Requesting the result
        // register is safe since we know ahead-of-time the builtin's signature.
        self.context.spill(self.masm)?;
        let elem_value: Reg = self.context.reg(
            builtin.sig().results.unwrap_singleton().unwrap_reg(),
            self.masm,
        )?;

        let index = self.context.pop_to_reg(self.masm, None)?;
        let base = self.context.any_gpr(self.masm)?;

        let elem_addr = self.emit_compute_table_elem_addr(index.into(), base, &table_data)?;
        self.masm.load_ptr(elem_addr, writable!(elem_value))?;
        // Free the register used as base, once we have loaded the element
        // address into the element value register.
        self.context.free_reg(base);

        let (defined, cont) = (self.masm.get_label()?, self.masm.get_label()?);

        // Push the built-in arguments to the stack.
        self.context
            .stack
            .extend([table_index.as_u32().try_into().unwrap(), index.into()]);

        self.masm.branch(
            IntCmpKind::Ne,
            elem_value,
            elem_value.into(),
            defined,
            ptr_type.try_into()?,
        )?;
        // Free the element value register.
        // This is safe since the FnCall::emit call below, will ensure
        // that the result register is placed on the value stack.
        self.context.free_reg(elem_value);
        FnCall::emit::<M>(
            &mut self.env,
            self.masm,
            &mut self.context,
            Callee::Builtin(builtin.clone()),
        )?;

        // We know the signature of the libcall in this case, so we assert that there's
        // one element in the stack and that it's  the ABI signature's result register.
        let top = self
            .context
            .stack
            .peek()
            .ok_or_else(|| CodeGenError::missing_values_in_stack())?;
        let top = top.unwrap_reg();
        ensure!(
            top.reg == elem_value,
            CodeGenError::table_element_value_expected()
        );
        self.masm.jmp(cont)?;

        // In the defined case, mask the funcref address in place, by peeking into the
        // last element of the value stack, which was pushed by the `indirect` function
        // call above.
        //
        // Note that `FUNCREF_MASK` as type `usize` but here we want a 64-bit
        // value so assert its actual value and then use a `-2` literal.
        self.masm.bind(defined)?;
        assert_eq!(FUNCREF_MASK as isize, -2);
        let imm = RegImm::i64(-2);
        let dst = top.into();
        self.masm
            .and(writable!(dst), dst, imm, top.ty.try_into()?)?;

        self.masm.bind(cont)
    }

    /// Emit the `table.set` operation for a function-reference table.
    ///
    /// Expects the value stack to contain `[index, value]` (with `value` on
    /// top) and consumes both.
    pub fn emit_table_set(&mut self, table_index: TableIndex) -> Result<()> {
        let table = self.env.table(table_index);
        ensure!(
            table.ref_type.heap_type == WasmHeapType::Func,
            CodeGenError::unsupported_wasm_type()
        );
        ensure!(
            self.tunables.table_lazy_init,
            CodeGenError::unsupported_table_eager_init()
        );
        let ptr_type = self.env.ptr_type();
        let table_data = self.env.resolve_table_data(table_index);
        let value = self.context.pop_to_reg(self.masm, None)?;
        let index = self.context.pop_to_reg(self.masm, None)?;
        let base = self.context.any_gpr(self.masm)?;
        let elem_addr = self.emit_compute_table_elem_addr(index.into(), base, &table_data)?;
        // Set the initialized bit.
        self.masm.or(
            writable!(value.into()),
            value.into(),
            RegImm::i64(FUNCREF_INIT_BIT as i64),
            ptr_type.try_into()?,
        )?;

        self.masm.store_ptr(value.into(), elem_addr)?;

        self.context.free_reg(value);
        self.context.free_reg(index);
        self.context.free_reg(base);
        Ok(())
    }

    /// Emit the `table.grow` operation.
    pub fn emit_table_grow(&mut self, table_index: TableIndex) -> Result<()> {
        let ptr_type = self.env.ptr_type();
        let idx_type = self.env.table(table_index).idx_type;

        // Duplicate the `delta` argument on the stack since we'll need it at
        // the end if growth succeeds.
        let delta = self.context.pop_to_reg(self.masm, None)?;
        let tmp = self.context.any_gpr(self.masm)?;
        self.masm
            .mov(writable!(tmp), delta.reg.into(), delta.ty.try_into()?)?;
        self.context.stack.push(TypedReg::new(delta.ty, tmp).into());
        self.context.stack.push(delta.into());

        // Invoke the `table.grow` builtin on the host which will return whether
        // the growth succeeded, and if so where it's located.
        let at = self.context.stack.ensure_index_at(1)?;
        let builtin = self.env.builtins.table_grow::<M::ABI>()?;
        let builtin = self.prepare_builtin_defined_table_arg(table_index, at, builtin)?;
        FnCall::emit::<M>(&mut self.env, self.masm, &mut self.context, builtin)?;

        // Pop everything that's on the stack now. The builtin took `delta` and
        // pushed a result, and then peel off our duplicate of `delta` plus the
        // initialization element of `table.grow` itself.
        let result = self.context.pop_to_reg(self.masm, None)?;
        let len = self.context.pop_to_reg(self.masm, None)?;
        let init = self.context.pop_to_reg(self.masm, None)?;

        // Save a copy of `result` on the stack since we'll need it after
        // `table.fill` is done.
        let tmp_result = self.context.any_gpr(self.masm)?;
        self.masm.mov(
            writable!(tmp_result),
            result.reg.into(),
            result.ty.try_into()?,
        )?;
        self.context
            .stack
            .push(TypedReg::new(result.ty, tmp_result).into());

        // Test if the result of growth is -1. If it is, then we're done.
        // Otherwise fall through to `table.fill`.
        let done = self.masm.get_label()?;
        self.masm.branch(
            IntCmpKind::Eq,
            result.reg,
            RegImm::i64(-1),
            done,
            OperandSize::S64,
        )?;

        // Prepare the arguments for `table.fill` in the order the wasm
        // instruction expects.
        self.context.stack.push(result.into());
        self.context.stack.push(init.into());
        self.context.stack.push(len.into());
        self.emit_table_fill(table_index)?;

        self.masm.bind(done)?;

        // Similar to the memory.grow builtin, `table.grow` returns a
        // pointer, however, we need to ensure that the returned index
        // is representative of the address space for tables.
        match (ptr_type, idx_type) {
            (WasmValType::I64, IndexType::I64) => Ok(()),
            (WasmValType::I64, IndexType::I32) => {
                let top: Reg = self.context.pop_to_reg(self.masm, None)?.into();
                self.masm.wrap(writable!(top), top)?;
                self.context.stack.push(TypedReg::i32(top).into());
                Ok(())
            }

            _ => Err(format_err!(CodeGenError::unsupported_32_bit_platform())),
        }
    }

    /// Emit the `table.fill` operation.
    pub fn emit_table_fill(&mut self, table_index: TableIndex) -> Result<()> {
        // Put all of this opcode's arguments into registers.
        let len = self.context.pop_to_reg(self.masm, None)?;
        let init = self.context.pop_to_reg(self.masm, None)?;
        let offset = self.context.pop_to_reg(self.masm, None)?;

        // Perform a bounds check to see if `offset+len` is inbounds.
        let table_data = self.env.resolve_table_data(table_index);
        self.emit_compute_table_size(&table_data)?;
        let table_size = self.context.pop_to_reg(self.masm, None)?;
        let tmp = self.context.any_gpr(self.masm)?;
        let idx_size = table_data.index_type().try_into()?;
        self.masm.mov(writable!(tmp), offset.reg.into(), idx_size)?;
        self.masm.checked_uadd(
            writable!(tmp),
            tmp,
            len.reg.into(),
            idx_size,
            TRAP_TABLE_OUT_OF_BOUNDS,
        )?;
        self.masm.cmp(tmp, table_size.reg.into(), idx_size)?;
        self.masm
            .trapif(IntCmpKind::GtU, TRAP_TABLE_OUT_OF_BOUNDS)?;
        self.context.free_reg(tmp);
        self.context.free_reg(table_size);

        let header = self.masm.get_label()?;
        let exit = self.masm.get_label()?;

        self.masm.bind(header)?;

        // Exit the loop once there are no more elements to copy.
        self.masm.branch(
            IntCmpKind::Eq,
            len.reg,
            RegImm::i64(0),
            exit,
            OperandSize::S64,
        )?;

        // Duplicate `offset`, where we're writing, and `init` what we're
        // writing, into temporary registers. These are used by `emit_table_set`
        // below.
        let tmp_index = self.context.any_gpr(self.masm)?;
        let tmp_init = self.context.any_gpr(self.masm)?;
        self.masm
            .mov(writable!(tmp_index), offset.reg.into(), OperandSize::S64)?;
        self.masm
            .mov(writable!(tmp_init), init.reg.into(), OperandSize::S64)?;

        // Spill all this loop's variables onto the stack.
        self.context.stack.push(TypedReg::i64(len.reg).into());
        self.context.stack.push(TypedReg::i64(init.reg).into());
        self.context.stack.push(TypedReg::i64(offset.reg).into());

        // Emit `table.set`, consuming our temporary registers.
        self.context.stack.push(TypedReg::i64(tmp_index).into());
        self.context.stack.push(TypedReg::i64(tmp_init).into());
        self.emit_table_set(table_index)?;

        // Reload this loop's variables into the same registers as the start of
        // the loop.
        self.context.pop_to_reg(self.masm, Some(offset.reg))?;
        self.context.pop_to_reg(self.masm, Some(init.reg))?;
        self.context.pop_to_reg(self.masm, Some(len.reg))?;

        // Advance the destination we're writing to, and decrement the number of
        // elements left to write.
        self.masm.add(
            writable!(offset.reg),
            offset.reg,
            RegImm::i64(1),
            OperandSize::S64,
        )?;
        self.masm.sub(
            writable!(len.reg),
            len.reg,
            RegImm::i64(1),
            OperandSize::S64,
        )?;
        self.masm.jmp(header)?;

        self.masm.bind(exit)?;

        self.context.free_reg(offset);
        self.context.free_reg(init);
        self.context.free_reg(len);
        Ok(())
    }

    /// Emits a bounds check for the range `[idx, idx + len)` against the
    /// current size of `table_data`, trapping with `TRAP_TABLE_OUT_OF_BOUNDS`
    /// if the range is out-of-bounds.
    ///
    /// Both `idx` and `len` are expected to be 64-bit values.
    fn emit_table_range_bounds_check(
        &mut self,
        table_data: &TableData,
        idx: Reg,
        len: Reg,
    ) -> Result<()> {
        self.emit_compute_table_size(table_data)?;
        let size = self.context.pop_to_reg(self.masm, None)?;

        // Compute `end = idx + len`, trapping on overflow, and then trap if
        // `end > size`.
        let end = self.context.any_gpr(self.masm)?;
        self.masm
            .mov(writable!(end), idx.into(), OperandSize::S64)?;
        self.masm.checked_uadd(
            writable!(end),
            end,
            len.into(),
            OperandSize::S64,
            TRAP_TABLE_OUT_OF_BOUNDS,
        )?;
        self.masm.cmp(end, size.reg.into(), OperandSize::S64)?;
        self.masm
            .trapif(IntCmpKind::GtU, TRAP_TABLE_OUT_OF_BOUNDS)?;

        self.context.free_reg(size);
        self.context.free_reg(end);
        Ok(())
    }

    /// Emit the `table.copy` operation.
    pub fn emit_table_copy(&mut self, dst_table: TableIndex, src_table: TableIndex) -> Result<()> {
        let dst_data = self.env.resolve_table_data(dst_table);
        let src_data = self.env.resolve_table_data(src_table);

        // The value stack contains `[dst, src, len]` (top is `len`).
        let len = self.context.pop_to_reg(self.masm, None)?;
        let src = self.context.pop_to_reg(self.masm, None)?;
        let dst = self.context.pop_to_reg(self.masm, None)?;

        // Zero-extend each operand to a full 64-bit value so that the
        // arithmetic and bounds checks below can uniformly operate on 64-bit
        // quantities regardless of the table's index type.
        for op in [&len, &src, &dst] {
            if op.ty == WasmValType::I32 {
                self.masm.extend(
                    writable!(op.reg),
                    op.reg,
                    Extend::<Zero>::I64Extend32.into(),
                )?;
            }
        }

        // Bounds check both ranges up-front; `table.copy` traps without
        // copying anything if either range is out-of-bounds.
        self.emit_table_range_bounds_check(&src_data, src.reg, len.reg)?;
        self.emit_table_range_bounds_check(&dst_data, dst.reg, len.reg)?;

        // Decide the copy direction. If `dst <= src` then do a forwards copy
        // and otherwise it's backwards.
        let step = self.context.any_gpr(self.masm)?;
        let forward = self.masm.get_label()?;
        let setup_done = self.masm.get_label()?;
        self.masm.branch(
            IntCmpKind::LeU,
            dst.reg,
            src.reg.into(),
            forward,
            OperandSize::S64,
        )?;
        // Backwards: start at the last element and walk down.
        {
            self.masm
                .mov(writable!(step), RegImm::i64(-1), OperandSize::S64)?;
            self.masm.add(
                writable!(src.reg),
                src.reg,
                len.reg.into(),
                OperandSize::S64,
            )?;
            self.masm.sub(
                writable!(src.reg),
                src.reg,
                RegImm::i64(1),
                OperandSize::S64,
            )?;
            self.masm.add(
                writable!(dst.reg),
                dst.reg,
                len.reg.into(),
                OperandSize::S64,
            )?;
            self.masm.sub(
                writable!(dst.reg),
                dst.reg,
                RegImm::i64(1),
                OperandSize::S64,
            )?;
        }
        self.masm.jmp(setup_done)?;
        // Forwards: start at the first element and walk up.
        self.masm.bind(forward)?;
        {
            self.masm
                .mov(writable!(step), RegImm::i64(1), OperandSize::S64)?;
        }

        self.masm.bind(setup_done)?;

        let header = self.masm.get_label()?;
        let exit = self.masm.get_label()?;

        self.masm.bind(header)?;

        // Exit the loop once there are no more elements to copy.
        self.masm.branch(
            IntCmpKind::Eq,
            len.reg,
            RegImm::i64(0),
            exit,
            OperandSize::S64,
        )?;

        // Spill all loop variables to the stack for the body of the loop.
        // These will get reloaded back into the same registers at the end of
        // the loop.
        self.context.stack.push(TypedReg::i64(step).into());
        self.context.stack.push(TypedReg::i64(len.reg).into());
        self.context.stack.push(TypedReg::i64(dst.reg).into());
        self.context.stack.push(TypedReg::i64(src.reg).into());

        // Do a `table.get` followed by a `table.set`. Note that this'll redo
        // bounds checks which technically aren't necessary, but it's less code
        // duplication/complexity in Winch.
        //
        // Note that `dst` and `src` are on the stack and are needed for these
        // operations. They're also needed at the end of the loop, so some
        // stack-shuffling is necessary to "dup" the right values and get
        // everything in the expected shapes for `emit_table_{get,set}`.
        {
            let tmp_src = self.context.pop_to_reg(self.masm, None)?;
            let s = self.context.any_gpr(self.masm)?;
            self.masm
                .mov(writable!(s), tmp_src.reg.into(), OperandSize::S64)?;
            self.context.stack.push(tmp_src.into());
            self.context.stack.push(TypedReg::i64(s).into());
            self.emit_table_get(src_table)?;
            let funcref = self.context.pop_to_reg(self.masm, None)?;

            let tmp_src = self.context.pop_to_reg(self.masm, None)?;
            let tmp_dst = self.context.pop_to_reg(self.masm, None)?;

            let d = self.context.any_gpr(self.masm)?;
            self.masm
                .mov(writable!(d), tmp_dst.reg.into(), OperandSize::S64)?;
            self.context.stack.push(tmp_dst.into());
            self.context.stack.push(tmp_src.into());
            self.context.stack.push(TypedReg::i64(d).into());
            self.context.stack.push(funcref.into());
            self.emit_table_set(dst_table)?;
        }

        // Reload loop variables specifically back into the same registers to
        // ensure that modifications below are picked up on the next iteration.
        self.context.pop_to_reg(self.masm, Some(src.reg))?;
        self.context.pop_to_reg(self.masm, Some(dst.reg))?;
        self.context.pop_to_reg(self.masm, Some(len.reg))?;
        self.context.pop_to_reg(self.masm, Some(step))?;

        // Advance the running indices and decrement the remaining count.
        self.masm
            .add(writable!(dst.reg), dst.reg, step.into(), OperandSize::S64)?;
        self.masm
            .add(writable!(src.reg), src.reg, step.into(), OperandSize::S64)?;
        self.masm.sub(
            writable!(len.reg),
            len.reg,
            RegImm::i64(1),
            OperandSize::S64,
        )?;

        self.masm.jmp(header)?;

        self.masm.bind(exit)?;

        self.context.free_reg(src);
        self.context.free_reg(dst);
        self.context.free_reg(len);
        self.context.free_reg(step);
        Ok(())
    }

    /// Emits a series of instructions to bounds check and calculate the address
    /// of the given WebAssembly memory.
    /// This function returns a register containing the requested address.
    ///
    /// In essence, when computing the heap address for a WebAssembly load or
    /// store instruction the objective is to ensure that such access is safe,
    /// but also to perform the least amount of checks, and rely on the system to
    /// detect illegal memory accesses where applicable.
    ///
    /// Winch follows almost the same principles as Cranelift when it comes to
    /// bounds checks, for a more detailed explanation refer to
    /// prepare_addr in wasmtime-cranelift.
    ///
    /// Winch implementation differs in that, it defaults to the general case
    /// for dynamic heaps rather than optimizing for doing the least amount of
    /// work possible at runtime, this is done to align with Winch's principle
    /// of doing the least amount of work possible at compile time. For static
    /// heaps, Winch does a bit more of work, given that some of the cases that
    /// are checked against, can benefit compilation times, like for example,
    /// detecting an out of bounds access at compile time.
    pub fn emit_compute_heap_address(
        &mut self,
        heap: &HeapData,
        memarg: &MemArg,
        access_size: OperandSize,
    ) -> Result<Option<Reg>> {
        let ptr_size: OperandSize = self.env.ptr_type().try_into()?;
        let enable_spectre_mitigation = self.env.heap_access_spectre_mitigation();
        let add_offset_and_access_size = |offset: ImmOffset, access_size: OperandSize| {
            (access_size.bytes() as u64) + (offset.as_u32() as u64)
        };

        let index = Index::from_typed_reg(self.context.pop_to_reg(self.masm, None)?);

        let offset = bounds::ensure_index_and_offset(
            self.masm,
            index,
            memarg.offset,
            heap.index_type().try_into()?,
        )?;
        let offset_with_access_size = add_offset_and_access_size(offset, access_size);

        let memory_tunables = MemoryTunables::new(self.tunables, MemoryKind::LinearMemory);
        let can_elide_bounds_check = heap
            .memory
            .can_elide_bounds_check(&memory_tunables, self.env.page_size_log2);

        let addr = if offset_with_access_size > heap.memory.maximum_byte_size().unwrap_or(u64::MAX)
            || (!self.tunables.memory_may_move
                && offset_with_access_size > self.tunables.memory_reservation)
        {
            // Detect at compile time if the access is out of bounds.
            // Doing so will put the compiler in an unreachable code state,
            // optimizing the work that the compiler has to do until the
            // reachability is restored or when reaching the end of the
            // function.

            self.emit_fuel_increment()?;
            self.masm.trap(TrapCode::HEAP_OUT_OF_BOUNDS)?;
            self.context.reachable = false;
            None

        // Account for the case in which we can completely elide the bounds
        // checks.
        //
        // This case, makes use of the fact that if a memory access uses
        // a 32-bit index, then we be certain that
        //
        //      index <= u32::MAX
        //
        // Therefore if any 32-bit index access occurs in the region
        // represented by
        //
        //      bound + guard_size - (offset + access_size)
        //
        // We are certain that it's in bounds or that the underlying virtual
        // memory subsystem will report an illegal access at runtime.
        //
        // Note:
        //
        // * bound - (offset + access_size) cannot wrap, because it's checked
        // in the condition above.
        // * bound + heap.offset_guard_size is guaranteed to not overflow if
        // the heap configuration is correct, given that it's address must
        // fit in 64-bits.
        // * If the heap type is 32-bits, the offset is at most u32::MAX, so
        // no  adjustment is needed as part of
        // [bounds::ensure_index_and_offset].
        } else if can_elide_bounds_check
            && u64::from(u32::MAX)
                <= self.tunables.memory_reservation + self.tunables.memory_guard_size
                    - offset_with_access_size
        {
            assert!(can_elide_bounds_check);
            assert!(heap.index_type() == WasmValType::I32);
            let addr = self.context.any_gpr(self.masm)?;
            bounds::load_heap_addr_unchecked(self.masm, &heap, index, offset, addr, ptr_size)?;
            Some(addr)

        // Account for the case of a static memory size. The access is out
        // of bounds if:
        //
        // index > bound - (offset + access_size)
        //
        // bound - (offset + access_size) cannot wrap, because we already
        // checked that (offset + access_size) > bound, above.
        } else if let Some(static_size) = heap.memory.static_heap_size() {
            let bounds = Bounds::from_u64(static_size);
            let addr = bounds::load_heap_addr_checked(
                self.masm,
                &mut self.context,
                ptr_size,
                &heap,
                enable_spectre_mitigation,
                bounds,
                index,
                offset,
                |masm, bounds, index| {
                    let adjusted_bounds = bounds.as_u64() - offset_with_access_size;
                    let index_reg = index.as_typed_reg().reg;
                    masm.cmp(
                        index_reg,
                        RegImm::i64(adjusted_bounds as i64),
                        // Similar to the dynamic heap case, even though the
                        // offset and access size are bound through the heap
                        // type, when added they can overflow, resulting in
                        // an erroneous comparison, therefore we rely on the
                        // target pointer size.
                        ptr_size,
                    )?;
                    Ok(IntCmpKind::GtU)
                },
            )?;
            Some(addr)
        } else {
            // Account for the general case for bounds-checked memories. The
            // access is out of bounds if:
            // * index + offset + access_size overflows
            //   OR
            // * index + offset + access_size > bound
            let bounds = bounds::load_dynamic_heap_bounds::<_>(
                &mut self.context,
                self.masm,
                &heap,
                ptr_size,
            )?;

            let index_reg = index.as_typed_reg().reg;
            // Allocate a temporary register to hold
            //      index + offset + access_size
            //  which will serve as the check condition.
            let index_offset_and_access_size = self.context.any_gpr(self.masm)?;

            // Move the value of the index to the
            // index_offset_and_access_size register to perform the overflow
            // check to avoid clobbering the initial index value.
            //
            // We derive size of the operation from the heap type since:
            //
            // * This is the first assignment to the
            // `index_offset_and_access_size` register
            //
            // * The memory64 proposal specifies that the index is bound to
            // the heap type instead of hardcoding it to 32-bits (i32).
            self.masm.mov(
                writable!(index_offset_and_access_size),
                index_reg.into(),
                heap.index_type().try_into()?,
            )?;
            // Perform
            // index = index + offset + access_size, trapping if the
            // addition overflows.
            //
            // We use the target's pointer size rather than depending on the heap
            // type since we want to check for overflow; even though the
            // offset and access size are guaranteed to be bounded by the heap
            // type, when added, if used with the wrong operand size, their
            // result could be clamped, resulting in an erroneous overflow
            // check.
            self.masm.checked_uadd(
                writable!(index_offset_and_access_size),
                index_offset_and_access_size,
                RegImm::i64(offset_with_access_size as i64),
                ptr_size,
                TrapCode::HEAP_OUT_OF_BOUNDS,
            )?;

            let addr = bounds::load_heap_addr_checked(
                self.masm,
                &mut self.context,
                ptr_size,
                &heap,
                enable_spectre_mitigation,
                bounds,
                index,
                offset,
                |masm, bounds, _| {
                    let bounds_reg = bounds.as_typed_reg().reg;
                    masm.cmp(
                        index_offset_and_access_size,
                        bounds_reg.into(),
                        // We use the pointer size to keep the bounds
                        // comparison consistent with the result of the
                        // overflow check above.
                        ptr_size,
                    )?;
                    Ok(IntCmpKind::GtU)
                },
            )?;
            self.context.free_reg(bounds.as_typed_reg().reg);
            self.context.free_reg(index_offset_and_access_size);
            Some(addr)
        };

        self.context.free_reg(index.as_typed_reg().reg);
        Ok(addr)
    }

    /// Emit checks to ensure that the address at `memarg` is
    /// correctly aligned for the access size.
    fn emit_check_align(
        &mut self,
        heap: &HeapData,
        memarg: &MemArg,
        access_size: OperandSize,
    ) -> Result<()> {
        if access_size.bytes() > 1 {
            let heap_ty_size: OperandSize = heap.index_type().try_into()?;
            let addr = *self
                .context
                .stack
                .peek()
                .ok_or_else(|| CodeGenError::missing_values_in_stack())?;
            let tmp = self.context.any_gpr(self.masm)?;
            self.context.move_val_to_reg(&addr, tmp, self.masm)?;

            if memarg.offset != 0 {
                self.masm.add(
                    writable!(tmp),
                    tmp,
                    RegImm::Imm(Imm::I64(memarg.offset)),
                    heap_ty_size,
                )?;
            }

            self.masm.and(
                writable!(tmp),
                tmp,
                RegImm::Imm(Imm::I32(access_size.bytes() - 1)),
                heap_ty_size,
            )?;

            self.masm.cmp(tmp, RegImm::Imm(Imm::i64(0)), heap_ty_size)?;
            self.masm.trapif(IntCmpKind::Ne, TRAP_HEAP_MISALIGNED)?;
            self.context.free_reg(tmp);
        }

        Ok(())
    }

    pub fn emit_compute_heap_address_align_checked(
        &mut self,
        heap: &HeapData,
        memarg: &MemArg,
        access_size: OperandSize,
    ) -> Result<Option<Reg>> {
        self.emit_check_align(heap, memarg, access_size)?;
        self.emit_compute_heap_address(heap, memarg, access_size)
    }

    /// Emit a WebAssembly load.
    pub fn emit_wasm_load(
        &mut self,
        arg: &MemArg,
        target_type: WasmValType,
        kind: LoadKind,
    ) -> Result<()> {
        let emit_load = |this: &mut Self, dst, addr, kind| -> Result<()> {
            let src = this.masm.address_at_reg(addr, 0)?;
            this.masm.wasm_load(src, writable!(dst), kind)?;
            this.context
                .stack
                .push(TypedReg::new(target_type, dst).into());
            this.context.free_reg(addr);
            Ok(())
        };

        let memory_index = MemoryIndex::from_u32(arg.memory);
        let heap = self.env.resolve_heap(memory_index);

        // Ensure that the destination register is not allocated if
        // `emit_compute_heap_address` does not return an address.
        match kind {
            LoadKind::VectorLane(_) => {
                // Destination vector register is at the top of the stack and
                // `emit_compute_heap_address` expects an integer register
                // containing the address to load to be at the top of the stack.
                let dst = self.context.pop_to_reg(self.masm, None)?;
                let addr =
                    self.emit_compute_heap_address(&heap, &arg, kind.derive_operand_size())?;
                if let Some(addr) = addr {
                    emit_load(self, dst.reg, addr, kind)?;
                } else {
                    self.context.free_reg(dst);
                }
            }
            _ => {
                let maybe_addr = match kind {
                    LoadKind::Atomic(_, _) => self.emit_compute_heap_address_align_checked(
                        &heap,
                        &arg,
                        kind.derive_operand_size(),
                    )?,
                    _ => self.emit_compute_heap_address(&heap, &arg, kind.derive_operand_size())?,
                };

                if let Some(addr) = maybe_addr {
                    let dst = match target_type {
                        WasmValType::I32 | WasmValType::I64 => self.context.any_gpr(self.masm)?,
                        WasmValType::F32 | WasmValType::F64 => self.context.any_fpr(self.masm)?,
                        WasmValType::V128 => self.context.reg_for_type(target_type, self.masm)?,
                        _ => bail!(CodeGenError::unsupported_wasm_type()),
                    };

                    emit_load(self, dst, addr, kind)?;
                }
            }
        }

        Ok(())
    }

    /// Emit a WebAssembly store.
    pub fn emit_wasm_store(&mut self, arg: &MemArg, kind: StoreKind) -> Result<()> {
        let memory_index = MemoryIndex::from_u32(arg.memory);
        let heap = self.env.resolve_heap(memory_index);
        let src = self.context.pop_to_reg(self.masm, None)?;

        let maybe_addr = match kind {
            StoreKind::Atomic(size) => {
                self.emit_compute_heap_address_align_checked(&heap, &arg, size)?
            }
            StoreKind::Operand(size) | StoreKind::VectorLane(LaneSelector { size, .. }) => {
                self.emit_compute_heap_address(&heap, &arg, size)?
            }
        };

        if let Some(addr) = maybe_addr {
            self.masm
                .wasm_store(src.reg, self.masm.address_at_reg(addr, 0)?, kind)?;

            self.context.free_reg(addr);
        }
        self.context.free_reg(src);

        Ok(())
    }

    /// Loads the address of the table element at a given index. Returns the
    /// address of the table element using the provided register as base.
    pub fn emit_compute_table_elem_addr(
        &mut self,
        index: Reg,
        base: Reg,
        table_data: &TableData,
    ) -> Result<M::Address> {
        let bound = self.context.any_gpr(self.masm)?;
        let tmp = self.context.any_gpr(self.masm)?;
        let ptr_size: OperandSize = self.env.ptr_type().try_into()?;

        if let Some(offset) = table_data.import_from {
            // If the table data declares a particular offset base,
            // load the address into a register to further use it as
            // the table address.
            self.masm
                .load_ptr(self.masm.address_at_vmctx(offset)?, writable!(base))?;
        } else {
            // Else, simply move the vmctx register into the addr register as
            // the base to calculate the table address.
            self.masm.mov(writable!(base), vmctx!(M).into(), ptr_size)?;
        };

        // OOB check.
        let bound_addr = self
            .masm
            .address_at_reg(base, table_data.current_elems_offset)?;
        let bound_size = table_data.current_elements_size;
        self.masm.load(bound_addr, writable!(bound), bound_size)?;
        self.masm.cmp(index, bound.into(), bound_size)?;
        self.masm
            .trapif(IntCmpKind::GeU, TRAP_TABLE_OUT_OF_BOUNDS)?;

        // Move the index into the scratch register to calculate the table
        // element address.
        // Moving the value of the index register to the scratch register
        // also avoids overwriting the context of the index register.
        self.masm.with_scratch::<IntScratch, _>(|masm, scratch| {
            masm.mov(scratch.writable(), index.into(), bound_size)?;
            masm.mul(
                scratch.writable(),
                scratch.inner(),
                RegImm::i32(table_data.element_size.bytes() as i32),
                table_data.element_size,
            )?;
            masm.load_ptr(
                masm.address_at_reg(base, table_data.offset)?,
                writable!(base),
            )?;
            // Copy the value of the table base into a temporary register
            // so that we can use it later in case of a misspeculation.
            masm.mov(writable!(tmp), base.into(), ptr_size)?;
            // Calculate the address of the table element.
            masm.add(writable!(base), base, scratch.inner().into(), ptr_size)
        })?;
        if self.env.table_access_spectre_mitigation() {
            // Perform a bounds check and override the value of the
            // table element address in case the index is out of bounds.
            self.masm.cmp(index, bound.into(), bound_size)?;
            self.masm
                .cmov(writable!(base), tmp, IntCmpKind::GeU, ptr_size)?;
        }
        self.context.free_reg(bound);
        self.context.free_reg(tmp);
        self.masm.address_at_reg(base, 0)
    }

    /// Retrieves the size of the table, pushing the result to the value stack.
    pub fn emit_compute_table_size(&mut self, table_data: &TableData) -> Result<()> {
        let size = self.context.any_gpr(self.masm)?;
        let ptr_size: OperandSize = self.env.ptr_type().try_into()?;

        self.masm.with_scratch::<IntScratch, _>(|masm, scratch| {
            if let Some(offset) = table_data.import_from {
                masm.load_ptr(masm.address_at_vmctx(offset)?, scratch.writable())?;
            } else {
                masm.mov(scratch.writable(), vmctx!(M).into(), ptr_size)?;
            };

            let size_addr =
                masm.address_at_reg(scratch.inner(), table_data.current_elems_offset)?;
            masm.load(size_addr, writable!(size), table_data.current_elements_size)
        })?;

        let dst = TypedReg::new(table_data.index_type(), size);
        self.context.stack.push(dst.into());
        Ok(())
    }

    /// Retrieves the size of the memory, pushing the result to the value stack.
    fn load_memory_length(&mut self, heap_data: &HeapData, size_reg: Reg) -> Result<()> {
        self.masm.with_scratch::<IntScratch, _>(|masm, scratch| {
            let base = if let Some(offset) = heap_data.import_from {
                masm.load_ptr(masm.address_at_vmctx(offset)?, scratch.writable())?;
                scratch.inner()
            } else {
                vmctx!(M)
            };

            let size_addr = masm.address_at_reg(base, heap_data.current_length_offset)?;
            masm.load_ptr(size_addr, writable!(size_reg))
        })?;
        Ok(())
    }

    /// Retrieves the size of the memory, pushing the result to the value stack.
    pub fn emit_compute_memory_size(&mut self, heap_data: &HeapData) -> Result<()> {
        let size_reg = self.context.any_gpr(self.masm)?;
        self.load_memory_length(heap_data, size_reg)?;

        // Emit a shift to get the size in pages rather than in bytes.
        let dst = TypedReg::new(heap_data.index_type(), size_reg);
        let pow = heap_data.memory.page_size_log2;
        self.masm.shift_ir(
            writable!(dst.reg),
            Imm::i32(pow as i32),
            dst.into(),
            ShiftKind::ShrU,
            self.env.ptr_type().try_into()?,
        )?;
        self.context.stack.push(dst.into());
        Ok(())
    }

    /// Emit a bounds check for `ptr+len` and put the native address for this
    /// wasm address into `dst`.
    fn emit_bounds_check_and_compute_addr(
        &mut self,
        heap: &HeapData,
        dst: Reg,
        ptr: Reg,
        len: Reg,
    ) -> Result<()> {
        let ptr_size: OperandSize = self.env.ptr_type().try_into()?;
        let idx_size: OperandSize = heap.index_type().try_into()?;
        // Compute `dst = ptr + len` trapping on overflow. For an `i32` index
        // type the operands are zero-extended to 64-bit so overflow is
        // impossible.
        match idx_size {
            OperandSize::S32 => {
                self.masm
                    .extend(writable!(dst), ptr, Extend::<Zero>::I64Extend32.into())?;
                self.masm.add_uextend(
                    writable!(dst),
                    dst,
                    len,
                    OperandSize::S32,
                    OperandSize::S64,
                )?;
            }
            OperandSize::S64 => {
                self.masm
                    .mov(writable!(dst), ptr.into(), OperandSize::S64)?;
                self.masm.checked_uadd(
                    writable!(dst),
                    dst,
                    len.into(),
                    OperandSize::S64,
                    TrapCode::HEAP_OUT_OF_BOUNDS,
                )?;
            }
            _ => unreachable!(),
        }

        // Load the current size in bytes of the memory, and trap if
        // `dst > size_in_bytes`.
        let size_in_bytes = self.context.any_gpr(self.masm)?;
        self.load_memory_length(&heap, size_in_bytes)?;
        assert!(ptr_size == OperandSize::S64);
        self.masm.cmp(dst, size_in_bytes.into(), ptr_size)?;
        self.masm
            .trapif(IntCmpKind::GtU, TrapCode::HEAP_OUT_OF_BOUNDS)?;
        self.context.free_reg(size_in_bytes);

        // Compute `dst = memory_base + ptr`.
        bounds::load_heap_addr_unchecked(
            self.masm,
            &heap,
            Index::from_typed_reg(TypedReg::new(heap.index_type(), ptr)),
            ImmOffset::from_u32(0),
            dst,
            ptr_size,
        )?;
        Ok(())
    }

    /// Emit the `memory.copy` operation.
    pub fn emit_memory_copy(&mut self, dst_mem: MemoryIndex, src_mem: MemoryIndex) -> Result<()> {
        let dst_heap = self.env.resolve_heap(dst_mem);
        let src_heap = self.env.resolve_heap(src_mem);
        let dst_idx_size: OperandSize = dst_heap.index_type().try_into()?;
        let src_idx_size: OperandSize = src_heap.index_type().try_into()?;

        let len = self.context.pop_to_reg(self.masm, None)?;
        let src = self.context.pop_to_reg(self.masm, None)?;
        let dst = self.context.pop_to_reg(self.masm, None)?;

        // For 32-bit linear memories go ahead and make sure `len` is zero
        // extended within its register ensuring that the full 64-bits of the
        // register are defined. This assists in situations like cross-memory
        // copies where one memory is 32-bit and one is 64-bit and the same
        // register can be used for the length in both bounds checks below.
        if dst_idx_size == OperandSize::S32 || src_idx_size == OperandSize::S32 {
            self.masm.extend(
                writable!(len.reg),
                len.reg,
                Extend::<Zero>::I64Extend32.into(),
            )?;
        }

        let dst_raw_addr = self.context.any_gpr(self.masm)?;
        self.emit_bounds_check_and_compute_addr(&dst_heap, dst_raw_addr, dst.reg, len.reg)?;
        self.context.free_reg(dst);

        let src_raw_addr = self.context.any_gpr(self.masm)?;
        self.emit_bounds_check_and_compute_addr(&src_heap, src_raw_addr, src.reg, len.reg)?;
        self.context.free_reg(src);

        self.context
            .stack
            .push(TypedReg::new(self.env.ptr_type(), dst_raw_addr).into());
        self.context
            .stack
            .push(TypedReg::new(self.env.ptr_type(), src_raw_addr).into());
        self.context
            .stack
            .push(TypedReg::new(self.env.ptr_type(), len.reg).into());

        let builtin = self.env.builtins.memory_copy::<M::ABI>()?;
        FnCall::emit::<M>(
            &mut self.env,
            self.masm,
            &mut self.context,
            Callee::Builtin(builtin),
        )?;
        Ok(())
    }

    /// Emit the `memory.fill` operation.
    pub fn emit_memory_fill(&mut self, mem: MemoryIndex) -> Result<()> {
        let heap = self.env.resolve_heap(mem);
        let ptr_size: OperandSize = self.env.ptr_type().try_into()?;
        let idx_size: OperandSize = heap.index_type().try_into()?;

        // The wasm stack at this point is `[dst, val, len]`.
        let len = self.context.pop_to_reg(self.masm, None)?;
        let val = self.context.pop_to_reg(self.masm, None)?;
        let dst = self.context.pop_to_reg(self.masm, None)?;

        let raw_addr = self.context.any_gpr(self.masm)?;
        self.emit_bounds_check_and_compute_addr(&heap, raw_addr, dst.reg, len.reg)?;
        self.context.free_reg(dst);

        // The libcall takes the length as a host-pointer-sized integer, so
        // zero-extend if the wasm index type is smaller.
        let len_reg = len.reg;
        if idx_size == OperandSize::S32 && ptr_size == OperandSize::S64 {
            self.masm.extend(
                writable!(len_reg),
                len_reg,
                Extend::<Zero>::I64Extend32.into(),
            )?;
        }

        // Set up the call arguments: `[dst_ptr, val, len]`.
        self.context
            .stack
            .push(TypedReg::new(self.env.ptr_type(), raw_addr).into());
        self.context.stack.push(val.into());
        self.context
            .stack
            .push(TypedReg::new(self.env.ptr_type(), len_reg).into());

        let builtin = self.env.builtins.memory_fill::<M::ABI>()?;
        FnCall::emit::<M>(
            &mut self.env,
            self.masm,
            &mut self.context,
            Callee::Builtin(builtin),
        )?;
        Ok(())
    }

    /// Emit the `memory.init` operation.
    pub fn emit_memory_init(&mut self, segment: DataIndex, mem: MemoryIndex) -> Result<()> {
        let dst_heap = self.env.resolve_heap(mem);

        let len = self.context.pop_to_reg(self.masm, None)?;
        let src = self.context.pop_to_reg(self.masm, None)?;
        let dst = self.context.pop_to_reg(self.masm, None)?;

        // Make sure `len` is zero extended within its register ensuring that
        // the full 64-bits of the register are defined. This assists in
        // situations like cross-memory copies where one memory is 32-bit and
        // one is 64-bit and the same register can be used for the length in
        // both bounds checks below.
        self.masm.extend(
            writable!(len.reg),
            len.reg,
            Extend::<Zero>::I64Extend32.into(),
        )?;

        let dst_raw_addr = self.context.any_gpr(self.masm)?;
        self.emit_bounds_check_and_compute_addr(&dst_heap, dst_raw_addr, dst.reg, len.reg)?;
        self.context.free_reg(dst);

        let runtime_data_index = match self.env.translation.runtime_data_map[segment] {
            Some(i) => i,

            // Active data segments always have length zero, so this is only
            // valid of src and len are both zero.
            None => {
                self.masm.cmp(src.reg, RegImm::i32(0), OperandSize::S32)?;
                self.masm
                    .trapif(IntCmpKind::Ne, TrapCode::HEAP_OUT_OF_BOUNDS)?;
                self.masm.cmp(len.reg, RegImm::i32(0), OperandSize::S32)?;
                self.masm
                    .trapif(IntCmpKind::Ne, TrapCode::HEAP_OUT_OF_BOUNDS)?;
                self.context.free_reg(dst_raw_addr);
                self.context.free_reg(src);
                self.context.free_reg(len);
                return Ok(());
            }
        };

        // Bounds check this passive data segment. Load its
        // dynamically-specified length and see if that's in the range
        // of `src+len`.
        let data_segment_length_offset = self
            .env
            .vmoffsets
            .runtime_data_lengths()
            .at(runtime_data_index);
        let tmp1 = self.context.any_gpr(self.masm)?;
        let tmp2 = self.context.any_gpr(self.masm)?;
        self.masm.load(
            self.masm.address_at_vmctx(data_segment_length_offset)?,
            writable!(tmp1),
            OperandSize::S32,
        )?;
        self.masm
            .mov(writable!(tmp2), src.reg.into(), OperandSize::S32)?;
        self.masm.checked_uadd(
            writable!(tmp2),
            tmp2,
            len.reg.into(),
            OperandSize::S32,
            TrapCode::HEAP_OUT_OF_BOUNDS,
        )?;
        self.masm.cmp(tmp2, tmp1.into(), OperandSize::S32)?;
        self.masm
            .trapif(IntCmpKind::GtU, TrapCode::HEAP_OUT_OF_BOUNDS)?;
        self.context.free_reg(tmp2);

        // Calculate the src pointer by loading the base of the passive segment
        // and adding in the `src` offset.
        let data_segment_base_offset = self
            .env
            .vmoffsets
            .runtime_data_bases()
            .at(runtime_data_index);
        self.masm.load(
            self.masm.address_at_vmctx(data_segment_base_offset)?,
            writable!(tmp1),
            OperandSize::S64,
        )?;
        self.masm.add_uextend(
            writable!(tmp1),
            tmp1,
            src.reg,
            OperandSize::S32,
            OperandSize::S64,
        )?;
        self.context.free_reg(src);

        // And finally, the final step is calling the `memory_copy` libcall.
        self.context.stack.push(TypedReg::i64(dst_raw_addr).into());
        self.context.stack.push(TypedReg::i64(tmp1).into());
        self.context.stack.push(len.into());
        let builtin = self.env.builtins.memory_copy::<M::ABI>()?;
        FnCall::emit::<M>(
            &mut self.env,
            self.masm,
            &mut self.context,
            Callee::Builtin(builtin),
        )?;
        Ok(())
    }

    pub fn emit_data_drop(&mut self, data_index: DataIndex) -> Result<()> {
        let runtime_data_index = match self.env.translation.runtime_data_map[data_index] {
            Some(idx) => idx,
            // Active data segments do nothing when dropped, so this is a noop.
            None => return Ok(()),
        };
        let data_segment_offset = self
            .env
            .vmoffsets
            .runtime_data_lengths()
            .at(runtime_data_index);
        let len_addr = self.masm.address_at_vmctx(data_segment_offset)?;
        self.masm.store(RegImm::i32(0), len_addr, OperandSize::S32)
    }

    /// Implementation of `table.init`
    pub fn emit_table_init(
        &mut self,
        elem_index: ElemIndex,
        table_index: TableIndex,
    ) -> Result<()> {
        let builtin_base = self.env.builtins.passive_elem_segment_base::<M::ABI>()?;
        let builtin_len = self.env.builtins.passive_elem_segment_len::<M::ABI>()?;

        // Push the passive segment's length and base onto the stack.
        match self.env.translation.passive_elem_map[elem_index] {
            Some(idx) => {
                self.context.stack.extend([idx.as_u32().try_into()?]);
                FnCall::emit::<M>(
                    &mut self.env,
                    self.masm,
                    &mut self.context,
                    Callee::Builtin(builtin_len),
                )?;
                self.context.stack.extend([idx.as_u32().try_into()?]);
                FnCall::emit::<M>(
                    &mut self.env,
                    self.masm,
                    &mut self.context,
                    Callee::Builtin(builtin_base),
                )?;
            }
            // Active data segments have 0 length and a null base pointer.
            None => {
                let tmp = self.context.any_gpr(self.masm)?;
                self.masm
                    .mov(writable!(tmp), RegImm::i64(0), OperandSize::S64)?;
                self.context
                    .stack
                    .push(TypedReg::new(WasmValType::I64, tmp).into());

                let tmp = self.context.any_gpr(self.masm)?;
                self.masm
                    .mov(writable!(tmp), RegImm::i64(0), OperandSize::S64)?;
                self.context
                    .stack
                    .push(TypedReg::new(WasmValType::I64, tmp).into());
            }
        };

        // Push the table's current length onto the stack.
        let table_data = self.env.resolve_table_data(table_index);
        let idx_size = table_data.index_type().try_into()?;
        self.emit_compute_table_size(&table_data)?;

        // And now pop off everything we have for this instruction to work with
        // it all below.
        let table_size = self.context.pop_to_reg(self.masm, None)?;
        let segment_base = self.context.pop_to_reg(self.masm, None)?;
        let segment_len = self.context.pop_to_reg(self.masm, None)?;
        let len = self.context.pop_to_reg(self.masm, None)?;
        let segment_off = self.context.pop_to_reg(self.masm, None)?;
        let table_off = self.context.pop_to_reg(self.masm, None)?;

        // Zero-extend the length to make it easier to work with below for
        // 64-bit tables.
        if len.ty == WasmValType::I32 {
            self.masm.extend(
                writable!(len.reg),
                len.reg,
                Extend::<Zero>::I64Extend32.into(),
            )?;
        }

        // Perform a bounds check to see if `segment_off+len` is inbounds.
        let tmp = self.context.any_gpr(self.masm)?;
        {
            self.masm
                .mov(writable!(tmp), segment_off.reg.into(), OperandSize::S32)?;
            self.masm.checked_uadd(
                writable!(tmp),
                tmp,
                len.reg.into(),
                OperandSize::S32,
                TRAP_TABLE_OUT_OF_BOUNDS,
            )?;
            self.masm
                .cmp(tmp, segment_len.reg.into(), OperandSize::S32)?;
            self.masm
                .trapif(IntCmpKind::GtU, TRAP_TABLE_OUT_OF_BOUNDS)?;
            self.context.free_reg(segment_len);
        }

        // Perform a bounds check to see if `table_off+len` is inbounds.
        {
            self.masm
                .mov(writable!(tmp), table_off.reg.into(), idx_size)?;
            self.masm.checked_uadd(
                writable!(tmp),
                tmp,
                len.reg.into(),
                idx_size,
                TRAP_TABLE_OUT_OF_BOUNDS,
            )?;
            self.masm.cmp(tmp, table_size.reg.into(), idx_size)?;
            self.masm
                .trapif(IntCmpKind::GtU, TRAP_TABLE_OUT_OF_BOUNDS)?;
            self.context.free_reg(table_size);
        }
        self.context.free_reg(tmp);

        // Calculate the base address of the segment that we're reading from.
        {
            self.masm.extend(
                writable!(segment_off.reg),
                segment_off.reg,
                Extend::<Zero>::I64Extend32.into(),
            )?;
            self.masm.mul(
                writable!(segment_off.reg),
                segment_off.reg,
                RegImm::i64(16),
                OperandSize::S64,
            )?;
            self.masm.add(
                writable!(segment_base.reg),
                segment_base.reg,
                segment_off.reg.into(),
                OperandSize::S64,
            )?;
            self.context.free_reg(segment_off);
        }

        // Now run `table.set` in a loop with the values read from the element
        // segment.
        let header = self.masm.get_label()?;
        let exit = self.masm.get_label()?;

        self.masm.bind(header)?;
        {
            self.masm.branch(
                IntCmpKind::Eq,
                len.reg,
                RegImm::i64(0),
                exit,
                OperandSize::S64,
            )?;

            // Read `*mut VMFuncRef` from `ValRaw`, and then increment the
            // `segment_base` pointer.
            let funcref = self.context.any_gpr(self.masm)?;
            self.masm.load_ptr(
                self.masm.address_at_reg(segment_base.reg, 0)?,
                writable!(funcref),
            )?;
            self.masm.add(
                writable!(segment_base.reg),
                segment_base.reg,
                RegImm::i64(16),
                OperandSize::S64,
            )?;

            // Spill context/variables for the table.set, and note that
            // `table_off` is duplicated here as one version is consumed by the
            // `table.set` and the other persists across the loop.
            self.context.stack.push(segment_base.into());
            self.context.stack.push(len.into());
            let table_off_copy = self.context.any_gpr(self.masm)?;
            self.masm.mov(
                writable!(table_off_copy),
                table_off.reg.into(),
                table_off.ty.try_into()?,
            )?;
            self.context.stack.push(table_off.into());
            self.context
                .stack
                .push(TypedReg::new(table_off.ty, table_off_copy).into());
            self.context
                .stack
                .push(TypedReg::new(WasmValType::FUNCREF, funcref).into());
            self.emit_table_set(table_index)?;

            // Pop loop variables into their original registers for the loop.
            self.context.pop_to_reg(self.masm, Some(table_off.reg))?;
            self.context.pop_to_reg(self.masm, Some(len.reg))?;
            self.context.pop_to_reg(self.masm, Some(segment_base.reg))?;

            // Increment the table index to copy next
            self.masm.add(
                writable!(table_off.reg),
                table_off.reg,
                RegImm::i64(1),
                table_off.ty.try_into()?,
            )?;

            // Decrement the number of remaining elements to copy, used as the
            // loop's exit condition above.
            self.masm.sub(
                writable!(len.reg),
                len.reg,
                RegImm::i64(1),
                OperandSize::S64,
            )?;
        }
        self.masm.jmp(header)?;

        self.masm.bind(exit)?;

        self.context.free_reg(segment_base);
        self.context.free_reg(len);
        self.context.free_reg(table_off);
        Ok(())
    }

    /// Implementation of `elem.drop`
    pub fn emit_elem_drop(&mut self, elem_index: ElemIndex) -> Result<()> {
        let passive_elem_index = match self.env.translation.passive_elem_map[elem_index] {
            Some(idx) => idx,
            // Active elem segments do nothing when dropped, so this is a noop.
            None => return Ok(()),
        };
        let builtin = self.env.builtins.passive_elem_segment_drop::<M::ABI>()?;
        self.context
            .stack
            .extend([passive_elem_index.as_u32().try_into()?]);
        FnCall::emit::<M>(
            &mut self.env,
            self.masm,
            &mut self.context,
            Callee::Builtin(builtin),
        )?;
        self.context.pop_and_free(self.masm)
    }

    /// Checks if fuel consumption is enabled and emits a series of instructions
    /// that check the current fuel usage by performing a zero-comparison with
    /// the number of units stored in `VMStoreContext`.
    pub fn maybe_emit_fuel_check(&mut self) -> Result<()> {
        if !self.tunables.consume_fuel {
            return Ok(());
        }

        self.emit_fuel_increment()?;
        let out_of_fuel = self.env.builtins.out_of_gas::<M::ABI>()?;
        let fuel_reg = self.context.without::<Result<Reg>, M, _>(
            &out_of_fuel.sig().regs,
            self.masm,
            |cx, masm| cx.any_gpr(masm),
        )??;

        self.emit_load_fuel_consumed(fuel_reg)?;

        // The  continuation label if the current fuel is under the limit.
        let continuation = self.masm.get_label()?;

        // Spill locals and registers to avoid conflicts at the out-of-fuel
        // control flow merge.
        self.context.spill(self.masm)?;
        // Fuel is stored as a negative i64, so if the number is less than zero,
        // we're still under the fuel limits.
        self.masm.branch(
            IntCmpKind::LtS,
            fuel_reg,
            RegImm::i64(0),
            continuation,
            OperandSize::S64,
        )?;
        // Out-of-fuel branch.
        FnCall::emit::<M>(
            &mut self.env,
            self.masm,
            &mut self.context,
            Callee::Builtin(out_of_fuel.clone()),
        )?;
        self.context.pop_and_free(self.masm)?;

        // Under fuel limits branch.
        self.masm.bind(continuation)?;
        self.context.free_reg(fuel_reg);

        Ok(())
    }

    /// Emits a series of instructions that load the `fuel_consumed` field from
    /// `VMStoreContext`.
    fn emit_load_fuel_consumed(&mut self, fuel_reg: Reg) -> Result<()> {
        let store_context_offset = self.env.vmoffsets.ptr.vmctx().store_context();
        let fuel_offset = self.env.vmoffsets.ptr.vm_store_context().fuel_consumed();
        self.masm.load_ptr(
            self.masm
                .address_at_vmctx(u32::from(store_context_offset))?,
            writable!(fuel_reg),
        )?;

        self.masm.load(
            self.masm.address_at_reg(fuel_reg, u32::from(fuel_offset))?,
            writable!(fuel_reg),
            // Fuel is an i64.
            OperandSize::S64,
        )
    }

    /// Checks if epoch interruption is configured and emits a series of
    /// instructions that check the current epoch against its deadline.
    pub fn maybe_emit_epoch_check(&mut self) -> Result<()> {
        if !self.tunables.epoch_interruption {
            return Ok(());
        }

        // The continuation branch if the current epoch hasn't reached the
        // configured deadline.
        let cont = self.masm.get_label()?;
        let new_epoch = self.env.builtins.new_epoch::<M::ABI>()?;

        // Checks for runtime limits (e.g., fuel, epoch) are special since they
        // require inserting arbitrary function calls and control flow.
        // Special care must be taken to ensure that all invariants are met. In
        // this case, since `new_epoch` takes an argument and returns a value,
        // we must ensure that any registers used to hold the current epoch
        // value and deadline are not going to be needed later on by the
        // function call.
        let (epoch_deadline_reg, epoch_counter_reg) =
            self.context.without::<Result<(Reg, Reg)>, M, _>(
                &new_epoch.sig().regs,
                self.masm,
                |cx, masm| Ok((cx.any_gpr(masm)?, cx.any_gpr(masm)?)),
            )??;

        self.emit_load_epoch_deadline_and_counter(epoch_deadline_reg, epoch_counter_reg)?;

        // Spill locals and registers to avoid conflicts at the control flow
        // merge below.
        self.context.spill(self.masm)?;
        self.masm.branch(
            IntCmpKind::LtU,
            epoch_counter_reg,
            RegImm::reg(epoch_deadline_reg),
            cont,
            OperandSize::S64,
        )?;
        // Epoch deadline reached branch.
        FnCall::emit::<M>(
            &mut self.env,
            self.masm,
            &mut self.context,
            Callee::Builtin(new_epoch.clone()),
        )?;
        // `new_epoch` returns the new deadline. However we don't
        // perform any caching, so we simply drop this value.
        self.visit_drop()?;

        // Under epoch deadline branch.
        self.masm.bind(cont)?;

        self.context.free_reg(epoch_deadline_reg);
        self.context.free_reg(epoch_counter_reg);
        Ok(())
    }

    fn emit_load_epoch_deadline_and_counter(
        &mut self,
        epoch_deadline_reg: Reg,
        epoch_counter_reg: Reg,
    ) -> Result<()> {
        let epoch_ptr_offset = self.env.vmoffsets.ptr.vmctx().epoch_ptr();
        let store_context_offset = self.env.vmoffsets.ptr.vmctx().store_context();
        let epoch_deadline_offset = self.env.vmoffsets.ptr.vm_store_context().epoch_deadline();

        // Load the current epoch value into `epoch_counter_var`.
        self.masm.load_ptr(
            self.masm.address_at_vmctx(u32::from(epoch_ptr_offset))?,
            writable!(epoch_counter_reg),
        )?;

        // `epoch_deadline_var` contains the address of the value, so we need
        // to extract it.
        self.masm.load(
            self.masm.address_at_reg(epoch_counter_reg, 0)?,
            writable!(epoch_counter_reg),
            OperandSize::S64,
        )?;

        // Load the `VMStoreContext`.
        self.masm.load_ptr(
            self.masm
                .address_at_vmctx(u32::from(store_context_offset))?,
            writable!(epoch_deadline_reg),
        )?;

        self.masm.load(
            self.masm
                .address_at_reg(epoch_deadline_reg, u32::from(epoch_deadline_offset))?,
            writable!(epoch_deadline_reg),
            // The deadline value is a u64.
            OperandSize::S64,
        )
    }

    /// Increments the fuel consumed in `VMStoreContext` by flushing
    /// `self.fuel_consumed` to memory.
    fn emit_fuel_increment(&mut self) -> Result<()> {
        let fuel_at_point = std::mem::replace(&mut self.fuel_consumed, 0);
        if fuel_at_point == 0 {
            return Ok(());
        }

        let store_context_offset = self.env.vmoffsets.ptr.vmctx().store_context();
        let fuel_offset = self.env.vmoffsets.ptr.vm_store_context().fuel_consumed();
        let limits_reg = self.context.any_gpr(self.masm)?;

        // Load `VMStoreContext` into the `limits_reg` reg.
        self.masm.load_ptr(
            self.masm
                .address_at_vmctx(u32::from(store_context_offset))?,
            writable!(limits_reg),
        )?;

        self.masm.with_scratch::<IntScratch, _>(|masm, scratch| {
            // Load the fuel consumed at point into the scratch register.
            masm.load(
                masm.address_at_reg(limits_reg, u32::from(fuel_offset))?,
                scratch.writable(),
                OperandSize::S64,
            )?;

            // Add the fuel consumed at point with the value in the scratch
            // register.
            masm.add(
                scratch.writable(),
                scratch.inner(),
                RegImm::i64(fuel_at_point),
                OperandSize::S64,
            )?;

            // Store the updated fuel consumed to `VMStoreContext`.
            masm.store(
                scratch.inner().into(),
                masm.address_at_reg(limits_reg, u32::from(fuel_offset))?,
                OperandSize::S64,
            )
        })?;

        self.context.free_reg(limits_reg);

        Ok(())
    }

    /// Hook to handle fuel before visiting an operator.
    fn fuel_before_visit_op(&mut self, op: &Operator) -> Result<()> {
        if !self.context.reachable {
            // `self.fuel_consumed` must be correctly flushed to memory when
            // entering an unreachable state.
            ensure!(self.fuel_consumed == 0, CodeGenError::illegal_fuel_state());
            // Control operators are still visited to track nesting and restore
            // reachability at `else` or `end`, but those visits must not charge fuel.
            return Ok(());
        }

        // Generally, most instructions require 1 fuel unit.
        //
        // However, there are exceptions, which are detailed in the code below.
        // Note that the fuel accounting semantics align with those of
        // Cranelift; for further information, refer to
        // `crates/cranelift/src/func_environ.rs`.
        //
        // The primary distinction between the two implementations is that Winch
        // does not utilize a local-based cache to track fuel consumption.
        // Instead, each increase in fuel necessitates loading from and storing
        // to memory.
        //
        // Memory traffic will undoubtedly impact runtime performance. One
        // potential optimization is to designate a register as non-allocatable,
        // when fuel consumption is enabled, effectively using it as a local
        // fuel cache.
        self.fuel_consumed += self.tunables.operator_cost.cost(op);

        match op {
            Operator::Unreachable
            | Operator::Loop { .. }
            | Operator::If { .. }
            | Operator::Else { .. }
            | Operator::Br { .. }
            | Operator::BrIf { .. }
            | Operator::BrTable { .. }
            | Operator::End
            | Operator::Return
            | Operator::CallIndirect { .. }
            | Operator::Call { .. }
            | Operator::ReturnCall { .. }
            | Operator::ReturnCallIndirect { .. }
            | Operator::Throw { .. }
            | Operator::ThrowRef => self.emit_fuel_increment(),
            _ => Ok(()),
        }
    }

    // Hook to handle source location mapping before visiting an operator.
    fn source_location_before_visit_op(&mut self, offset: u64) -> Result<()> {
        let loc = SourceLoc::new(offset as u32);
        let rel = self.source_loc_from(loc);
        self.source_location.current = self.masm.start_source_loc(rel)?;
        Ok(())
    }

    // Hook to handle source location mapping after visiting an operator.
    fn source_location_after_visit_op(&mut self) -> Result<()> {
        // Because in Winch binary emission is done in a single pass
        // and because the MachBuffer performs optimizations during
        // emission, we have to be careful when calling
        // [`MacroAssembler::end_source_location`] to avoid breaking the
        // invariant that checks that the end [CodeOffset] must be equal
        // or greater than the start [CodeOffset].
        if self.masm.current_code_offset()? >= self.source_location.current.0 {
            self.masm.end_source_loc()?;
        }

        Ok(())
    }

    pub(crate) fn emit_atomic_rmw(
        &mut self,
        arg: &MemArg,
        op: RmwOp,
        size: OperandSize,
        extend: Option<Extend<Zero>>,
    ) -> Result<()> {
        let memory_index = MemoryIndex::from_u32(arg.memory);
        let heap = self.env.resolve_heap(memory_index);
        // We need to pop-push the operand to compute the address before passing control over to
        // masm, because some architectures may have specific requirements for the registers used
        // in some atomic operations. The computed address is pushed back to the context's stack
        // too, rather than handed over as a register, since registers that are not tracked by the
        // value stack can't be spilled, so an untracked address register would make any request
        // for a fixed register fail if the address happened to be allocated to it. For this
        // reason, the address is pushed as a register to be dereferenced prior to emission, after
        // all the ISA-specifc constraints have been solved.
        let operand = self.context.pop_to_reg(self.masm, None)?;
        if let Some(addr) = self.emit_compute_heap_address_align_checked(&heap, arg, size)? {
            self.context
                .stack
                .push(TypedReg::new(self.env.ptr_type(), addr).into());
            self.context.stack.push(operand.into());
            self.masm
                .atomic_rmw(&mut self.context, size, op, UNTRUSTED_FLAGS, extend)?;
        } else {
            // Ensure that the operand register is not left allocated if the access was proven to
            // be out of bounds at compile time.
            self.context.free_reg(operand);
        }

        Ok(())
    }

    pub(crate) fn emit_atomic_cmpxchg(
        &mut self,
        arg: &MemArg,
        size: OperandSize,
        extend: Option<Extend<Zero>>,
    ) -> Result<()> {
        // At this point in the stack we have:
        //    [ address, expected, replacement ]
        //
        // Therefore, emission for this instruction is a bit
        // trickier. The address for the CAS is the 3rd from the top
        // of the stack, and we must emit instruction to compute the
        // actual address with
        // `emit_compute_heap_address_align_checked`, while we still
        // have access to self. However, some ISAs have requirements
        // with regard to the registers used for some arguments, so we
        // need to pass the context to the masm. To solve this issue,
        // we pop the two first arguments from the stack, compute the
        // address, push back the address and the arguments, and hand
        // over the control to masm. The implementer of `atomic_cas`
        // can expect to find `address`, `expected` and `replacement`
        // at the top the context's stack.
        //
        // The computed address is pushed back to the stack as a
        // register, rather than handed over directly, so that the
        // register allocator is able to spill it if the target
        // requires a fixed register.

        let replacement = self.context.pop_to_reg(self.masm, None)?;
        let expected = self.context.pop_to_reg(self.masm, None)?;

        let memory_index = MemoryIndex::from_u32(arg.memory);
        let heap = self.env.resolve_heap(memory_index);
        if let Some(addr) = self.emit_compute_heap_address_align_checked(&heap, arg, size)? {
            self.context
                .stack
                .push(TypedReg::new(self.env.ptr_type(), addr).into());
            self.context.stack.push(expected.into());
            self.context.stack.push(replacement.into());

            self.masm
                .atomic_cas(&mut self.context, size, UNTRUSTED_FLAGS, extend)?;
        } else {
            // Ensure that the argument registers are not left allocated if the access was proven
            // to be out of bounds at compile time.
            self.context.free_reg(expected);
            self.context.free_reg(replacement);
        }
        Ok(())
    }

    /// Emit the sequence of instruction for a `memory.atomic.wait*`.
    pub fn emit_atomic_wait(&mut self, arg: &MemArg, kind: AtomicWaitKind) -> Result<()> {
        // The `memory_atomic_wait*` builtins expect the following arguments:
        // - `memory`, as u32
        // - `address`, as u64
        // - `expected`, as either u64 or u32
        // - `timeout`, as u64
        // At this point our stack only contains the `timeout`, the `expected` and the address, so
        // we need to:
        // - insert the memory as the first argument
        // - compute the actual memory offset from the `MemArg`, if necessary.
        // Note that the builtin function performs the alignment and bounds checks for us, so we
        // don't need to emit that.

        let timeout = self.context.pop_to_reg(self.masm, None)?;
        let expected = self.context.pop_to_reg(self.masm, None)?;
        let addr = self.context.pop_to_reg(self.masm, None)?;

        // Put the target memory index as the first argument.
        let stack_len = self.context.stack.len();
        let builtin = match kind {
            AtomicWaitKind::Wait32 => self.env.builtins.memory_atomic_wait32::<M::ABI>()?,
            AtomicWaitKind::Wait64 => self.env.builtins.memory_atomic_wait64::<M::ABI>()?,
        };
        let builtin = self.prepare_builtin_defined_memory_arg(
            MemoryIndex::from_u32(arg.memory),
            stack_len,
            builtin,
        )?;

        if arg.offset != 0 {
            self.masm.checked_uadd(
                writable!(addr.reg),
                addr.reg,
                RegImm::i64(arg.offset as i64),
                OperandSize::S64,
                TrapCode::HEAP_OUT_OF_BOUNDS,
            )?;
        }

        self.context
            .stack
            .push(TypedReg::new(WasmValType::I64, addr.reg).into());
        self.context.stack.push(expected.into());
        self.context.stack.push(timeout.into());

        FnCall::emit::<M>(&mut self.env, self.masm, &mut self.context, builtin)?;

        Ok(())
    }

    pub fn emit_atomic_notify(&mut self, arg: &MemArg) -> Result<()> {
        // The memory `memory_atomic_notify` builtin expects the following arguments:
        // - `memory`, as u32
        // - `address`, as u64
        // - `count`: as u32
        // At this point our stack only contains the `count` and the `address`, so we need to:
        // - insert the memory as the first argument
        // - compute the actual memory offset from the `MemArg`, if necessary.
        // Note that the builtin function performs the alignment and bounds checks for us, so we
        // don't need to emit that.

        // pop the arguments from the stack.
        let count = self.context.pop_to_reg(self.masm, None)?;
        let addr = self.context.pop_to_reg(self.masm, None)?;

        // Put the target memory index as the first argument.
        let builtin = self.env.builtins.memory_atomic_notify::<M::ABI>()?;
        let stack_len = self.context.stack.len();
        let builtin = self.prepare_builtin_defined_memory_arg(
            MemoryIndex::from_u32(arg.memory),
            stack_len,
            builtin,
        )?;

        if arg.offset != 0 {
            self.masm.checked_uadd(
                writable!(addr.reg),
                addr.reg,
                RegImm::i64(arg.offset as i64),
                OperandSize::S64,
                TrapCode::HEAP_OUT_OF_BOUNDS,
            )?;
        }

        // push remaining arguments.
        self.context
            .stack
            .push(TypedReg::new(WasmValType::I64, addr.reg).into());
        self.context.stack.push(count.into());

        FnCall::emit::<M>(&mut self.env, self.masm, &mut self.context, builtin)?;

        Ok(())
    }

    pub fn prepare_builtin_defined_memory_arg(
        &mut self,
        mem: MemoryIndex,
        defined_index_at: usize,
        builtin: BuiltinFunction,
    ) -> Result<Callee> {
        match self.env.translation.module.defined_memory_index(mem) {
            // This memory is defined in this module, so the vmctx is this
            // module's vmctx and the memory index is `defined` as returned here.
            Some(defined) => {
                self.context
                    .stack
                    .insert_many(defined_index_at, &[defined.as_u32().try_into()?]);
                Ok(Callee::Builtin(builtin))
            }

            // This memory is not defined in this module, so the defined index
            // is loaded from the `VMMemoryImport` and the vmctx is loaded from
            // the vmctx itself.
            None => {
                let vmimport = self.env.vmoffsets.imported_memories().at(mem);
                let vmctx_offset =
                    vmimport + u32::from(self.env.vmoffsets.ptr.vm_memory_import().vmctx());
                let index_offset =
                    vmimport + u32::from(self.env.vmoffsets.ptr.vm_memory_import().index());
                let index_addr = self.masm.address_at_vmctx(index_offset)?;
                let index_dst = self.context.reg_for_class(RegClass::Int, self.masm)?;
                self.masm
                    .load(index_addr, writable!(index_dst), OperandSize::S32)?;
                self.context
                    .stack
                    .insert_many(defined_index_at, &[Val::reg(index_dst, WasmValType::I32)]);
                Ok(Callee::BuiltinWithDifferentVmctx(builtin, vmctx_offset))
            }
        }
    }

    /// Same as `prepare_builtin_defined_memory_arg`, but for tables.
    pub fn prepare_builtin_defined_table_arg(
        &mut self,
        table: TableIndex,
        defined_index_at: usize,
        builtin: BuiltinFunction,
    ) -> Result<Callee> {
        match self.env.translation.module.defined_table_index(table) {
            Some(defined) => {
                self.context
                    .stack
                    .insert_many(defined_index_at, &[defined.as_u32().try_into()?]);
                Ok(Callee::Builtin(builtin))
            }
            None => {
                let vmimport = self.env.vmoffsets.imported_tables().at(table);
                let vmctx_offset =
                    vmimport + u32::from(self.env.vmoffsets.ptr.vm_table_import().vmctx());
                let index_offset =
                    vmimport + u32::from(self.env.vmoffsets.ptr.vm_table_import().index());
                let index_addr = self.masm.address_at_vmctx(index_offset)?;
                let index_dst = self.context.reg_for_class(RegClass::Int, self.masm)?;
                self.masm
                    .load(index_addr, writable!(index_dst), OperandSize::S32)?;
                self.context
                    .stack
                    .insert_many(defined_index_at, &[Val::reg(index_dst, WasmValType::I32)]);
                Ok(Callee::BuiltinWithDifferentVmctx(builtin, vmctx_offset))
            }
        }
    }
}

/// Returns the index of the [`ControlStackFrame`] for the given
/// depth.
pub fn control_index(depth: u32, control_length: usize) -> Result<usize> {
    (control_length - 1)
        .checked_sub(depth as usize)
        .ok_or_else(|| format_err!(CodeGenError::control_frame_expected()))
}