refactor(ecs): create archetype lookup entries on-the-go

8 files changed, 224 insertions, 403 deletions

diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index 057b5ff..67ae453 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -6,7 +6,7 @@ use seq_macro::seq;
 use crate::lock::WriteGuard;
 use crate::system::{ComponentRefMut, Input as SystemInput};
 use crate::type_name::TypeName;
-use crate::{EntityComponent, WorldData};
+use crate::EntityComponent;
 
 pub mod local;
 
@@ -36,12 +36,6 @@ pub trait Component: SystemInput + Any + TypeName
     {
         false
     }
-
-    fn prepare(_world_data: &WorldData)
-    where
-        Self: Sized,
-    {
-    }
 }
 
 impl dyn Component
@@ -150,8 +144,6 @@ pub trait Sequence
     fn from_components<'component>(
         components: impl Iterator<Item = &'component EntityComponent>,
     ) -> Self::Refs<'component>;
-
-    fn prepare(_world_data: &WorldData);
 }
 
 /// [`Component`] metadata.
@@ -262,13 +254,6 @@ macro_rules! inner {
                         Comp~I::RefMut::from_optional_component(comp_~I),
                     )*)
                 }
-
-                fn prepare(world_data: &WorldData)
-                {
-                    #(
-                        Comp~I::prepare(world_data);
-                    )*
-                }
             }
         });
     };
diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 2bd3826..6e06ded 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -1,4 +1,6 @@
 use std::any::type_name;
+use std::borrow::Borrow;
+use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
 use std::slice::Iter as SliceIter;
 
