wasmi 2.0.0-beta.8

use crate::limits::LimitsError;
#[cfg(doc)]
use crate::{ElementSegment, Func, Global, InstanceEntity, Memory, Table, memory::DataSegment};

/// Offsets within [`InstanceEntity::handles`] buffer for various handle types.
///
/// Objects in instances are layout out by-group in the following order:
///
/// 1. `memories`: Linear Memories
/// 2. `globals`: Global Variables
/// 3. `tables`: Tables
/// 4. `funcs`: Functions
/// 5. `elems`: Element Segments
/// 6. `datas`: Data Segments
///
/// # Note
///
/// - Linear memories must be placed first because Wasmi IR uses 16-bit memory addresses.
/// - Modules without `data_count` section do not allow to query the addresses of data segments.
#[derive(Debug, Copy, Clone)]
pub struct InstanceLayout {
    /// The start offset within `InstanceEntity::handles` for [`Memory`] handles.
    end_memories: u32,
    /// The start offset within `InstanceEntity::handles` for [`Global`] handles.
    end_globals: u32,
    /// The start offset within `InstanceEntity::handles` for [`Table`] handles.
    end_tables: u32,
    /// The start offset within `InstanceEntity::handles` for [`Func`] handles.
    end_funcs: u32,
    /// The start offset within `InstanceEntity::handles` for [`ElementSegment`] handles.
    end_elems: u32,
    /// The start offset within `InstanceEntity::handles` for [`DataSegment`] handles.
    ///
    /// This may be `0` if the `data_count` section does not exist.
    end_datas: u32,
}

macro_rules! define_addr_types {
    (
        $( pub struct $name:ident(u32) = $ty:ty );* $(;)?
    ) => {
        $(
            #[doc = concat!("The 32-bit address of a [`", stringify!($ty), "`] within the [`InstanceEntity::handles`] buffer.")]
            #[derive(Debug, Copy, Clone)]
            pub struct $name(u32);

            impl From<u32> for $name {
                #[inline]
                fn from(value: u32) -> Self {
                    Self(value)
                }
            }

            impl From<$name> for u32 {
                #[inline]
                fn from(addr: $name) -> Self {
                    addr.0
                }
            }
        )*
    };
}
define_addr_types! {
    pub struct GlobalAddr(u32) = Global;
    pub struct MemoryAddr(u32) = Memory;
    pub struct TableAddr(u32) = Table;
    pub struct FuncAddr(u32) = Func;
    pub struct DataAddr(u32) = DataSegment;
    pub struct ElemAddr(u32) = ElementSegment;
}

impl InstanceLayout {
    /// Creates a new [`InstanceLayoutBuilder`].
    pub(crate) fn build() -> InstanceLayoutBuilder {
        InstanceLayoutBuilder::default()
    }

    /// Creates an uninitialized [`InstanceLayout`].
    pub(crate) fn uninit() -> Self {
        Self {
            end_globals: 0,
            end_memories: 0,
            end_tables: 0,
            end_funcs: 0,
            end_elems: 0,
            end_datas: 0,
        }
    }

    /// Returns the number of linear memories in the associated instance.
    #[inline]
    fn len_memories(&self) -> u32 {
        // Note: memories are placed first.
        self.end_memories
    }

    /// Returns the number of global variables in the associated instance.
    #[inline]
    fn len_globals(&self) -> u32 {
        // Note: globals are placed directly after memories.
        self.end_globals - self.end_memories
    }

    /// Returns the number of tables in the associated instance.
    #[inline]
    fn len_tables(&self) -> u32 {
        // Note: tables are placed directly after global variables.
        self.end_tables - self.end_globals
    }

    /// Returns the number of functions in the associated instance.
    #[inline]
    fn len_funcs(&self) -> u32 {
        // Note: functions are placed directly after tables.
        self.end_funcs - self.end_tables
    }

    /// Returns the number of element segments in the associated instance.
    #[inline]
    fn len_elems(&self) -> u32 {
        // Note: element segments are placed directly after functions.
        self.end_elems - self.end_funcs
    }

    /// Returns the number of data segments in the associated instance.
    #[inline]
    fn len_datas(&self) -> u32 {
        // Note: data segments are placed directly after element segments.
        self.end_datas - self.end_elems
    }

    /// Returns the [`MemoryAddr`] for a [`Memory`] Wasm index.
    #[inline]
    pub fn memory_addr(&self, index: u32) -> Option<MemoryAddr> {
        if index >= self.len_memories() {
            return None;
        }
        // Note: memories are placed first in the instance.
        Some(MemoryAddr(index))
    }

    /// Returns the [`GlobalAddr`] for a [`Global`] Wasm index.
    #[inline]
    pub fn global_addr(&self, index: u32) -> Option<GlobalAddr> {
        if index >= self.len_globals() {
            return None;
        }
        // Note: globals are placed directly after memories.
        Some(GlobalAddr(self.end_memories + index))
    }

