7 files changed, 331 insertions, 21 deletions

diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index 67ae453..057b5ff 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;
+use crate::{EntityComponent, WorldData};
 
 pub mod local;
 
@@ -36,6 +36,12 @@ pub trait Component: SystemInput + Any + TypeName
     {
         false
     }
+
+    fn prepare(_world_data: &WorldData)
+    where
+        Self: Sized,
+    {
+    }
 }
 
 impl dyn Component
@@ -144,6 +150,8 @@ pub trait Sequence
     fn from_components<'component>(
         components: impl Iterator<Item = &'component EntityComponent>,
     ) -> Self::Refs<'component>;
+
+    fn prepare(_world_data: &WorldData);
 }
 
 /// [`Component`] metadata.
@@ -254,6 +262,13 @@ 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 76701ac..b227337 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -9,6 +9,7 @@ use crate::component::{
     IsOptional as ComponentIsOptional,
     Metadata as ComponentMetadata,
 };
+use crate::entity::Uid as EntityUid;
 use crate::lock::Lock;
 use crate::type_name::TypeName;
 use crate::EntityComponent;
@@ -39,8 +40,13 @@ impl Storage
     }
 
     #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
-    pub fn push_entity(&mut self, components: Vec<Box<dyn Component>>)
+    pub fn push_entity(
+        &mut self,
+        mut components: Vec<Box<dyn Component>>,
+    ) -> (ArchetypeId, EntityUid)
     {
+        components.sort_by_key(|component| component.id());
+
         #[cfg(feature = "debug")]
         tracing::debug!(
             "Pushing entity with components: ({})",
@@ -51,20 +57,22 @@ impl Storage
                 .join(", ")
         );
 
-        let archetype_indices = self
-            .archetype_lookup
-            .entry(ArchetypeId::from_components_metadata(
-                components
-                    .iter()
-                    .map(|component| ComponentMetadata::of(&**component)),
-            ))
-            .or_insert_with(|| {
-                self.archetypes.push(Archetype::new(
-                    components.iter().map(|component| component.id()),
-                ));
-
-                vec![self.archetypes.len() - 1]
-            });
+        let archetype_id = ArchetypeId::from_components_metadata(
+            components
+                .iter()
+                .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]
+                });
 
         if archetype_indices.is_empty() {
             self.archetypes.push(Archetype::new(
@@ -94,6 +102,14 @@ impl Storage
                 })
                 .collect(),
         );