@@ -12,35 +14,83 @@ use crate::component::{
 use crate::entity::Uid as EntityUid;
 use crate::lock::Lock;
 use crate::type_name::TypeName;
-use crate::util::Sortable;
+use crate::util::{RefCellRefMap, Sortable};
 use crate::EntityComponent;
 
 #[derive(Debug, Default)]
 pub struct Storage
 {
     archetypes: Vec<Archetype>,
-    archetype_lookup: HashMap<ArchetypeId, Vec<usize>>,
-    pending_archetype_lookup_entries: Vec<Vec<ComponentMetadata>>,
+    archetype_lookup: RefCell<HashMap<ArchetypeId, ArchetypeLookupEntry>>,
 }
 
 impl Storage
 {
-    pub fn find_entities(
+    pub fn find_entities<CompMetadataList>(
         &self,
-        mut components_metadata: impl IntoIterator<Item = ComponentMetadata>
-            + Sortable<Item = ComponentMetadata>,
+        mut components_metadata: CompMetadataList,
     ) -> ArchetypeRefIter<'_>
+    where
+        CompMetadataList: Sortable<Item = ComponentMetadata>,
+        CompMetadataList: Borrow<[ComponentMetadata]>,
     {
         components_metadata.sort_by_key_b(|component_metadata| component_metadata.id);
 
-        self.archetype_lookup
-            .get(&ArchetypeId::from_components_metadata(components_metadata))
-            .map_or_else(ArchetypeRefIter::new_empty, |archetypes_indices| {
-                ArchetypeRefIter {
-                    inner: archetypes_indices.iter(),
-                    archetypes: &self.archetypes,
+        let archetype_id = ArchetypeId::from_components_metadata(
+            components_metadata.borrow().into_iter().cloned(),
+        );
+
+        // This looks stupid but the borrow checker complains otherwise
+        if self.archetype_lookup.borrow().contains_key(&archetype_id) {
+            let archetype_lookup = self.archetype_lookup.borrow();
+
+            return ArchetypeRefIter {
+                inner: RefCellRefMap::new(archetype_lookup, |archetype_lookup| {
+                    archetype_lookup
+                        .get(&archetype_id)
+                        .unwrap()
+                        .archetype_indices
+                        .iter()
+                }),
+                archetypes: &self.archetypes,
+            };
+        }
+
+        let comp_ids_set = create_non_opt_component_id_set(components_metadata.borrow());
+
+        let matching_archetype_indices = self
+            .archetypes
+            .iter()
+            .enumerate()
+            .filter_map(|(index, archetype)| {
+                if archetype.component_ids_is_superset(&comp_ids_set) {
+                    return Some(index);
                 }
+
+                None
             })
+            .collect();
+
+        self.archetype_lookup.borrow_mut().insert(
+            archetype_id,
+            ArchetypeLookupEntry {
+                component_ids: comp_ids_set,
+                archetype_indices: matching_archetype_indices,
+            },
+        );
+
+        let archetype_lookup = self.archetype_lookup.borrow();
+
+        ArchetypeRefIter {
+            inner: RefCellRefMap::new(archetype_lookup, |archetype_lookup| {
+                archetype_lookup
+                    .get(&archetype_id)
+                    .unwrap()
+                    .archetype_indices
+                    .iter()
+            }),
+            archetypes: &self.archetypes,
+        }
     }
 
     #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
@@ -67,48 +117,25 @@ impl Storage
                 .map(|component| ComponentMetadata::of(&**component)),
         );
 
-        let archetype_indices =
-            self.archetype_lookup
-                .entry(archetype_id)
-                .or_insert_with(|| {
-                    self.archetypes.push(Archetype::new(
-                        components.iter().map(|component| component.id()),
-                    ));
-
-                    vec![self.archetypes.len() - 1]
-                });
+        let comp_ids_set = create_non_opt_component_id_set(
+            components
+                .iter()
+                .map(|component| ComponentMetadata::of(&**component)),
+        );
 
-        if archetype_indices.is_empty() {
-            self.archetypes.push(Archetype::new(
-                components.iter().map(|component| component.id()),
-            ));
+        let archetype_index =
+            self.get_or_create_archetype(archetype_id, &comp_ids_set, &components);
 
-            archetype_indices.push(self.archetypes.len() - 1);
-        }
+        self.populate_matching_archetype_lookup_entries(&comp_ids_set, archetype_index);
 
         let archetype = self
             .archetypes
-            .get_mut(
-                archetype_indices
-                    .first()
-                    .copied()
-                    .expect("No archetype index found"),
-            )
+            .get_mut(archetype_index)
             .expect("Archetype is gone");
 
         let entity_uid = EntityUid::new_unique();
 
-        archetype.entities.push(ArchetypeEntity {
-            uid: entity_uid,
-            components: components
-                .into_iter()
-                .map(|component| EntityComponent {
-                    id: component.id(),
-                    name: component.type_name(),
-                    component: Lock::new(component),
-                })
-                .collect(),
-        });
+        archetype.push_entity(entity_uid, components);
 
         archetype
             .entity_lookup
@@ -117,54 +144,54 @@ impl Storage
         (archetype_id, entity_uid)
     }
 
-    pub fn add_archetype_lookup_entry(
+    fn populate_matching_archetype_lookup_entries(
         &mut self,
-        mut components_metadata: impl IntoIterator<Item = ComponentMetadata>
-            + Sortable<Item = ComponentMetadata>,
+        comp_ids_set: &HashSet<ComponentId>,
+        archetype_index: usize,
     )
     {
-        components_metadata.sort_by_key_b(|component_metadata| component_metadata.id);
+        let mut archetype_lookup = self.archetype_lookup.borrow_mut();
 
-        self.pending_archetype_lookup_entries
-            .push(components_metadata.into_iter().collect());
+        for (_, lookup_entry) in archetype_lookup.iter_mut() {
+            if &lookup_entry.component_ids == comp_ids_set {
+                continue;
+            }
+
+            if lookup_entry.component_ids.is_subset(&comp_ids_set) {
+                lookup_entry.archetype_indices.push(archetype_index);
+            }
+        }
     }
 
-    pub fn make_archetype_lookup_entries(&mut self)
+    fn get_or_create_archetype(
+        &mut self,
+        archetype_id: ArchetypeId,
+        comp_ids_set: &HashSet<ComponentId>,
+        components: &[Box<dyn Component>],
+    ) -> usize
     {
-        for pending_entry in self.pending_archetype_lookup_entries.drain(..) {
-            let pending_entry_ids_set = pending_entry
-                .iter()
-                .filter_map(|component_metadata| {
-                    if component_metadata.is_optional == ComponentIsOptional::Yes {
-                        return None;
-                    }
-
-                    Some(component_metadata.id)
-                })
-                .collect::<HashSet<_>>();
-
-            let matching_archetype_indices = self
-                .archetypes
-                .iter()
-                .enumerate()
-                .filter_map(|(index, archetype)| {
-                    if archetype.component_ids_is_superset(&pending_entry_ids_set) {
-                        return Some(index);
-                    }
-
-                    None
-                });
+        let mut archetype_lookup = self.archetype_lookup.borrow_mut();
 
-            let archetype_id =
-                ArchetypeId::from_components_metadata(pending_entry.into_iter());
-
-            if self.archetype_lookup.contains_key(&archetype_id) {
-                continue;
+        let lookup_entry = archetype_lookup.entry(archetype_id).or_insert_with(|| {
+            ArchetypeLookupEntry {
+                component_ids: comp_ids_set.clone(),
+                archetype_indices: Vec::with_capacity(1),
             }
+        });
 
-            self.archetype_lookup
-                .insert(archetype_id, matching_archetype_indices.collect());
+        if lookup_entry.archetype_indices.is_empty() {
+            self.archetypes.push(Archetype::new(
+                components.iter().map(|component| component.id()),
+            ));
+
+            lookup_entry
+                .archetype_indices
+                .push(self.archetypes.len() - 1);
         }
+
+        // SAFETY: Above, we push a archetype index if archetype_indices is empty so this
+        // cannot fail
+        unsafe { *lookup_entry.archetype_indices.first().unwrap_unchecked() }
     }
 }
 
@@ -177,6 +204,13 @@ impl TypeName for Storage
 }
 
 #[derive(Debug)]
+struct ArchetypeLookupEntry
+{
+    component_ids: HashSet<ComponentId>,
+    archetype_indices: Vec<usize>,
+}
+
+#[derive(Debug)]
 pub struct Archetype
 {
     component_ids: HashMap<ComponentId, usize>,
@@ -229,6 +263,25 @@ impl Archetype
     {
         self.component_ids.get(component_id).copied()
     }
+
+    fn push_entity(
+        &mut self,
+        entity_uid: EntityUid,
+        components: impl IntoIterator<Item = Box<dyn Component>>,
+    )
+    {
+        self.entities.push(ArchetypeEntity {
+            uid: entity_uid,
+            components: components
+                .into_iter()
+                .map(|component| EntityComponent {
+                    id: component.id(),
+                    name: component.type_name(),
+                    component: Lock::new(component),
+                })
+                .collect(),
+        });
+    }
 }
 
 #[derive(Debug)]
@@ -259,18 +312,14 @@ impl ArchetypeEntity
 #[derive(Debug)]
 pub struct ArchetypeRefIter<'component_storage>
 {
-    inner: SliceIter<'component_storage, usize>,
+    inner: RefCellRefMap<
+        'component_storage,
+        HashMap<ArchetypeId, ArchetypeLookupEntry>,
+        SliceIter<'component_storage, usize>,
+    >,
     archetypes: &'component_storage [Archetype],
 }
 
-impl<'component_storage> ArchetypeRefIter<'component_storage>
-{
-    fn new_empty() -> Self
-    {
-        Self { inner: [].iter(), archetypes: &[] }
-    }
-}
-
 impl<'component_storage> Iterator for ArchetypeRefIter<'component_storage>
 {
     type Item = &'component_storage Archetype;
@@ -303,24 +352,40 @@ impl<'archetype> Iterator for EntityIter<'archetype>
     }
 }
 
+fn create_non_opt_component_id_set<'a, Item>(
+    component_metadata_iter: impl IntoIterator<Item = Item>,
+) -> HashSet<ComponentId>
+where
+    Item: Borrow<ComponentMetadata>,
+{
+    component_metadata_iter
+        .into_iter()
+        .filter_map(|item| {
+            let component_metadata = item.borrow();
+
+            if component_metadata.is_optional == ComponentIsOptional::Yes {
+                return None;
+            }
+
+            Some(component_metadata.id)
+        })
+        .collect::<HashSet<_>>()
+}
+
 #[cfg(test)]
 mod tests
 {
-    use std::collections::HashMap;
 
     use ecs_macros::Component;
 
-    use super::{Archetype, Storage};
+    use super::Storage;
     use crate::archetype::Id as ArchetypeId;
-    use crate::component::storage::ArchetypeEntity;
     use crate::component::{
         Id as ComponentId,
         IsOptional as ComponentIsOptional,
         Metadata as ComponentMetadata,
     };
-    use crate::entity::Uid as EntityUid;
-    use crate::lock::Lock;
-    use crate::{self as ecs, EntityComponent};
+    use crate::{self as ecs};
 
     #[derive(Debug, Component)]
     struct HealthPotion
@@ -378,7 +443,7 @@ mod tests
 
         assert_eq!(entity_components.components.len(), 2);
 
-        assert_eq!(component_storage.archetype_lookup.len(), 1);
+        assert_eq!(component_storage.archetype_lookup.borrow().len(), 1);
 
         let mut components_metadata = [
             ComponentMetadata {
@@ -393,160 +458,17 @@ mod tests
 
         components_metadata.sort_by_key(|comp_metadata| comp_metadata.id);
 
-        let lookup = component_storage
-            .archetype_lookup
+        let archetype_lookup = component_storage.archetype_lookup.borrow();
+
+        let lookup_entry = archetype_lookup
             .get(&ArchetypeId::from_components_metadata(components_metadata))
             .expect("Expected entry in archetype lookup map");
 
-        let first_archetype_index = lookup
+        let first_archetype_index = lookup_entry
+            .archetype_indices
             .first()
             .expect("Expected archetype lookup to contain a archetype reference");
 
         assert_eq!(*first_archetype_index, 0);
     }
-
-    #[test]
-    fn lookup_works()
-    {
-        let mut component_storage = Storage::default();
-
-        let entity_uid_a = EntityUid::new_unique();
-        let entity_uid_b = EntityUid::new_unique();
-        let entity_uid_c = EntityUid::new_unique();
-
-        component_storage.archetypes.push(Archetype {
-            component_ids: HashMap::from([
-                (ComponentId::of::<IronBoots>(), 0),
-                (ComponentId::of::<HealthPotion>(), 1),
-                (ComponentId::of::<Hookshot>(), 2),
-            ]),
-            entity_lookup: HashMap::from([
-                (entity_uid_a, 0),
-                (entity_uid_b, 1),
-                (entity_uid_c, 2),
-            ]),
-            entities: vec![
-                ArchetypeEntity {
-                    uid: entity_uid_a,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<IronBoots>(),
-                        name: "Iron boots",
-                        component: Lock::new(Box::new(IronBoots)),
-                    }],
-                },
-                ArchetypeEntity {
-                    uid: entity_uid_b,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<HealthPotion>(),
-                        name: "Health potion",
-                        component: Lock::new(Box::new(HealthPotion {
-                            _hp_restoration: 20,
-                        })),
-                    }],
-                },
-                ArchetypeEntity {
-                    uid: entity_uid_c,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<Hookshot>(),
-                        name: "Hookshot",
-                        component: Lock::new(Box::new(Hookshot { _range: 67 })),
-                    }],
-                },
-            ],
-        });
-
-        let entity_uid_d = EntityUid::new_unique();
-        let entity_uid_e = EntityUid::new_unique();
-        let entity_uid_f = EntityUid::new_unique();
-        let entity_uid_g = EntityUid::new_unique();
-
-        component_storage.archetypes.push(Archetype {
-            component_ids: HashMap::from([
-                (ComponentId::of::<DekuNut>(), 0),
-                (ComponentId::of::<IronBoots>(), 1),
-                (ComponentId::of::<Bow>(), 2),
-                (ComponentId::of::<Hookshot>(), 3),
-            ]),
-            entity_lookup: HashMap::from([
-                (entity_uid_d, 0),
-                (entity_uid_e, 1),
-                (entity_uid_f, 2),
-                (entity_uid_g, 3),
-            ]),
-            entities: vec![
-                ArchetypeEntity {
-                    uid: entity_uid_d,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<DekuNut>(),
-                        name: "Deku nut",
-                        component: Lock::new(Box::new(DekuNut { _throwing_damage: 5 })),
-                    }],
-                },
-                ArchetypeEntity {
-                    uid: entity_uid_e,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<IronBoots>(),
-                        name: "Iron boots",
-                        component: Lock::new(Box::new(IronBoots)),
-                    }],
-                },
-                ArchetypeEntity {
-                    uid: entity_uid_f,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<Bow>(),
-                        name: "Bow",
-                        component: Lock::new(Box::new(Bow { _damage: 20 })),
-                    }],
-                },
-                ArchetypeEntity {
-                    uid: entity_uid_g,
-                    components: vec![EntityComponent {
-                        id: ComponentId::of::<Hookshot>(),
-                        name: "Hookshot",
-                        component: Lock::new(Box::new(Hookshot { _range: 67 })),
-                    }],
-                },
-            ],
-        });
-
-        component_storage.add_archetype_lookup_entry([
-            ComponentMetadata {
-                id: ComponentId::of::<IronBoots>(),
-                is_optional: ComponentIsOptional::No,
-            },
-            ComponentMetadata {
-                id: ComponentId::of::<Hookshot>(),
-                is_optional: ComponentIsOptional::No,
-            },
-        ]);
-
-        assert_eq!(component_storage.pending_archetype_lookup_entries.len(), 1);
-
-        component_storage.make_archetype_lookup_entries();
-
-        assert_eq!(component_storage.archetype_lookup.len(), 1);
-
-        let mut comps_metadata = [
-            ComponentMetadata {
-                id: ComponentId::of::<IronBoots>(),
-                is_optional: ComponentIsOptional::No,
-            },
-            ComponentMetadata {
-                id: ComponentId::of::<Hookshot>(),
-                is_optional: ComponentIsOptional::No,
-            },
-        ];
-
-        comps_metadata.sort_by_key(|comp_metadata| comp_metadata.id);
-
-        let archetypes = component_storage
-            .archetype_lookup
-            .get(&ArchetypeId::from_components_metadata(comps_metadata))
-            .expect(concat!(
-                "Expected a archetype for IronBoots & Hookshot to be found in the ",
-                "archetype lookup map"
-            ));
-
-        assert_eq!(archetypes.len(), 2);
-    }
 }
diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index 0da8b71..34d6b7b 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -94,8 +94,6 @@ impl World
     ) where
         EventT: Event,
     {
-        system.prepare(&self.data);
-
         self.systems.push(system.into_type_erased());
 
         self.data
@@ -155,27 +153,12 @@ impl World
         {
             match action {
                 Action::Spawn(components) => {
-                    {
-                        let mut component_storage_lock =
-                            self.data.component_storage.write_nonblock().expect(
-                                "Failed to acquire read-write component storage lock",
-                            );
-
-                        component_storage_lock.push_entity(components);
-                    }
-
-                    for system in &self.systems {
-                        unsafe {
-                            system.prepare(&self.data);
-                        }
-                    }
-
                     let mut component_storage_lock =
                         self.data.component_storage.write_nonblock().expect(
                             "Failed to acquire read-write component storage lock",
                         );
 
-                    component_storage_lock.make_archetype_lookup_entries();
+                    component_storage_lock.push_entity(components);
                 }
                 Action::Stop => {
                     self.stop.store(true, Ordering::Relaxed);
@@ -184,18 +167,6 @@ impl World
         }
     }
 
-    /// Prepares the world. Should be called before manually emitting events and after
-    /// creating entities, registering systems & adding extensions. You do not need to
-    /// call this function if you use [`event_loop`].
-    pub fn prepare(&self)
-    {
-        self.data
-            .component_storage
-            .write_nonblock()
-            .expect("Failed to acquire read-write component storage lock")
-            .make_archetype_lookup_entries();
-    }
-
     /// A event loop which runs until a stop is issued with [`Flags::stop`]. Before the
     /// loop begins, [`StartEvent`] is emitted.
     ///
@@ -203,8 +174,6 @@ impl World
     /// Will panic if a internal lock cannot be acquired.
     pub fn event_loop<EventSeq: EventSequence>(&self)
     {
-        self.prepare();
-
         self.emit(StartEvent);
 
         let event_seq = EventSeq::ids();
diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index 60d4210..ea61640 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -20,7 +20,7 @@ use crate::system::{
     Param as SystemParam,
     System,
 };
-use crate::{World, WorldData};
+use crate::World;
 
 pub mod options;
 
@@ -138,27 +138,6 @@ where
     {
         Box::new(QueryComponentIds { component_ids: Comps::metadata() })
     }
-
-    #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
-    fn prepare(world_data: &WorldData)
-    {
-        let mut component_storage_lock = world_data
-            .component_storage
-            .write_nonblock()
-            .expect("Failed to acquire read-write component storage lock");
-
-        #[cfg(feature = "debug")]
-        tracing::debug!(
-            "Adding archetypes lookup entry for components: ({})",
-            std::any::type_name::<Comps>()
-        );
-
-        component_storage_lock.add_archetype_lookup_entry(Comps::metadata());
-
-        drop(component_storage_lock);
-
-        Comps::prepare(world_data);
-    }
 }
 
 type ComponentIterMapFn = for<'a> fn(&'a Archetype) -> EntityIter<'a>;
diff --git a/ecs/src/relationship.rs b/ecs/src/relationship.rs
index 964a47f..b7c5c02 100644
--- a/ecs/src/relationship.rs
+++ b/ecs/src/relationship.rs
@@ -14,7 +14,7 @@ use crate::entity::Uid as EntityUid;
 use crate::lock::{Lock, ReadGuard};
 use crate::system::{ComponentRefMut, Input as SystemInput};
 use crate::type_name::TypeName;
-use crate::{World, WorldData};
+use crate::World;
 
 #[derive(Debug)]
 pub struct Relationship<Kind, ComponentT: Component>
@@ -60,28 +60,6 @@ where
     {
         self
     }
-
-    #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
-    fn prepare(world_data: &WorldData)
-    where
-        Self: Sized,
-    {
-        let mut component_storage_lock = world_data
-            .component_storage
-            .write_nonblock()
-            .expect("Failed to acquire read-write component storage lock");
-
-        #[cfg(feature = "debug")]
-        tracing::debug!(
-            "Adding archetypes lookup entry for component: {}",
-            std::any::type_name::<ComponentT>()
-        );
-
-        component_storage_lock.add_archetype_lookup_entry([ComponentMetadata {
-            id: ComponentId::of::<ComponentT>(),
-            is_optional: component_is_optional::<ComponentT>().into(),
-        }]);
-    }
 }
 
 impl<Kind, ComponentT> SystemInput for Relationship<Kind, ComponentT>
