wasm/execution/runtime_structure/

value_stack.rs

use alloc::vec::Vec;
use core::mem::MaybeUninit;

use crate::{
    core::{
        fixed_capacity_vec::FixedCapacityVec, structure::modules::indices::LocalIdx,
        utils::ToUsizeExt,
    },
    execution::assert_validated::UnwrapValidatedExt,
    Config, FuncAddr, FuncType, RuntimeError, ValType, Value,
};

/// The stack at runtime containing
/// 1. Values
/// 2. Labels
/// 3. Activations
///
/// See <https://webassembly.github.io/spec/core/exec/runtime.html#stack>
#[derive(Debug, Clone)]
pub(crate) struct Stack {
    /// WASM values on the stack, i.e. the actual data that instructions operate on
    values: FixedCapacityVec<Value>,

    /// Call frames
    ///
    /// Each time a function is called, a new frame is pushed, whenever a function returns, a frame is popped
    frames: FixedCapacityVec<CallFrame>,
}

impl Stack {
    pub fn new<T: Config>(
        params_to_base_call_frame: Vec<Value>,
        base_call_frame_func_ty: &FuncType,
        base_call_frame_remaining_locals: &[ValType],
    ) -> Result<Self, RuntimeError> {
        let mut stack = Self {
            values: FixedCapacityVec::with_capacity(T::MAX_VALUE_STACK_SIZE),
            frames: FixedCapacityVec::with_capacity(T::MAX_CALL_STACK_SIZE),
        };

        stack.values.push_from_slice(&params_to_base_call_frame)?;

        debug_assert!(
            stack.values.len() >= base_call_frame_func_ty.params.valtypes.len(),
            "when pushing a new call frame, at least as many values need to be on the stack as required by the new call frames's function"
        );

        // the topmost `param_count` values are transferred into/consumed by this new call frame
        let param_count = base_call_frame_func_ty.params.valtypes.len();
        let call_frame_base_idx = stack.values.len() - param_count;

        // verify that enough space is on the stack to push the remaining locals
        if base_call_frame_remaining_locals.len() + stack.values.len() >= stack.values.capacity() {
            return Err(RuntimeError::StackExhaustion);
        }

        // after the params, put the additional locals
        for local in base_call_frame_remaining_locals {
            // SAFETY: the previous if checks that sufficient space is one the stack
            unsafe { stack.push_value_unchecked(Value::default_from_ty(*local)) };
        }

        // now that the locals are all populated, the actual stack section of this call frame begins
        let value_stack_base_idx = stack.values.len();

        stack.frames.push(CallFrame {
            return_func_addr: MaybeUninit::uninit(),
            return_addr: MaybeUninit::uninit(),
            value_stack_base_idx,
            call_frame_base_idx,
            return_value_count: base_call_frame_func_ty.returns.valtypes.len(),
            return_stp: MaybeUninit::uninit(),
        })?;

        Ok(stack)
    }

    pub fn into_values(mut self) -> Vec<Value> {
        self.values
            .pop_into_slice(self.values.len())
            .unwrap_validated()
            .to_vec()
    }

    /// Pop a value from the value stack
    #[inline(always)]
    pub fn pop_value(&mut self) -> Value {
        // If there is at least one call frame, we shall not pop values past the current
        // call frame. However, there is one legitimate reason to pop when there is **no** current
        // call frame: after the outermost function returns, to extract the final return values of
        // this interpreter invocation.
        debug_assert!(
            if !self.frames.is_empty() {
                self.values.len() > self.current_call_frame().value_stack_base_idx
            } else {
                true
            },
            "can not pop values past the current call frame"
        );

        debug_assert!(!self.values.is_empty(), "pop results in stack underflow");

        // SAFETY: this is safe only if there is at least one `Value` on the stack, which must
        // belong to the current `CallFrame`.
        unsafe { self.pop_value_unchecked() }
    }

    /// Pop a [`Value`] from the value stack
    ///
    /// # Safety
    ///
    /// This will underflow the stack and cause undefined behavior, if the stack is empty. To
    /// check, before pushing, assure that [`Self::values`] is not empty.
    #[inline(always)]
    unsafe fn pop_value_unchecked(&mut self) -> Value {
        // SAFETY: safe only if stack contains at least one element
        unsafe { self.values.pop_unchecked() }
    }

