wasmi 2.0.0-beta.8

use super::InstanceEntity;
use crate::{
    ElementSegment,
    Error,
    Extern,
    ExternType,
    Func,
    Global,
    Memory,
    Module,
    Table,
    collections::Map,
    instance::{AnyHandleAndEntity, InstanceLayout, handle::AnyHandle},
    memory::DataSegment,
    module::FuncIdx,
};
use alloc::{boxed::Box, vec::Vec};
use core::iter::FusedIterator;

/// A module instance entity builder.
#[derive(Debug)]
pub struct InstanceEntityBuilder {
    tables: Vec<Table>,
    funcs: Vec<Func>,
    memories: Vec<Memory>,
    globals: Vec<Global>,
    start_fn: Option<FuncIdx>,
    exports: Map<Box<str>, Extern>,
    data_segments: Vec<DataSegment>,
    elem_segments: Vec<ElementSegment>,
    layout: InstanceLayout,
}

impl InstanceEntityBuilder {
    /// Creates a new [`InstanceEntityBuilder`] optimized for the [`Module`].
    pub fn new(module: &Module) -> Self {
        fn vec_with_capacity_exact<T>(capacity: usize) -> Vec<T> {
            let mut v = Vec::new();
            v.reserve_exact(capacity);
            v
        }
        let mut len_funcs = module.len_funcs();
        let mut len_globals = module.len_globals();
        let mut len_tables = module.len_tables();
        let mut len_memories = module.len_memories();
        for import in module.imports() {
            match import.ty() {
                ExternType::Func(_) => {
                    len_funcs += 1;
                }
                ExternType::Table(_) => {
                    len_tables += 1;
                }
                ExternType::Memory(_) => {
                    len_memories += 1;
                }
                ExternType::Global(_) => {
                    len_globals += 1;
                }
            }
        }
        let layout = *module.instance_layout();
        Self {
            tables: vec_with_capacity_exact(len_tables),
            funcs: vec_with_capacity_exact(len_funcs),
            memories: vec_with_capacity_exact(len_memories),
            globals: vec_with_capacity_exact(len_globals),
            start_fn: None,
            exports: Map::default(),
            data_segments: Vec::new(),
            elem_segments: Vec::new(),
            layout,
        }
    }

    /// Sets the start function of the built instance.
    ///
    /// # Panics
    ///
    /// If the start function has already been set.
    pub fn set_start(&mut self, start_fn: FuncIdx) {
        match &mut self.start_fn {
            Some(index) => panic!("already set start function {index:?}"),
            None => {
                self.start_fn = Some(start_fn);
            }
        }
    }

    /// Returns the start function index if any.
    pub fn get_start(&self) -> Option<FuncIdx> {
        self.start_fn
    }

    /// Returns the [`Memory`] at the `index`.
    ///
    /// # Panics
    ///
    /// If there is no [`Memory`] at the given `index.
    pub fn get_memory(&self, index: u32) -> Memory {
        self.memories
            .get(index as usize)
            .copied()
            .unwrap_or_else(|| panic!("missing `Memory` at index: {index}"))
    }

    /// Returns the [`Table`] at the `index`.
    ///
    /// # Panics
    ///
    /// If there is no [`Table`] at the given `index.
    pub fn get_table(&self, index: u32) -> Table {
        self.tables
            .get(index as usize)
            .copied()
            .unwrap_or_else(|| panic!("missing `Table` at index: {index}"))
    }

    /// Returns the [`Global`] at the `index`.
    ///
    /// # Panics
    ///
    /// If there is no [`Global`] at the given `index.
    pub fn get_global(&self, index: u32) -> Global {
        self.globals
            .get(index as usize)
            .copied()
            .unwrap_or_else(|| panic!("missing `Global` at index: {index}"))
    }

    /// Returns the function at the `index`.
    ///
    /// # Panics
    ///
    /// If there is no function at the given `index.
    pub fn get_func(&self, index: u32) -> Func {
        self.funcs
            .get(index as usize)
            .copied()
            .unwrap_or_else(|| panic!("missing `Func` at index: {index}"))
    }

    /// Pushes a new [`Memory`] to the [`InstanceEntity`] under construction.
    pub fn push_memory(&mut self, memory: Memory) {
        self.memories.push(memory);
    }

    /// Pushes a new [`Table`] to the [`InstanceEntity`] under construction.
    pub fn push_table(&mut self, table: Table) {
        self.tables.push(table);
    }

    /// Pushes a new [`Global`] to the [`InstanceEntity`] under construction.
    pub fn push_global(&mut self, global: Global) {
        self.globals.push(global);
    }

    /// Pushes a new [`Func`] to the [`InstanceEntity`] under construction.
    pub fn push_func(&mut self, func: Func) {
        self.funcs.push(func);
    }

    /// Pushes a new [`Extern`] under the given `name` to the [`InstanceEntity`] under construction.
    ///
    /// # Panics
    ///
    /// If the name has already been used by an already pushed [`Extern`].
    pub fn push_export(&mut self, name: &str, new_value: Extern) {
        if let Some(old_value) = self.exports.get(name) {
            panic!(
                "tried to register {new_value:?} for name {name} \
                but name is already used by {old_value:?}",
            )
        }
        self.exports.insert(name.into(), new_value);
    }

    /// Pushes the [`DataSegment`] to the [`InstanceEntity`] under construction.
    pub fn push_data_segment(&mut self, segment: DataSegment) {
        self.data_segments.push(segment);
    }

    /// Pushes the [`ElementSegment`] to the [`InstanceEntity`] under construction.
    pub fn push_element_segment(&mut self, segment: ElementSegment) {
        self.elem_segments.push(segment);
    }

    /// Finishes constructing the [`InstanceEntity`].
    pub fn finish(self) -> Result<Box<InstanceEntity>, Error> {
        let handles = self.finish_handles();
        Ok(InstanceEntity::new_init(self.exports, self.layout, handles))
    }

    /// Finishes construction of the [`AnyHandle`] buffer.
    fn finish_handles(&self) -> Vec<AnyHandleAndEntity> {
        fn handle_and_cache_iter<T>(
            slice: &[T],
        ) -> impl ExactSizeIterator<Item = AnyHandleAndEntity> + FusedIterator
        where
            T: Clone,
            AnyHandle: From<T>,
        {
            slice
                .iter()
                .cloned()
                .map(AnyHandle::from)
                .map(AnyHandleAndEntity::new)
        }

        // Note: this is not `InstanceLayout` based since modules without a `data_count`
        //       section do not expose addresses for their data segments.
        let len_handles = self.memories.len()
            + self.globals.len()
            + self.tables.len()
            + self.funcs.len()
            + self.elem_segments.len()
            + self.data_segments.len();
        let mut handles = Vec::with_capacity(len_handles);
        handles.extend(handle_and_cache_iter(&self.memories));
        handles.extend(handle_and_cache_iter(&self.globals));
        handles.extend(handle_and_cache_iter(&self.tables));
        handles.extend(handle_and_cache_iter(&self.funcs));
        handles.extend(handle_and_cache_iter(&self.elem_segments));
        handles.extend(handle_and_cache_iter(&self.data_segments));
        handles
    }
}