    /// Returns the [`TableAddr`] for a [`Table`] Wasm index.
    #[inline]
    pub fn table_addr(&self, index: u32) -> Option<TableAddr> {
        if index >= self.len_tables() {
            return None;
        }
        // Note: tables are placed directly after globals.
        Some(TableAddr(self.end_globals + index))
    }

    /// Returns the [`FuncAddr`] for a [`Func`] Wasm index.
    #[inline]
    pub fn func_addr(&self, index: u32) -> Option<FuncAddr> {
        if index >= self.len_funcs() {
            return None;
        }
        // Note: functions are placed directly after tables.
        Some(FuncAddr(self.end_tables + index))
    }

    /// Returns the [`ElemAddr`] for a [`ElementSegment`] Wasm index.
    #[inline]
    pub fn elem_addr(&self, index: u32) -> Option<ElemAddr> {
        if index >= self.len_elems() {
            return None;
        }
        // Note: element segments are placed directly after functions.
        Some(ElemAddr(self.end_funcs + index))
    }

    /// Returns the [`DataAddr`] for a [`DataSegment`] Wasm index.
    #[inline]
    pub fn data_addr(&self, index: u32) -> Option<DataAddr> {
        if index >= self.len_datas() {
            return None;
        }
        // Note: data segments are placed directly after element segments.
        Some(DataAddr(self.end_elems + index))
    }
}

#[derive(Debug, Default)]
pub struct InstanceLayoutBuilder {
    /// The start offset within `InstanceEntity::handles` for [`Memory`] handles.
    memories: Option<u32>,
    /// The start offset within `InstanceEntity::handles` for [`Global`] handles.
    globals: Option<u32>,
    /// The start offset within `InstanceEntity::handles` for [`Table`] handles.
    tables: Option<u32>,
    /// The start offset within `InstanceEntity::handles` for [`DataSegment`] handles.
    datas: Option<u32>,
    /// The start offset within `InstanceEntity::handles` for [`ElementSegment`] handles.
    elems: Option<u32>,
    /// The start offset within `InstanceEntity::handles` for [`Func`] handles.
    funcs: Option<u32>,
}

macro_rules! impl_builder {
    (
        $( pub fn $name:ident(&mut self, $len:ident: usize) -> Result<&mut Self, LimitsError> = ($ty:ident, $check_limit:expr));*
        $(;)?
    ) => {
        $(
            #[doc = concat!("Initializes the number of [`", stringify!($ty), "`] handles in `self`.")]
            pub fn $name(&mut self, $len: usize) -> Result<&mut Self, LimitsError> {
                assert!(self.$name.is_none());
                $check_limit($len)?;
                let $name = $len as u32;
                self.$name = Some($name);
                Ok(self)
            }
        )*
    };
}
impl InstanceLayoutBuilder {
    impl_builder! {
        pub fn memories(&mut self, len_memories: usize) -> Result<&mut Self, LimitsError> = (Memory, LimitsError::max_memory_count);
        pub fn globals(&mut self, len_globals: usize) -> Result<&mut Self, LimitsError> = (Global, LimitsError::max_global_count);
        pub fn tables(&mut self, len_tables: usize) -> Result<&mut Self, LimitsError> = (Table, LimitsError::max_table_count);
        pub fn datas(&mut self, len_datas: usize) -> Result<&mut Self, LimitsError> = (DataSegment, LimitsError::max_data_count);
        pub fn elems(&mut self, len_elems: usize) -> Result<&mut Self, LimitsError> = (ElementSegment, LimitsError::max_elem_count);
        pub fn funcs(&mut self, len_funcs: usize) -> Result<&mut Self, LimitsError> = (Func, LimitsError::max_func_count);
    }

    pub fn finish(self) -> Result<InstanceLayout, LimitsError> {
        let err = LimitsError::TooManyInstanceHandles;
        let len_memories = self.memories.unwrap_or(0);
        let len_globals = self.globals.unwrap_or(0);
        let len_tables = self.tables.unwrap_or(0);
        let len_datas = self.datas.unwrap_or(0);
        let len_elems = self.elems.unwrap_or(0);
        let len_funcs = self.funcs.unwrap_or(0);
        let end_memories = len_memories;
        let end_globals = end_memories.checked_add(len_globals).ok_or(err)?;
        let end_tables = end_globals.checked_add(len_tables).ok_or(err)?;
        let end_funcs = end_tables.checked_add(len_funcs).ok_or(err)?;
        let end_elems = end_funcs.checked_add(len_elems).ok_or(err)?;
        let end_datas = end_elems.checked_add(len_datas).ok_or(err)?;
        Ok(InstanceLayout {
            end_memories,
            end_globals,
            end_tables,
            end_funcs,
            end_elems,
            end_datas,
        })
    }
}