fix(ecs): correct oversights made in component storage rewrite

1 files changed, 206 insertions, 76 deletions

diff --git a/ecs/src/component/storage/graph.rs b/ecs/src/component/storage/graph.rs
index fe97933..a8e3f17 100644
--- a/ecs/src/component/storage/graph.rs
+++ b/ecs/src/component/storage/graph.rs
@@ -1,8 +1,10 @@
-use hashbrown::hash_map::Values as HashMapValueIter;
-use hashbrown::HashMap;
+use std::vec::IntoIter as VecIntoIter;
+
+use hashbrown::{HashMap, HashSet};
 
 use crate::component::storage::archetype::{Archetype, Id as ArchetypeId};
 use crate::uid::{Kind as UidKind, Uid};
+use crate::util::{BorrowedOrOwned, StreamingIterator};
 
 #[derive(Debug, Default)]
 pub struct Graph
@@ -26,6 +28,8 @@ impl Graph
         component_ids: &impl AsRef<[Uid]>,
     ) -> &mut ArchetypeNode
     {
+        let exists_before = self.archetype_index_lookup.contains_key(&id);
+
         let index = *self.archetype_index_lookup.entry(id).or_insert_with(|| {
             self.nodes.push(ArchetypeNode {
                 archetype: Archetype::new(id, component_ids.as_ref()),
@@ -35,7 +39,9 @@ impl Graph
             self.nodes.len() - 1
         });
 
-        self.create_missing_edges(id);
+        if !exists_before {
+            self.create_missing_edges(id);
+        }
 
         self.nodes
             .get_mut(index)
@@ -65,56 +71,29 @@ impl Graph
         }))
     }
 
-    pub fn recurse_archetype_add_edges(
+    pub fn dfs_archetype_add_edges(
         &self,
         archetype_id: ArchetypeId,
-    ) -> Option<ArchetypeAddEdgeRecursiveIter>
+    ) -> Option<ArchetypeAddEdgeDfsIter>
     {
-        Some(ArchetypeAddEdgeRecursiveIter {
+        let node = self.get_node_by_id(archetype_id)?;
+
+        Some(ArchetypeAddEdgeDfsIter {
             graph: self,
-            stack: vec![self.get_node_by_id(archetype_id)?.edges.values()],
+            stack: vec![(
+                BorrowedOrOwned::Borrowned(node.archetype()),
+                node.edges
+                    .iter()
+                    .map(|(comp_id, edges)| (*comp_id, edges.clone()))
+                    .collect::<Vec<_>>()
+                    .into_iter(),
+            )],
+            visited: HashSet::new(),
         })
     }
 
     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)
@@ -126,34 +105,101 @@ impl Graph
             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
+            if archetype_node.archetype().component_cnt()
+                > other_archetype_node.archetype().component_cnt()
+                && other_archetype_node
                     .archetype()
-                    .is_superset(&archetype_node.archetype())
+                    .is_subset(archetype_node.archetype())
             {
+                Self::create_missing_subset_node_edges(
+                    archetype_node,
+                    other_archetype_node,
+                );
+
                 continue;
             }
 
-            let extra_comp_id = other_archetype_node
+            if other_archetype_node
+                .archetype()
+                .is_superset(&archetype_node.archetype())
+            {
+                Self::create_missing_superset_node_edges(
+                    archetype_node,
+                    other_archetype_node,
+                );
+            }
+        }
+    }
+
+    fn create_missing_subset_node_edges(
+        target_node: &ArchetypeNode,
+        subset_node: &mut ArchetypeNode,
+    )
+    {
+        let uniq_comp_id = target_node
+            .archetype()
+            .component_ids_sorted()
+            .find(|id| !subset_node.archetype().has_component_with_id(*id))
+            .unwrap();
+
+        subset_node
+            .get_or_insert_edges(uniq_comp_id, ArchetypeEdges::default)
+            .add = Some(subset_node.make_add_edge(uniq_comp_id).0);
+    }
+
+    fn create_missing_superset_node_edges(
+        target_node: &mut ArchetypeNode,
+        superset_node: &mut ArchetypeNode,
+    )
+    {
+        if superset_node.archetype().component_cnt()
+            >= target_node.archetype().component_cnt() + 1
+        {
+            let first_unique_comp_id = superset_node
                 .archetype()
-                .component_ids_unsorted()
-                .find(|comp_id| {
-                    !archetype_node.archetype().has_component_with_id(*comp_id)
+                .component_ids_sorted()
+                .find(|other_archetype_comp_id| {
+                    !target_node
+                        .archetype()
+                        .has_component_with_id(*other_archetype_comp_id)
                 })
-                .expect("Archetype should contain one extra component ID");
+                .or_else(|| {
+                    if target_node.archetype().component_cnt() != 0 {
+                        return None;
+                    }
 
-            other_archetype_node
-                .get_or_insert_edges(extra_comp_id, ArchetypeEdges::default)
-                .remove = Some(archetype_id);
+                    superset_node.archetype().component_ids_sorted().next()
+                })
+                .expect("Not possible");
+
+            target_node
+                .get_or_insert_edges(first_unique_comp_id, ArchetypeEdges::default)
+                .add = Some(target_node.make_add_edge(first_unique_comp_id).0);
 
-            archetype_node
-                .get_or_insert_edges(extra_comp_id, ArchetypeEdges::default)
-                .add = Some(other_archetype_node.archetype().id());
+            return;
+        }
+
+        if superset_node.archetype().component_cnt()
+            != target_node.archetype().component_cnt() + 1
+        {
+            return;
         }
+
+        let extra_comp_id = superset_node
+            .archetype()
+            .component_ids_unsorted()
+            .find(|comp_id| !target_node.archetype().has_component_with_id(*comp_id))
+            .expect("Archetype should contain one extra component ID");
+
+        superset_node
+            .get_or_insert_edges(extra_comp_id, ArchetypeEdges::default)
+            .remove = Some(target_node.archetype().id());
+
+        target_node
+            .get_or_insert_edges(extra_comp_id, ArchetypeEdges::default)
+            .add = Some(superset_node.archetype().id());
     }
 }
 
@@ -225,7 +271,7 @@ impl ArchetypeNode
     }
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug, Default, Clone)]
 pub struct ArchetypeEdges
 {
     pub add: Option<ArchetypeId>,
@@ -233,37 +279,121 @@ pub struct ArchetypeEdges
 }
 
 #[derive(Debug)]