    /// Returns a cloned copy of the top value on the stack, or `None` if the stack is empty
    pub fn peek_value(&self) -> Option<Value> {
        self.values.peek().ok().cloned()
    }

    /// Push a value to the value stack after veryfing that this will not overflow the stack
    pub fn push_value(&mut self, value: Value) -> Result<(), RuntimeError> {
        // check for value stack exhaustion
        if self.values.len() >= self.values.capacity() {
            return Err(RuntimeError::StackExhaustion);
        }

        // SAFETY: The prior if statement verifies that the stack is not yet exhausted
        unsafe { self.push_value_unchecked(value) };

        Ok(())
    }

    /// Push a [`Value`] to the value stack
    ///
    /// # Safety
    ///
    /// This will overflow the stack and cause undefined behavior, if the stack is already full. To
    /// check, before pushing, assure that the [`Self::values`] is not full.
    #[inline(always)]
    unsafe fn push_value_unchecked(&mut self, value: Value) {
        // SAFETY: safe when called with remaining space on the stack.
        unsafe { self.values.push_unchecked(value) }
    }

    /// Returns a shared reference to a specific local by its index in the current call frame.
    pub fn get_local(&self, idx: LocalIdx) -> &Value {
        let idx = idx.into_inner().into_usize();
        let call_frame_base_idx = self.current_call_frame().call_frame_base_idx;
        self.values
            .get(call_frame_base_idx + idx)
            .unwrap_validated()
    }

    /// Returns a mutable reference to a specific local by its index in the current call frame.
    pub fn get_local_mut(&mut self, idx: LocalIdx) -> &mut Value {
        let idx = idx.into_inner().into_usize();
        let call_frame_base_idx = self.current_call_frame().call_frame_base_idx;
        self.values
            .get_mut(call_frame_base_idx + idx)
            .unwrap_validated()
    }

    /// Get a shared reference to the current [`CallFrame`]
    ///
    /// # Safety
    ///
    /// This will underflow if no active call frame is on the stack.
    pub fn current_call_frame(&self) -> &CallFrame {
        // SAFETY: must only be called if there is at least one callframe on the stack.
        unsafe { self.frames.peek_unchecked() }
    }

    /// Pop a [`CallFrame`] from the call stack, returning the caller function store address, return address, and the return stp
    ///
    /// Returns `None` if the base call frame was popped. Its information cannot
    /// be retrieved.
    pub fn pop_call_frame(&mut self) -> Option<(FuncAddr, usize, usize)> {
        let CallFrame {
            return_func_addr,
            return_addr,
            call_frame_base_idx,
            return_value_count,
            return_stp,
            ..
        } = self.frames.pop().unwrap_validated();

        let remove_count = self.values.len() - call_frame_base_idx - return_value_count;

        self.remove_in_between(remove_count, return_value_count);

        debug_assert_eq!(
            self.values.len(),
            call_frame_base_idx + return_value_count,
            "after a function call finished, the stack must have exactly as many values as it had before calling the function plus the number of function return values"
        );

        // If this was the base call frame, do not return its uninitialized fields
        if self.call_frame_count() == 0 {
            None
        } else {
            // SAFETY: This is sound, because we just checked that the call
            // frame which we just popped is not the base call frame and only
            // the base call frame can contain uninitialized fields.
            unsafe {
                Some((
                    return_func_addr.assume_init(),
                    return_addr.assume_init(),
                    return_stp.assume_init(),
                ))
            }
        }
    }

