lyra-engine/lyra-ecs/src/bundle.rs

use std::{ptr::NonNull, mem::size_of, alloc::Layout};

use crate::{component::Component, component_info::ComponentInfo, DynTypeId};

pub trait Bundle {
    /// Get a list of type ids that this bundle is storing
    fn type_ids(&self) -> Vec<DynTypeId>;

    /// Get ComponentInfo's for the components in this bundle
    fn info(&self) -> Vec<ComponentInfo>;

    /// Take the bundle by calling the closure with pointers to each component, its type and size.
    /// The closure is expected to take ownership of the pointer.
    fn take(self, f: impl FnMut(NonNull<u8>, DynTypeId, usize));

    /// Returns a boolean indicating if this Bundle is dynamic. See [`DynamicBundle`]
    fn is_dynamic(&self) -> bool;
}

impl<C: Component> Bundle for C {
    fn type_ids(&self) -> Vec<DynTypeId> {
        vec![DynTypeId::of::<C>()]
    }

    fn info(&self) -> Vec<ComponentInfo> {
        vec![ComponentInfo::new::<C>()]
    }

    fn take(self, mut f: impl FnMut(NonNull<u8>, DynTypeId, usize)) {
        f(NonNull::from(&self).cast(), DynTypeId::of::<C>(), size_of::<C>());

        // this must be done to avoid calling drop on heap memory that the component
        // may manage. So something like a Vec, or HashMap, etc.
        std::mem::forget(self);
    }

    fn is_dynamic(&self) -> bool {
        false
    }
}

macro_rules! impl_bundle_tuple {
    ( $($name: ident),+ ) => (
        #[allow(non_snake_case)]
        impl<$($name: Component),+> Bundle for ($($name,)+) {
            fn type_ids(&self) -> Vec<DynTypeId> {
                // these names wont follow rust convention, but its a macro so deal with it
                vec![$(DynTypeId::of::<$name>()),+]
            }

            fn info(&self) -> Vec<ComponentInfo> {
                vec![$(ComponentInfo::new::<$name>()),+]
            }

            fn take(self, mut f: impl FnMut(NonNull<u8>, DynTypeId, usize)) {
                // these names wont follow rust convention, but its a macro so deal with it
                let ($($name,)+) = self;

                $(
                    f(NonNull::from(&$name).cast(), DynTypeId::of::<$name>(), size_of::<$name>());
                    // this must be done to avoid calling drop on heap memory that the component
                    // may manage. So something like a Vec, or HashMap, etc.
                    std::mem::forget($name);
                )+
            }

            fn is_dynamic(&self) -> bool {
                false
            }
        }
    );
}

// hopefully 16 components in a bundle is enough
impl_bundle_tuple! { C1 }
impl_bundle_tuple! { C1, C2 }
impl_bundle_tuple! { C1, C2, C3 }
impl_bundle_tuple! { C1, C2, C3, C4 }
impl_bundle_tuple! { C1, C2, C3, C4, C5 }
impl_bundle_tuple! { C1, C2, C3, C4, C5, C6 }
impl_bundle_tuple! { C1, C2, C3, C4, C5, C6, C7 }
impl_bundle_tuple! { C1, C2, C3, C4, C5, C6, C7, C8 }
impl_bundle_tuple! { C1, C2, C3, C4, C5, C6, C7, C8, C9 }
impl_bundle_tuple! { C1, C2, C3, C4, C5, C6, C7, C8, C9, C10 }
impl_bundle_tuple! { C1, C2, C3, C4, C5, C6, C7, C8, C9, C10, C11 }
impl_bundle_tuple! { C1, C2, C3, C4, C5, C6, C7, C8, C9, C10, C11, C12 }
impl_bundle_tuple! { C1, C2, C3, C4, C5, C6, C7, C8, C9, C10, C11, C12, C13 }
impl_bundle_tuple! { C1, C2, C3, C4, C5, C6, C7, C8, C9, C10, C11, C12, C13, C14 }
impl_bundle_tuple! { C1, C2, C3, C4, C5, C6, C7, C8, C9, C10, C11, C12, C13, C14, C15 }
impl_bundle_tuple! { C1, C2, C3, C4, C5, C6, C7, C8, C9, C10, C11, C12, C13, C14, C15, C16 }

/// A bundle of a dynamic number of components. The types of the components may not be known to Rust.
/// 
/// # Safety
/// Do not drop this without inserting it into an archetype. It WILL cause a memory leak
#[derive(Default, Clone)]
pub struct DynamicBundle {
    bundle: Vec<(NonNull<u8>, ComponentInfo)>,
}

// TODO: When a bundle is dropped without being inserted into an archetype, it WILL cause a memory leak. Find a way around that.
// maybe it can be done with Rc, or Weak, or a mixture of both.
impl DynamicBundle {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn is_empty(&self) -> bool {
        self.bundle.len() == 0
    }

    pub fn len(&self) -> usize {
        self.bundle.len()
    }

    /// Push a type known to rust to this bundle
    pub fn push<C>(&mut self, comp: C)
    where
        C: Component
    {
        let info = ComponentInfo::new::<C>();
        
        // an owned pointer must be created from the provided component since comp would drop
        // out of scope and the data would become invalid
        let ptr = unsafe {
            let data = NonNull::from(&comp);

            let layout = Layout::new::<C>();
            let alloc_ptr = NonNull::new_unchecked(std::alloc::alloc(layout)).cast::<C>();
            std::ptr::copy_nonoverlapping(data.as_ptr(), alloc_ptr.as_ptr(), 1);

            alloc_ptr.cast()
        };

        self.bundle.push((ptr, info));
    }

    /// Push an unknown type to the bundle
    pub fn push_unknown(&mut self, data: NonNull<u8>, info: ComponentInfo) {
        self.bundle.push((data, info));
    }
}

impl Bundle for DynamicBundle {
    fn type_ids(&self) -> Vec<DynTypeId> {
        self.bundle.iter().map(|b| b.1.type_id).collect()
    }

    fn info(&self) -> Vec<ComponentInfo> {
        self.bundle.iter().map(|b| b.1).collect()
    }

    fn take(self, mut f: impl FnMut(NonNull<u8>, DynTypeId, usize)) {
        for (data, info) in self.bundle.iter() {
            f(*data, info.type_id, info.layout.size());
        }
    }

    fn is_dynamic(&self) -> bool {
        true
    }
}