-pub struct ArchetypeAddEdgeRecursiveIter<'graph>
+pub struct ArchetypeAddEdgeDfsIter<'graph>
 {
     graph: &'graph Graph,
-    stack: Vec<HashMapValueIter<'graph, Uid, ArchetypeEdges>>,
+    stack: Vec<(
+        BorrowedOrOwned<'graph, Archetype>,
+        VecIntoIter<(Uid, ArchetypeEdges)>,
+    )>,
+    visited: HashSet<ArchetypeId>,
+}
+
+impl<'graph> ArchetypeAddEdgeDfsIter<'graph>
+{
+    pub fn new(graph: &'graph Graph, start_nodes: &[ArchetypeId]) -> Self
+    {
+        Self {
+            graph,
+            stack: start_nodes
+                .into_iter()
+                .map(|start_node_id| {
+                    let start_node = graph
+                        .get_node_by_id(*start_node_id)
+                        .expect("Start node does not exist");
+
+                    (
+                        BorrowedOrOwned::Borrowned(start_node.archetype()),
+                        start_node
+                            .edges
+                            .iter()
+                            .map(|(comp_id, edges)| (*comp_id, edges.clone()))
+                            .collect::<Vec<_>>()
+                            .into_iter(),
+                    )
+                })
+                .collect(),
+            visited: HashSet::from_iter(start_nodes.iter().copied()),
+        }
+    }
+
+    pub fn push(
+        &mut self,
+        item: (
+            BorrowedOrOwned<'graph, Archetype>,
+            VecIntoIter<(Uid, ArchetypeEdges)>,
+        ),
+    )
+    {
+        self.stack.push(item);
+    }
+
+    pub fn pop(&mut self)
+    {
+        self.stack.pop();
+    }
 }
 
-impl<'graph> Iterator for ArchetypeAddEdgeRecursiveIter<'graph>
+impl<'graph> StreamingIterator for ArchetypeAddEdgeDfsIter<'graph>
 {
-    type Item = Option<ArchetypeId>;
+    type Item<'a>
+        = ArchetypeAddEdgeDfsIterResult<'graph, 'a>
+    where
+        Self: 'a;
 
-    fn next(&mut self) -> Option<Self::Item>
+    fn streaming_next(&mut self) -> Option<Self::Item<'_>>
     {
-        let iter = self.stack.last_mut()?;
+        let (_, edges_iter) = self.stack.last_mut()?;
 
-        let Some(edges) = iter.next() else {
+        let Some((component_id, edges)) = edges_iter.next() else {
             self.stack.pop();
 
-            return Some(None);
+            return Some(ArchetypeAddEdgeDfsIterResult::NoEdgesLeftForArchetype);
         };
 
         let Some(add_edge) = edges.add else {
-            return Some(None);
+            return Some(ArchetypeAddEdgeDfsIterResult::NoAddEdge);
         };
 
-        let add_edge_archetype = self
-            .graph
-            .get_node_by_id(add_edge)
-            .expect("Add edge archetype not found");
+        if self.visited.contains(&add_edge) {
+            return Some(ArchetypeAddEdgeDfsIterResult::AddEdgeAlreadyVisited);
+        }
 
-        self.stack.push(add_edge_archetype.edges.values());
+        self.visited.insert(add_edge);
 
-        Some(Some(add_edge))
+        let Some(add_edge_archetype) = self.graph.get_node_by_id(add_edge) else {
+            return Some(ArchetypeAddEdgeDfsIterResult::AddEdgeArchetypeNotFound {
+                archetype: &self.stack.last().unwrap().0,
+                add_edge_archetype_id: add_edge,
+                add_edge_component_id: component_id,
+            });
+        };
+
+        self.stack.push((
+            BorrowedOrOwned::Borrowned(add_edge_archetype.archetype()),
+            add_edge_archetype
+                .edges
+                .iter()
+                .map(|(comp_id, edges)| (*comp_id, edges.clone()))
+                .collect::<Vec<_>>()
+                .into_iter(),
+        ));
+
+        Some(ArchetypeAddEdgeDfsIterResult::AddEdge(add_edge))
     }
 }
+
+#[derive(Debug)]
+pub enum ArchetypeAddEdgeDfsIterResult<'graph, 'iter>
+{
+    AddEdge(ArchetypeId),
+    NoEdgesLeftForArchetype,
+    NoAddEdge,
+    AddEdgeAlreadyVisited,
+    AddEdgeArchetypeNotFound
+    {
+        archetype: &'iter BorrowedOrOwned<'graph, Archetype>,
+        add_edge_archetype_id: ArchetypeId,
+        add_edge_component_id: Uid,
+    },
+}