diff --git a/ecs/src/system.rs b/ecs/src/system.rs
index ba5ac96..36359c7 100644
--- a/ecs/src/system.rs
+++ b/ecs/src/system.rs
@@ -11,7 +11,7 @@ use crate::component::{Component, FromOptional as FromOptionalComponent};
 use crate::lock::WriteGuard;
 use crate::system::util::check_params_are_compatible;
 use crate::tuple::{ReduceElement as TupleReduceElement, With as TupleWith};
-use crate::{World, WorldData};
+use crate::World;
 
 pub mod stateful;
 
@@ -24,8 +24,6 @@ pub trait System<'world, Impl>: 'static
     #[must_use]
     fn initialize(self, input: Self::Input) -> Self;
 
-    fn prepare(&self, world_data: &WorldData);
-
     fn run<'this>(&'this self, world: &'world World)
     where
         'this: 'world;
@@ -58,13 +56,6 @@ macro_rules! impl_system {
                     self
                 }
 
-                fn prepare(&self, world_data: &WorldData)
-                {
-                    #(
-                        TParam~I::prepare(world_data);
-                    )*
-                }
-
                 fn run<'this>(&'this self, world: &'world World)
                 where
                     'this: 'world
@@ -99,21 +90,6 @@ macro_rules! impl_system {
 
                             me.run(world);
                         }),
