1 files changed, 269 insertions, 0 deletions
diff --git a/ecs/src/component/storage/graph.rs b/ecs/src/component/storage/graph.rs
new file mode 100644
index 0000000..fe97933
--- /dev/null
+++ b/ecs/src/component/storage/graph.rs
@@ -0,0 +1,269 @@
+use hashbrown::hash_map::Values as HashMapValueIter;
+use hashbrown::HashMap;
+
+use crate::component::storage::archetype::{Archetype, Id as ArchetypeId};
+use crate::uid::{Kind as UidKind, Uid};
+
+#[derive(Debug, Default)]
+pub struct Graph
+{
+    nodes: Vec<ArchetypeNode>,
+    archetype_index_lookup: HashMap<ArchetypeId, usize>,
+}
+
+impl Graph
+{
+    pub fn create_node(&mut self, id: ArchetypeId, component_ids: &impl AsRef<[Uid]>)
+    {
+        debug_assert!(!self.contains_archetype(id));
+
+        let _ = self.get_or_create_node(id, component_ids);
+    }
+
+    pub fn get_or_create_node(
+        &mut self,
+        id: ArchetypeId,
+        component_ids: &impl AsRef<[Uid]>,
+    ) -> &mut ArchetypeNode
+    {
+        let index = *self.archetype_index_lookup.entry(id).or_insert_with(|| {
+            self.nodes.push(ArchetypeNode {
+                archetype: Archetype::new(id, component_ids.as_ref()),
+                edges: HashMap::new(),
+            });
+
+            self.nodes.len() - 1
+        });
+
+        self.create_missing_edges(id);
+
+        self.nodes
+            .get_mut(index)
+            .expect("Archetype index from lookup is out of bounds")
+    }
+
+    pub fn contains_archetype(&self, id: ArchetypeId) -> bool
+    {
+        self.archetype_index_lookup.contains_key(&id)
+    }
+
+    pub fn get_node_by_id(&self, id: ArchetypeId) -> Option<&ArchetypeNode>
+    {
+        let index = self.archetype_index_lookup.get(&id)?;
+
+        Some(self.nodes.get(*index).unwrap_or_else(|| {
+            panic!("In invalid state! Index of archetype with ID {id:?} is out of bounds")
+        }))
+    }
+
+    pub fn get_node_by_id_mut(&mut self, id: ArchetypeId) -> Option<&mut ArchetypeNode>
+    {
+        let index = self.archetype_index_lookup.get(&id)?;
+
+        Some(self.nodes.get_mut(*index).unwrap_or_else(|| {
+            panic!("In invalid state! Index of archetype with ID {id:?} is out of bounds")
+        }))
+    }
+
+    pub fn recurse_archetype_add_edges(
+        &self,
+        archetype_id: ArchetypeId,
+    ) -> Option<ArchetypeAddEdgeRecursiveIter>
+    {
+        Some(ArchetypeAddEdgeRecursiveIter {
+            graph: self,
+            stack: vec![self.get_node_by_id(archetype_id)?.edges.values()],
+        })
+    }
+
+    fn create_missing_edges(&mut self, archetype_id: ArchetypeId)
+    {
+        let archetype_node = self
+            .get_node_by_id(archetype_id)
+            .expect("Archetype should exist");
+
+        let mut comp_ids = archetype_node
+            .archetype()
+            .component_ids_unsorted()
+            .collect::<Vec<_>>();
+
+        comp_ids.sort();
+
+        for component_id_index in 0..archetype_node.archetype().component_cnt() {
+            let mut comp_ids_without_comp = comp_ids.clone();
+
+            let removed_comp_id = comp_ids_without_comp.remove(component_id_index);
+
+            let archetype_without_comp_id = ArchetypeId::new(&comp_ids_without_comp);
+
+            let Some(archetype_node_without_comp) =
+                self.get_node_by_id_mut(archetype_without_comp_id)
+            else {
+                continue;
+            };
+
+            archetype_node_without_comp
+                .get_or_insert_edges(removed_comp_id, ArchetypeEdges::default)
+                .add = Some(archetype_id);
+
+            let archetype_node = self
+                .get_node_by_id_mut(archetype_id)
+                .expect("Archetype should exist");
+
+            archetype_node
+                .get_or_insert_edges(removed_comp_id, ArchetypeEdges::default)
+                .remove = Some(archetype_without_comp_id);
+        }
+
+        let archetype_node_index = *self
+            .archetype_index_lookup
+            .get(&archetype_id)
+            .expect("Archetype should exist");
+
+        let (nodes_before, nodes_rest) = self.nodes.split_at_mut(archetype_node_index);
+
+        let ([archetype_node], nodes_after) = nodes_rest.split_at_mut(1) else {
+            unreachable!();
+        };
+
+        let expected_component_cnt = archetype_node.archetype().component_cnt() + 1;
+
+        for other_archetype_node in nodes_before.iter_mut().chain(nodes_after.iter_mut())
+        {
+            if other_archetype_node.archetype().component_cnt() != expected_component_cnt
+                || !other_archetype_node
+                    .archetype()
+                    .is_superset(&archetype_node.archetype())
+            {
+                continue;
+            }
+
+            let extra_comp_id = other_archetype_node
+                .archetype()
+                .component_ids_unsorted()
+                .find(|comp_id| {
+                    !archetype_node.archetype().has_component_with_id(*comp_id)
+                })
+                .expect("Archetype should contain one extra component ID");
+
+            other_archetype_node
+                .get_or_insert_edges(extra_comp_id, ArchetypeEdges::default)
+                .remove = Some(archetype_id);
+
+            archetype_node
+                .get_or_insert_edges(extra_comp_id, ArchetypeEdges::default)
+                .add = Some(other_archetype_node.archetype().id());
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct ArchetypeNode
+{
+    archetype: Archetype,
+    edges: HashMap<Uid, ArchetypeEdges>,
+}
+
+impl ArchetypeNode
+{
+    pub fn archetype(&self) -> &Archetype
+    {
+        &self.archetype
+    }
+
+    pub fn archetype_mut(&mut self) -> &mut Archetype
+    {
+        &mut self.archetype
+    }
+
+    pub fn get_or_insert_edges(
+        &mut self,
+        component_id: Uid,
+        insert_fn: impl FnOnce() -> ArchetypeEdges,
+    ) -> &mut ArchetypeEdges
+    {
+        debug_assert_eq!(component_id.kind(), UidKind::Component);
+
+        self.edges.entry(component_id).or_insert_with(insert_fn)
+    }
+
+    pub fn get_edges_mut(&mut self, component_id: Uid) -> Option<&mut ArchetypeEdges>
+    {
+        debug_assert_eq!(component_id.kind(), UidKind::Component);
+
+        self.edges.get_mut(&component_id)
+    }
+
+    pub fn make_add_edge(&self, component_id: Uid) -> (ArchetypeId, Vec<Uid>)
+    {
+        let mut edge_comp_ids = self
+            .archetype()
+            .component_ids_unsorted()
+            .chain([component_id])
+            .collect::<Vec<_>>();
+
+        edge_comp_ids.sort();
+
+        let add_edge_id = ArchetypeId::new(&edge_comp_ids);
+
+        (add_edge_id, edge_comp_ids)
+    }
+
+    pub fn make_remove_edge(&self, component_id: Uid) -> (ArchetypeId, Vec<Uid>)
+    {
+        let mut edge_comp_ids = self
+            .archetype()
+            .component_ids_unsorted()
+            .filter(|id| *id != component_id)
+            .collect::<Vec<_>>();
+
+        edge_comp_ids.sort();
+
+        let remove_edge_id = ArchetypeId::new(&edge_comp_ids);
+
+        (remove_edge_id, edge_comp_ids)
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct ArchetypeEdges
+{
+    pub add: Option<ArchetypeId>,
+    pub remove: Option<ArchetypeId>,
+}
+
+#[derive(Debug)]
+pub struct ArchetypeAddEdgeRecursiveIter<'graph>
+{
+    graph: &'graph Graph,
+    stack: Vec<HashMapValueIter<'graph, Uid, ArchetypeEdges>>,
+}
+
+impl<'graph> Iterator for ArchetypeAddEdgeRecursiveIter<'graph>
+{
+    type Item = Option<ArchetypeId>;
+
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        let iter = self.stack.last_mut()?;
+
+        let Some(edges) = iter.next() else {
+            self.stack.pop();
+
+            return Some(None);
+        };
+
+        let Some(add_edge) = edges.add else {
+            return Some(None);
+        };
+
+        let add_edge_archetype = self
+            .graph
+            .get_node_by_id(add_edge)
+            .expect("Add edge archetype not found");
+
+        self.stack.push(add_edge_archetype.edges.values());
+
+        Some(Some(add_edge))
+    }
+}