2 files changed, 95 insertions, 11 deletions

diff --git a/ecs/src/component/storage/archetype.rs b/ecs/src/component/storage/archetype.rs
index f6f8132..5306cf9 100644
--- a/ecs/src/component/storage/archetype.rs
+++ b/ecs/src/component/storage/archetype.rs
@@ -3,16 +3,16 @@ use std::slice::Iter as SliceIter;
 
 use hashbrown::HashMap;
 
-use crate::component::Metadata as ComponentMetadata;
+use crate::component::{Component, Metadata as ComponentMetadata};
+use crate::lock::Lock;
 use crate::uid::{Kind as UidKind, Uid};
 use crate::util::HashMapExt;
-use crate::EntityComponent;
 
 #[derive(Debug)]
 pub struct Archetype
 {
     id: Id,
-    entities: Vec<ArchetypeEntity>,
+    entities: Vec<Entity>,
     entity_index_lookup: HashMap<Uid, usize>,
     component_index_lookup: HashMap<Uid, usize>,
     component_ids: Vec<Uid>,
@@ -53,7 +53,7 @@ impl Archetype
             .keys_is_subset(&other.component_index_lookup)
     }
 
-    pub fn get_entity_by_id(&self, entity_uid: Uid) -> Option<&ArchetypeEntity>
+    pub fn get_entity_by_id(&self, entity_uid: Uid) -> Option<&Entity>
     {
         let index = *self.entity_index_lookup.get(&entity_uid)?;
 
@@ -64,7 +64,7 @@ impl Archetype
         }))
     }
 
-    pub fn push_entity(&mut self, entity: ArchetypeEntity)
+    pub fn push_entity(&mut self, entity: Entity)
     {
         self.entity_index_lookup
             .insert(entity.uid, self.entities.len());
@@ -72,7 +72,7 @@ impl Archetype
         self.entities.push(entity);
     }
 
-    pub fn remove_entity(&mut self, entity_uid: Uid) -> Option<ArchetypeEntity>
+    pub fn remove_entity(&mut self, entity_uid: Uid) -> Option<Entity>
     {
         //debug_assert_eq!(entity_uid.kind(), UidKind::Entity);
 
@@ -142,12 +142,12 @@ impl Archetype
 #[derive(Debug)]
 pub struct EntityIter<'archetype>
 {
-    iter: SliceIter<'archetype, ArchetypeEntity>,
+    iter: SliceIter<'archetype, Entity>,
 }
 
 impl<'archetype> Iterator for EntityIter<'archetype>
 {
-    type Item = &'archetype ArchetypeEntity;
+    type Item = &'archetype Entity;
 
     fn next(&mut self) -> Option<Self::Item>
     {
@@ -156,10 +156,82 @@ impl<'archetype> Iterator for EntityIter<'archetype>
 }
 
 #[derive(Debug)]
-pub struct ArchetypeEntity
+pub struct Entity
 {
-    pub uid: Uid,
-    pub components: Vec<EntityComponent>,
+    uid: Uid,
+    components: Vec<EntityComponent>,
+}
+
+impl Entity
+{
+    pub fn new(uid: Uid, components: impl IntoIterator<Item = EntityComponent>) -> Self
+    {
+        Self {
+            uid,
+            components: components.into_iter().collect(),
+        }
+    }
+
+    pub fn uid(&self) -> Uid
+    {
+        self.uid
+    }
+
+    pub fn components(&self) -> &[EntityComponent]
+    {
+        &self.components
+    }
+
+    pub fn remove_component(&mut self, component_id: Uid, archetype: &Archetype)
+    {
+        let index = archetype
+            .get_index_for_component(component_id)
+            .expect("Archetype should contain component");
+
+        self.components.remove(index);
+    }
+
+    pub fn insert_component(
+        &mut self,
+        component_id: Uid,
+        component: EntityComponent,
+        archetype: &Archetype,
+    )
+    {
+        let index = archetype
+            .get_index_for_component(component_id)
+            .expect("Archetype should contain component");
+
+        self.components.insert(index, component);
+    }
+}
+
+#[derive(Debug)]
+pub struct EntityComponent
+{
+    name: &'static str,
+    component: Lock<Box<dyn Component>>,
+}
+
+impl EntityComponent
+{
+    pub fn new(component: Box<dyn Component>) -> Self
+    {
+        Self {
+            name: component.type_name(),
+            component: Lock::new(component),
+        }
+    }
+
+    pub fn name(&self) -> &str
+    {
+        self.name
+    }
+
+    pub fn component(&self) -> &Lock<Box<dyn Component>>
+    {
+        &self.component
+    }
 }
 
 /// Archetype ID.
diff --git a/ecs/src/component/storage/graph.rs b/ecs/src/component/storage/graph.rs
index a221877..11160e7 100644
--- a/ecs/src/component/storage/graph.rs
+++ b/ecs/src/component/storage/graph.rs
@@ -71,6 +71,12 @@ impl Graph
         }))
     }
 
+    #[cfg(feature = "vizoxide")]
+    pub fn iter_nodes(&self) -> impl Iterator<Item = &ArchetypeNode>
+    {
+        self.nodes.iter()
+    }
+
     pub fn dfs_archetype_add_edges(
         &self,
         archetype_id: ArchetypeId,
@@ -233,6 +239,12 @@ impl ArchetypeNode
         self.edges.entry(component_id).or_insert_with(insert_fn)
     }
 
+    #[cfg(feature = "vizoxide")]
+    pub fn iter_edges(&self) -> impl Iterator<Item = (&Uid, &ArchetypeEdges)>
+    {
+        self.edges.iter()
+    }
+
     pub fn get_edges_mut(&mut self, component_id: Uid) -> Option<&mut ArchetypeEdges>
     {
         debug_assert_eq!(component_id.kind(), UidKind::Component);