+
+        let entity_uid = EntityUid::new_unique();
+
+        archetype
+            .entity_lookup
+            .insert(entity_uid, archetype.components.len() - 1);
+
+        (archetype_id, entity_uid)
     }
 
     pub fn add_archetype_lookup_entry(
@@ -127,7 +143,7 @@ impl Storage
                 .iter()
                 .enumerate()
                 .filter_map(|(index, archetype)| {
-                    if archetype.component_ids.is_superset(&ids_set) {
+                    if archetype.component_ids_is_superset(&ids_set) {
                         return Some(index);
                     }
 
@@ -157,7 +173,8 @@ impl TypeName for Storage
 #[derive(Debug)]
 pub struct Archetype
 {
-    component_ids: HashSet<ComponentId>,
+    component_ids: HashMap<ComponentId, usize>,
+    entity_lookup: HashMap<EntityUid, usize>,
     pub components: Vec<Vec<EntityComponent>>,
 }
 
@@ -166,10 +183,41 @@ impl Archetype
     fn new(component_ids: impl IntoIterator<Item = ComponentId>) -> Self
     {
         Self {
-            component_ids: component_ids.into_iter().collect(),
+            component_ids: component_ids
+                .into_iter()
+                .enumerate()
+                .map(|(index, component_id)| (component_id, index))
+                .collect(),
+            entity_lookup: HashMap::new(),
             components: Vec::new(),
         }
     }
+
+    pub fn component_ids_is_superset(
+        &self,
+        other_component_ids: &HashSet<ComponentId>,
+    ) -> bool
+    {
+        if other_component_ids.len() <= self.component_ids.len() {
+            other_component_ids
+                .iter()
+                .all(|v| self.component_ids.contains_key(v))
+        } else {
+            false
+        }
+    }
+
+    pub fn get_entity(&self, entity_uid: EntityUid) -> Option<&[EntityComponent]>
+    {
+        let entity_index = *self.entity_lookup.get(&entity_uid)?;
+
+        Some(&self.components.get(entity_index)?)
+    }
+
+    pub fn get_index_for_component(&self, component_id: &ComponentId) -> Option<usize>
+    {
+        self.component_ids.get(component_id).copied()
+    }
 }
 
 #[derive(Debug)]
diff --git a/ecs/src/entity.rs b/ecs/src/entity.rs
new file mode 100644
index 0000000..abc5991
--- /dev/null
+++ b/ecs/src/entity.rs
@@ -0,0 +1,30 @@
+use std::sync::atomic::{AtomicU64, Ordering};
+
+static NEXT_UID: AtomicU64 = AtomicU64::new(0);
+
+/// Unique entity ID.
+#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
+pub struct Uid
+{
+    inner: u64,
+}
+
+impl Uid
+{
+    pub fn new(uid: u64) -> Self
+    {
+        debug_assert!(
+            uid < NEXT_UID.load(Ordering::Relaxed),
+            "Invalid entity UID {uid}"
+        );
+
+        Self { inner: uid }
+    }
+
+    pub fn new_unique() -> Self
+    {
+        Self {
+            inner: NEXT_UID.fetch_add(1, Ordering::Relaxed),
+        }
+    }
+}
diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index 741b555..b920b48 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -12,6 +12,7 @@ use std::vec::Drain;
 use crate::actions::Action;
 use crate::component::storage::Storage as ComponentStorage;
 use crate::component::{Component, Id as ComponentId, Sequence as ComponentSequence};
+use crate::entity::Uid as EntityUid;
 use crate::event::{Event, Id as EventId, Ids, Sequence as EventSequence};
 use crate::extension::{Collector as ExtensionCollector, Extension};
 use crate::lock::Lock;
@@ -22,10 +23,12 @@ use crate::type_name::TypeName;
 
 pub mod actions;
 pub mod component;
+pub mod entity;
 pub mod event;
 pub mod extension;
 pub mod lock;
 pub mod query;
+pub mod relationship;
 pub mod sole;
 pub mod system;
 pub mod tuple;
@@ -57,15 +60,18 @@ impl World
     ///
     /// # Panics
     /// Will panic if mutable internal lock cannot be acquired.
-    pub fn create_entity<Comps>(&mut self, components: Comps)
+    pub fn create_entity<Comps>(&mut self, components: Comps) -> EntityUid
     where
         Comps: ComponentSequence,
     {
-        self.data
+        let (_, entity_uid) = self
+            .data
             .component_storage
             .write_nonblock()
             .expect("Failed to acquire read-write component storage lock")
             .push_entity(components.into_vec());
+
+        entity_uid
     }
 
     /// Adds a globally shared singleton value.
diff --git a/ecs/src/lock.rs b/ecs/src/lock.rs
index c8a8495..b3c9f57 100644
--- a/ecs/src/lock.rs
+++ b/ecs/src/lock.rs
@@ -1,3 +1,4 @@
+use std::mem::transmute;
 use std::ops::{Deref, DerefMut};
 use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard, TryLockError};
 
@@ -73,6 +74,21 @@ where
     inner: RwLockReadGuard<'guard, Value>,
 }
 
+impl<'guard, Value> ReadGuard<'guard, Value>
+where
+    Value: TypeName,
+{
+    /// Converts the `ReadGuard` to a `ReadGuard` with a possibly longer lifetime.
+    ///
+    /// # Safety
+    /// The returned `ReadGuard` must **NOT** be used for longer than the original
+    /// lifetime.
+    pub unsafe fn upgrade_lifetime<'new>(self) -> ReadGuard<'new, Value>
+    {
+        unsafe { transmute(self) }
+    }
+}
+
 impl<'guard, Value> Deref for ReadGuard<'guard, Value>
 where
     Value: TypeName,
diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index bc98ac0..beb2478 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -123,6 +123,7 @@ 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
@@ -137,6 +138,10 @@ where
         );
 
         component_storage_lock.add_archetype_lookup_entry(Comps::metadata());
+
+        drop(component_storage_lock);
+
+        Comps::prepare(world_data);
     }
 }
 