-                        prepare: Box::new(|data, world| {
-                            // SAFETY: The caller of TypeErased::run ensures the lifetime
-                            // is correct
-                            let data = unsafe { &*std::ptr::from_ref(data) };
-
-                            let me = data
-                                .downcast_ref::<Func>()
-                                .expect("Function downcast failed");
-
-                            // SAFETY: The caller of TypeErased::run ensures the lifetime
-                            // is correct
-                            let world = unsafe { &*std::ptr::from_ref(world) };
-
-                            me.prepare(world);
-                        }),
                     }
                 }
 
@@ -150,7 +126,6 @@ pub struct TypeErased
 {
     data: Box<dyn Any + RefUnwindSafe + UnwindSafe>,
     run: Box<TypeErasedRunFn>,
-    prepare: Box<TypeErasedPrepareFn>,
 }
 
 impl TypeErased
@@ -166,18 +141,6 @@ impl TypeErased
 
         (self.run)(data, world);
     }
-
-    /// Prepares the system.
-    ///
-    /// # Safety
-    /// `world_data` must live at least as long as the [`World`] the system belongs to.
-    pub unsafe fn prepare(&self, world_data: &WorldData)
-    {
-        // You have to dereference for downcasting to work for some reason
-        let data = &*self.data;
-
-        (self.prepare)(data, world_data);
-    }
 }
 
 impl Debug for TypeErased
