roc/crates/wasm_interp/src/value_stack.rs

use bitvec::vec::BitVec;
use bumpalo::{collections::Vec, Bump};
use roc_wasm_module::ValueType;
use std::{fmt::Debug, mem::size_of};

use crate::Value;

/// Memory-efficient Struct-of-Arrays storage for the value stack.
/// Pack the values and their types as densely as possible,
/// to get better cache usage, at the expense of some extra logic.
pub struct ValueStack<'a> {
    bytes: Vec<'a, u8>,
    is_float: BitVec,
    is_64: BitVec,
}

macro_rules! pop_bytes {
    ($ty: ty, $bytes: expr) => {{
        const SIZE: usize = size_of::<$ty>();
        let bytes_idx = $bytes.len() - SIZE;
        let mut b = [0; SIZE];
        b.copy_from_slice(&$bytes[bytes_idx..][..SIZE]);
        $bytes.truncate(bytes_idx);
        <$ty>::from_ne_bytes(b)
    }};
}

impl<'a> ValueStack<'a> {
    pub fn new(arena: &'a Bump) -> Self {
        ValueStack {
            bytes: Vec::with_capacity_in(1024, arena),
            is_float: BitVec::with_capacity(1024),
            is_64: BitVec::with_capacity(1024),
        }
    }

    pub fn push(&mut self, value: Value) {
        match value {
            Value::I32(x) => {
                self.bytes.extend_from_slice(&x.to_ne_bytes());
                self.is_float.push(false);
                self.is_64.push(false);
            }
            Value::I64(x) => {
                self.bytes.extend_from_slice(&x.to_ne_bytes());
                self.is_float.push(false);
                self.is_64.push(true);
            }
            Value::F32(x) => {
                self.bytes.extend_from_slice(&x.to_ne_bytes());
                self.is_float.push(true);
                self.is_64.push(false);
            }
            Value::F64(x) => {
                self.bytes.extend_from_slice(&x.to_ne_bytes());
                self.is_float.push(true);
                self.is_64.push(true);
            }
        }
    }

    pub fn pop(&mut self) -> Value {
        let is_64 = self.is_64.pop().unwrap();
        let is_float = self.is_float.pop().unwrap();
        let size = if is_64 { 8 } else { 4 };
        let bytes_idx = self.bytes.len() - size;
        let value = self.get(is_64, is_float, bytes_idx);
        self.bytes.truncate(bytes_idx);
        value
    }

    fn get(&self, is_64: bool, is_float: bool, bytes_idx: usize) -> Value {
        if is_64 {
            let mut b = [0; 8];
            b.copy_from_slice(&self.bytes[bytes_idx..][..8]);
            if is_float {
                Value::F64(f64::from_ne_bytes(b))
            } else {
                Value::I64(i64::from_ne_bytes(b))
            }
        } else {
            let mut b = [0; 4];
            b.copy_from_slice(&self.bytes[bytes_idx..][..4]);
            if is_float {
                Value::F32(f32::from_ne_bytes(b))
            } else {
                Value::I32(i32::from_ne_bytes(b))
            }
        }
    }

    pub fn pop_i32(&mut self) -> i32 {
        match (self.is_float.pop(), self.is_64.pop()) {
            (Some(false), Some(false)) => pop_bytes!(i32, self.bytes),
            (Some(is_float), Some(is_64)) => panic!(
                "Expected I32 but found {:?}",
                type_from_flags(is_float, is_64)
            ),
            _ => panic!("Expected I32 but value stack was empty"),
        }
    }

    pub fn pop_i64(&mut self) -> i64 {
        match (self.is_float.pop(), self.is_64.pop()) {
            (Some(false), Some(true)) => pop_bytes!(i64, self.bytes),
            (Some(is_float), Some(is_64)) => panic!(
                "Expected I64 but found {:?}",
                type_from_flags(is_float, is_64)
            ),
            _ => panic!("Expected I64 but value stack was empty"),
        }
    }

    pub fn pop_f32(&mut self) -> f32 {
        match (self.is_float.pop(), self.is_64.pop()) {
            (Some(true), Some(false)) => pop_bytes!(f32, self.bytes),
            (Some(is_float), Some(is_64)) => panic!(
                "Expected F32 but found {:?}",
                type_from_flags(is_float, is_64)
            ),
            _ => panic!("Expected F32 but value stack was empty"),
        }
    }

    pub fn pop_f64(&mut self) -> f64 {
        match (self.is_float.pop(), self.is_64.pop()) {
            (Some(true), Some(true)) => pop_bytes!(f64, self.bytes),
            (Some(is_float), Some(is_64)) => panic!(
                "Expected F64 but found {:?}",
                type_from_flags(is_float, is_64)
            ),
            _ => panic!("Expected F64 but value stack was empty"),
        }
    }
}

fn type_from_flags(is_float: bool, is_64: bool) -> ValueType {
    match (is_float, is_64) {
        (false, false) => ValueType::I32,
        (false, true) => ValueType::I64,
        (true, false) => ValueType::F32,
        (true, true) => ValueType::F64,
    }
}

impl Debug for ValueStack<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "[")?;
        let mut index = 0;
        assert_eq!(self.is_64.len(), self.is_float.len());
        let iter_64 = self.is_64.iter().by_vals();
        let iter_float = self.is_float.iter().by_vals();
        for (i, (is_64, is_float)) in iter_64.zip(iter_float).enumerate() {
            let value = self.get(is_64, is_float, index);
            index += if is_64 { 8 } else { 4 };
            value.fmt(f)?;
            if i < self.is_64.len() - 1 {
                write!(f, ", ")?;
            }
        }
        write!(f, "]")
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const VALUES: [Value; 4] = [
        Value::I32(123),
        Value::I64(123456),
        Value::F32(3.14),
        Value::F64(-1.1),
    ];

    #[test]
    fn test_push_pop() {
        let arena = Bump::new();
        let mut stack = ValueStack::new(&arena);

        for val in VALUES {
            stack.push(val);
        }

        for val in VALUES.iter().rev() {
            let popped = stack.pop();
            assert_eq!(popped, *val);
        }
    }

    #[test]
    fn test_debug_fmt() {
        let arena = Bump::new();
        let mut stack = ValueStack::new(&arena);

        for val in VALUES {
            stack.push(val);
        }

        assert_eq!(format!("{:?}", VALUES), format!("{:?}", stack));
    }
}