diff --git a/ecs/src/relationship.rs b/ecs/src/relationship.rs
new file mode 100644
index 0000000..9004f0f
--- /dev/null
+++ b/ecs/src/relationship.rs
@@ -0,0 +1,190 @@
+use std::any::{type_name, Any};
+use std::marker::PhantomData;
+use std::sync::{Arc, Weak};
+
+use crate::component::storage::Storage as ComponentStorage;
+use crate::component::{
+    is_optional as component_is_optional,
+    Component,
+    FromOptional as ComponentFromOptional,
+    Id as ComponentId,
+    Metadata as ComponentMetadata,
+};
+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};
+
+#[derive(Debug)]
+pub struct Relationship<Kind, ComponentT: Component>
+{
+    entity_uid: EntityUid,
+    component_storage: Weak<Lock<ComponentStorage>>,
+    _pd: PhantomData<(Kind, ComponentT)>,
+}
+
+impl<Kind, ComponentT> Relationship<Kind, ComponentT>
+where
+    ComponentT: Component,
+{
+    pub fn new(world: &World, entity_uid: EntityUid) -> Self
+    {
+        Self {
+            entity_uid,
+            component_storage: Arc::downgrade(&world.data.component_storage),
+            _pd: PhantomData,
+        }
+    }
+}
+
+impl<Kind, ComponentT> Component for Relationship<Kind, ComponentT>
+where
+    Kind: 'static,
+    ComponentT: Component,
+{
+    type Component = Self;
+    type RefMut<'component> = Relation<'component, Kind, ComponentT>;
+
+    fn id(&self) -> ComponentId
+    {
+        ComponentId::of::<Self>()
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any
+    {
+        self
+    }
+
+    fn as_any(&self) -> &dyn Any
+    {
+        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>
+where
+    Kind: 'static,
+    ComponentT: Component,
+{
+}
+
+impl<Kind, ComponentT: Component> TypeName for Relationship<Kind, ComponentT>
+where
+    ComponentT: Component,
+{
+    fn type_name(&self) -> &'static str
+    {
+        type_name::<Self>()
+    }
+}
+
+pub struct Relation<'rel_comp, Kind, ComponentT>
+where
+    Kind: 'static,
+    ComponentT: Component,
+{
+    component_storage_lock: ReadGuard<'static, ComponentStorage>,
+    relationship_comp: ComponentRefMut<'rel_comp, Relationship<Kind, ComponentT>>,
+    _component_storage: Arc<Lock<ComponentStorage>>,
+}
+
+impl<'rel_comp, Kind, ComponentT> ComponentFromOptional<'rel_comp>
+    for Relation<'rel_comp, Kind, ComponentT>
+where
+    ComponentT: Component,
+{
+    fn from_optional_component(
+        optional_component: Option<
+            crate::lock::WriteGuard<'rel_comp, Box<dyn Component>>,
+        >,
+    ) -> Self
+    {
+        let relationship_comp =
+            ComponentRefMut::<Relationship<Kind, ComponentT>>::from_optional_component(
+                optional_component,
+            );
+
+        let component_storage = relationship_comp
+            .component_storage
+            .upgrade()
+            .expect("World has been dropped");
+
+        let component_storage_lock = component_storage
+            .read_nonblock()
+            .expect("Failed to aquire read-only component storage lock");
+
+        Self {
+            relationship_comp,
+            // SAFETY: The component lock is not used for longer than the original
+            // lifetime
+            component_storage_lock: unsafe { component_storage_lock.upgrade_lifetime() },
+            _component_storage: component_storage,
+        }
+    }
+}
+
+impl<'rel_comp, Kind, ComponentT> Relation<'rel_comp, Kind, ComponentT>
+where
+    ComponentT: Component,
+{
+    /// Retrieves the related-to component.
+    pub fn get(&self) -> Option<ComponentRefMut<'_, ComponentT>>
+    {
+        let mut archetype_iter =
+            self.component_storage_lock
+                .find_entities([ComponentMetadata {
+                    id: ComponentId::of::<ComponentT>(),
+                    is_optional: component_is_optional::<ComponentT>().into(),
+                }]);
+
+        let (entity, archetype) = archetype_iter.find_map(|archetype| {
+            let Some(entity) = archetype.get_entity(self.relationship_comp.entity_uid)
+            else {
+                return None;
+            };
+
+            Some((entity, archetype))
+        })?;
+
+        let component_index =
+            archetype.get_index_for_component(&ComponentId::of::<ComponentT>())?;
+
+        let component = ComponentRefMut::new(
+            entity
+                .get(component_index)?
+                .component
+                .write_nonblock()
+                .unwrap_or_else(|_| {
+                    panic!(
+                        "Failed to aquire read-write lock of component {}",
+                        type_name::<ComponentT>()
+                    )
+                }),
+        );
+
+        Some(component)
+    }
+}