@@ -191,9 +154,6 @@ impl Debug for TypeErased
 /// Function in [`TypeErased`] used to run the system.
 type TypeErasedRunFn = dyn Fn(&dyn Any, &World) + RefUnwindSafe + UnwindSafe;
 
-/// Function in [`TypeErased`] used to prepare the system.
-type TypeErasedPrepareFn = dyn Fn(&dyn Any, &WorldData) + RefUnwindSafe + UnwindSafe;
-
 /// A parameter to a [`System`].
 ///
 /// # Safety
@@ -216,8 +176,6 @@ pub unsafe trait Param<'world>
     fn is_compatible<Other: Param<'world>>() -> bool;
 
     fn get_comparable() -> Box<dyn Any>;
-
-    fn prepare(_world_data: &WorldData) {}
 }
 
 pub struct NoInitParamFlag {}
diff --git a/ecs/src/system/stateful.rs b/ecs/src/system/stateful.rs
index dd0c5e1..99b56c1 100644
--- a/ecs/src/system/stateful.rs
+++ b/ecs/src/system/stateful.rs
@@ -21,7 +21,7 @@ use crate::tuple::{
     TakeOptionElementResult as TupleTakeOptionElementResult,
     WithOptionElements as TupleWithOptionElements,
 };
-use crate::{World, WorldData};
+use crate::World;
 
 /// A stateful system.
 pub struct Stateful<Func>