    /// Push a call frame to the call stack
    ///
    /// Takes the current [`Self::values`]'s length as [`CallFrame::value_stack_base_idx`].
    pub fn push_call_frame<C: Config>(
        &mut self,
        return_func_addr: FuncAddr,
        func_ty: &FuncType,
        remaining_locals: &[ValType],
        return_addr: usize,
        return_stp: usize,
    ) -> Result<(), RuntimeError> {
        // check for call stack exhaustion
        if self.call_frame_count() > C::MAX_CALL_STACK_SIZE {
            return Err(RuntimeError::StackExhaustion);
        }

        debug_assert!(
            self.values.len()>= func_ty.params.valtypes.len(),
            "when pushing a new call frame, at least as many values need to be on the stack as required by the new call frames's function"
        );

        // the topmost `param_count` values are transferred into/consumed by this new call frame
        let param_count = func_ty.params.valtypes.len();
        let call_frame_base_idx = self.values.len() - param_count;

        // verify that enough space is on the stack to push the remaining locals
        if remaining_locals.len() + self.values.len() >= self.values.capacity() {
            return Err(RuntimeError::StackExhaustion);
        }

        // after the params, put the additional locals
        for local in remaining_locals {
            // SAFETY: the previous if checks that sufficient space is one the stack
            unsafe { self.push_value_unchecked(Value::default_from_ty(*local)) };
        }

        // now that the locals are all populated, the actual stack section of this call frame begins
        let value_stack_base_idx = self.values.len();

        self.frames.push(CallFrame {
            return_func_addr: MaybeUninit::new(return_func_addr),
            return_addr: MaybeUninit::new(return_addr),
            value_stack_base_idx,
            call_frame_base_idx,
            return_value_count: func_ty.returns.valtypes.len(),
            return_stp: MaybeUninit::new(return_stp),
        })?;

        Ok(())
    }

    /// Returns how many call frames are on the stack, in total.
    pub fn call_frame_count(&self) -> usize {
        self.frames.len()
    }

    /// Pop `n` elements from the value stack's tail as an iterator, with the first element being
    /// closest to the **bottom** of the value stack
    ///
    /// Note that this is providing the values in reverse order compared to popping `n` values
    /// (which would yield the element closest to the **top** of the value stack first).
    pub fn pop_tail_iter(&mut self, n: usize) -> Vec<Value> {
        let tail = self.values.pop_into_slice(n).unwrap_validated().to_vec();
        debug_assert_eq!(tail.len(), n);
        tail
    }

    /// Remove `remove_count` values from the stack, keeping the topmost `keep_count` values
    ///
    /// From the stack, remove `remove_count` elements, by sliding down the `keep_count` topmost
    /// values `remove_count` positions.
    ///
    /// **Effects**
    ///
    /// - after the operation, [`Stack`] will contain `remove_count` fewer elements
    /// - `keep_count` topmost elements will be identical before and after the operation
    /// - all elements below the `remove_count + keep_count` topmost stack entry remain
    pub fn remove_in_between(&mut self, remove_count: usize, keep_count: usize) {
        self.values.remove_in_between(remove_count, keep_count);
    }
}

/// The [WASM spec](https://webassembly.github.io/spec/core/exec/runtime.html#stack) calls this `Activations`, however it refers to the call frames of functions.
#[derive(Debug, Clone, Copy)]
pub(crate) struct CallFrame {
    /// Store address of the function that called this [`CallFrame`]'s function
    ///
    /// This is always uninitialized for the base call frame and initialized for
    /// all other call frames.
    pub return_func_addr: MaybeUninit<FuncAddr>,

    /// Value that the PC has to be set to when this function returns
    ///
    /// This is always uninitialized for the base call frame and initialized for
    /// all other call frames.
    pub return_addr: MaybeUninit<usize>,

    /// The index to the lowermost value in [`Stack::values`] belonging to this [`CallFrame`]'s
    /// stack
    ///
    /// Values below this may still belong to this [`CallFrame`], but they are locals. Consequently,
    /// this is the lowest index down to which the stack may be popped in this [`CallFrame`].
    /// However, clearing up this [`CallFrame`] may require further popping, down to (and
    /// including!) the index stored in [`Self::call_frame_base_idx`].
    pub value_stack_base_idx: usize,

    /// The index to the lowermost value on [`Stack::values`] that belongs to this [`CallFrame`]
    ///
    /// Clearing this [`CallFrame`] requires popping all elements on [`Stack::values`] down to (and
    /// including!) this index.
    pub call_frame_base_idx: usize,

    /// Number of return values to retain on [`Stack::values`] when unwinding/popping a [`CallFrame`]
    pub return_value_count: usize,

    /// Value that the stp has to be set to when this function returns
    ///
    /// This is always uninitialized for the base call frame and initialized for
    /// all other call frames.
    pub return_stp: MaybeUninit<usize>,
}