@@ -85,13 +85,6 @@ macro_rules! impl_system {
                     self
                 }
 
-                fn prepare(&self, world_data: &WorldData)
-                {
-                    #(
-                        TParam~I::prepare(world_data);
-                    )*
-                }
-
                 fn run<'this>(&'this self, world: &'world World)
                 where
                     'this: 'world
@@ -124,19 +117,6 @@ macro_rules! impl_system {
 
                             me.run(world);
                         }),
-                        prepare: Box::new(|data, world| {
-                            // SAFETY: The caller of TypeErased::run ensures the lifetime
-                            // is correct
-                            let data = unsafe { &*std::ptr::from_ref::<dyn Any>(data) };
-
-                            let me = data.downcast_ref::<Self>().unwrap();
-
-                            // SAFETY: The caller of TypeErased::run ensures the lifetime
-                            // is correct
-                            let world = unsafe { &*std::ptr::from_ref(world) };
-
-                            me.prepare(world);
-                        }),
                     }
                 }
 
diff --git a/ecs/src/util.rs b/ecs/src/util.rs
index f4f8632..5ec21c7 100644
--- a/ecs/src/util.rs
+++ b/ecs/src/util.rs
@@ -1,3 +1,6 @@
+use std::cell::Ref;
+use std::ops::{Deref, DerefMut};
+
 pub trait Sortable
 {
     type Item;
@@ -46,3 +49,50 @@ impl<Item> Sortable for Vec<Item>
         self.sort_by_key(func)
     }
 }
+
+#[derive(Debug)]
+pub struct RefCellRefMap<'a, Value: 'a, MappedValue: 'a>
+{
+    mapped_value: MappedValue,
+    _refcell_ref: Ref<'a, Value>,
+}
+
+impl<'a, Value: 'a, MappedValue: 'a> RefCellRefMap<'a, Value, MappedValue>
+{
+    pub fn new(
+        refcell_ref: Ref<'a, Value>,
+        map_fn: impl Fn(&'a Value) -> MappedValue,
+    ) -> Self
+    {
+        // Convert the lifetime to 'static. This is necessary to make the compiler not
+        // complain about refcell_ref being moved
+        //
+        // SAFETY: This is fine since we pass it to map_fn with the original lifetime 'a
+        let val_ref = unsafe { &*(&*refcell_ref as *const Value) };
+
+        let mapped_value = map_fn(val_ref);
+
+        Self {
+            mapped_value,
+            _refcell_ref: refcell_ref,
+        }
+    }
+}
+
+impl<'a, Value: 'a, MappedValue: 'a> Deref for RefCellRefMap<'a, Value, MappedValue>
+{
+    type Target = MappedValue;
+
+    fn deref(&self) -> &Self::Target
+    {
+        &self.mapped_value
+    }
+}
+
+impl<'a, Value: 'a, MappedValue: 'a> DerefMut for RefCellRefMap<'a, Value, MappedValue>
+{
+    fn deref_mut(&mut self) -> &mut Self::Target
+    {
+        &mut self.mapped_value
+    }
+}