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

author: HampusM <hampus@hampusmat.com> 2025-03-17 22:08:45 +0100
committer: HampusM <hampus@hampusmat.com> 2025-03-17 23:08:53 +0100
commit: ee7c0cb713891ae480407f56823289f3fe3b9807 (patch)
tree: 8c0fa0bd98050ccd74ab0c9100d12c057bdea6cd /ecs/src
parent: 64eddc633cea0f4bc5603cc2d4c4c6eafac5c177 (diff)
7 files changed, 611 insertions, 130 deletions
diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index b2ecf80..f77b55a 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -228,6 +228,11 @@ pub struct Metadata
 
 impl Metadata
 {
+    pub fn new_non_optional(id: Uid) -> Self
+    {
+        Self { id, is_optional: false }
+    }
+
     pub fn get<ComponentT: Component + ?Sized>(component: &ComponentT) -> Self
     {
         Self {
diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 0e1a03d..b68f89d 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -1,7 +1,9 @@
 use std::any::type_name;
-use std::iter::{Chain, Flatten};
+use std::array::IntoIter as ArrayIter;
+use std::cell::RefCell;
+use std::vec::IntoIter as VecIntoIter;
 
-use hashbrown::{HashMap, HashSet};
+use hashbrown::HashMap;
 
 use crate::component::storage::archetype::{
     Archetype,
@@ -9,13 +11,15 @@ use crate::component::storage::archetype::{
     Id as ArchetypeId,
 };
 use crate::component::storage::graph::{
-    ArchetypeAddEdgeRecursiveIter,
+    ArchetypeAddEdgeDfsIter,
+    ArchetypeAddEdgeDfsIterResult,
     ArchetypeEdges,
     Graph,
 };
 use crate::component::{Component, Metadata as ComponentMetadata};
 use crate::type_name::TypeName;
 use crate::uid::{Kind as UidKind, Uid};
+use crate::util::{BorrowedOrOwned, Either, StreamingIterator, VecExt};
 use crate::EntityComponent;
 
 pub mod archetype;
@@ -27,27 +31,50 @@ pub struct Storage
 {
     graph: Graph,
     entity_archetype_lookup: HashMap<Uid, ArchetypeId>,
+    imaginary_archetypes: RefCell<Vec<ImaginaryArchetype>>,
 }
 
 impl Storage
 {
-    pub fn iter_archetypes_with_comps(
+    pub fn search_archetypes(
         &self,
-        component_metadata: impl AsRef<[ComponentMetadata]>,
+        component_metadata: &[ComponentMetadata],
     ) -> ArchetypeRefIter<'_>
     {
-        let archetype_id = ArchetypeId::from_components_metadata(&component_metadata);
+        let archetype_id = ArchetypeId::from_components_metadata(component_metadata);
+
+        let Some(add_edge_recursive_iter) =
+            self.graph.dfs_archetype_add_edges(archetype_id)
+        else {
+            self.imaginary_archetypes
+                .borrow_mut()
+                .push(ImaginaryArchetype {
+                    id: archetype_id,
+                    component_ids: component_metadata
+                        .iter()
+                        .filter_map(|comp_metadata| {
+                            if comp_metadata.is_optional {
+                                return None;
+                            }
+
+                            Some(comp_metadata.id)
+                        })
+                        .collect::<Vec<_>>(),
+                });
+
+            let found_archetypes = self.find_all_archetype_with_comps(component_metadata);
+
+            return ArchetypeRefIter {
+                storage: self,
+                pre_iter: Either::B(found_archetypes.clone().into_iter()),
+                dfs_iter: ArchetypeAddEdgeDfsIter::new(&self.graph, &found_archetypes),
+            };
+        };
 
         ArchetypeRefIter {
             storage: self,
-            iter: [archetype_id].into_iter().chain(
-                self.graph
-                    .recurse_archetype_add_edges(archetype_id)
-                    .into_iter()
-                    .flatten()
-                    .flatten(),
-            ),
-            visited: HashSet::new(),
+            pre_iter: Either::A([archetype_id].into_iter()),
+            dfs_iter: add_edge_recursive_iter,
         }
     }
 
@@ -288,6 +315,69 @@ impl Storage
 
         Ok(())
     }
+
+    pub fn create_imaginary_archetypes(&mut self)
+    {
+        for imaginary_archetype in self.imaginary_archetypes.get_mut().drain(..) {
+            if self.graph.contains_archetype(imaginary_archetype.id) {
+                continue;
+            }
+
+            self.graph
+                .create_node(imaginary_archetype.id, &imaginary_archetype.component_ids);
+        }
+    }
+
+    fn find_all_archetype_with_comps(
+        &self,
+        component_metadata: &[ComponentMetadata],
+    ) -> Vec<ArchetypeId>
+    {
+        let comp_metadata_not_opt_cnt = component_metadata
+            .iter()
+            .filter(|comp_metadata| !comp_metadata.is_optional)
+            .count();
+
+        let Some(mut search_iter) =
+            self.graph.dfs_archetype_add_edges(ArchetypeId::new(&[]))
+        else {
+            // If the root archetype doesn't exist, no other archetype can exist either
+            return Vec::new();
+        };
+
+        let mut found = Vec::<ArchetypeId>::new();
+
+        while let Some(node_id) = search_iter.streaming_next() {
+            let ArchetypeAddEdgeDfsIterResult::AddEdge(node_id) = node_id else {
+                continue;
+            };
+
+            let node = self
+                .graph
+                .get_node_by_id(node_id)
+                .expect("Graph node found through DFS doesn't exist");
+
+            if node.archetype().component_cnt() < comp_metadata_not_opt_cnt {
+                continue;
+            }
+
+            if !component_metadata
+                .iter()
+                .filter(|comp_metadata| !comp_metadata.is_optional)
+                .all(|comp_metadata| {
+                    node.archetype().has_component_with_id(comp_metadata.id)
+                })
+            {
+                continue;
+            }
+
+            found.push(node.archetype().id());
+
+            search_iter.pop();
+        }
+
+        found
+    }
 }
 
 impl TypeName for Storage
@@ -302,11 +392,8 @@ impl TypeName for Storage
 pub struct ArchetypeRefIter<'storage>
 {
     storage: &'storage Storage,
-    iter: Chain<
-        std::array::IntoIter<ArchetypeId, 1>,
-        Flatten<Flatten<std::option::IntoIter<ArchetypeAddEdgeRecursiveIter<'storage>>>>,
-    >,
-    visited: HashSet<ArchetypeId>,
+    pre_iter: Either<ArrayIter<ArchetypeId, 1>, VecIntoIter<ArchetypeId>>,
+    dfs_iter: ArchetypeAddEdgeDfsIter<'storage>,
 }
 
 impl<'component_storage> Iterator for ArchetypeRefIter<'component_storage>
@@ -315,15 +402,122 @@ impl<'component_storage> Iterator for ArchetypeRefIter<'component_storage>
 
     fn next(&mut self) -> Option<Self::Item>
     {
-        let archetype_id = self
-            .iter
-            .find(|archetype_id| !self.visited.contains(archetype_id))?;
+        if let Some(pre_iter_archetype_id) = self.pre_iter.next() {
+            return Some(
+                self.storage
+                    .get_archetype_by_id(pre_iter_archetype_id)
+                    .expect("Archetype should exist"),
+            );
+        }
 
-        let archetype = self.storage.get_archetype_by_id(archetype_id)?;
+        let archetype_id = loop {
+            match self.dfs_iter.streaming_find(|res| {
+                matches!(
+                    res,
+                    ArchetypeAddEdgeDfsIterResult::AddEdge { .. }
+                        | ArchetypeAddEdgeDfsIterResult::AddEdgeArchetypeNotFound { .. }
+                )
+            })? {
+                ArchetypeAddEdgeDfsIterResult::AddEdge(add_edge_archetype_id) => {
+                    break add_edge_archetype_id
+                }
+                ArchetypeAddEdgeDfsIterResult::AddEdgeArchetypeNotFound {
+                    archetype,
+                    add_edge_archetype_id,
+                    add_edge_component_id,
+                } => {
+                    let mut add_edge_archetype_comps = archetype
+                        .component_ids_sorted()
+                        .map(|id| ComponentMetadata { id, is_optional: false })
+                        .collect::<Vec<_>>();
+
+                    add_edge_archetype_comps.insert_at_partition_point_by_key(
+                        ComponentMetadata::new_non_optional(add_edge_component_id),
+                        |comp_metadata| comp_metadata.id,
+                    );
+
+                    self.storage.imaginary_archetypes.borrow_mut().push(
+                        ImaginaryArchetype {
+                            id: add_edge_archetype_id,
+                            component_ids: add_edge_archetype_comps
+                                .iter()
+                                .map(|comp_metadata| comp_metadata.id)
+                                .collect::<Vec<_>>(),
+                        },
+                    );
+
+                    let found =
+                        self.find_edges_of_imaginary_archetype(&add_edge_archetype_comps);
+
+                    self.dfs_iter.push((
+                        BorrowedOrOwned::Owned(Archetype::new(
+                            add_edge_archetype_id,
+                            &add_edge_archetype_comps
+                                .iter()
+                                .map(|metadata| metadata.id)
+                                .collect::<Vec<_>>(),
+                        )),
+                        found.into_iter(),
+                    ));
+
+                    continue;
+                }
+                _ => {
+                    unreachable!();
+                }
+            }
+        };
 
-        self.visited.insert(archetype_id);
+        Some(
+            self.storage
+                .get_archetype_by_id(archetype_id)
+                .expect("Archetype should exist"),
+        )
+    }
+}
 
-        Some(archetype)
+impl<'storage> ArchetypeRefIter<'storage>
+{
+    fn find_edges_of_imaginary_archetype(
+        &self,
+        imaginary_archetype_comps: &[ComponentMetadata],
+    ) -> Vec<(Uid, ArchetypeEdges)>
+    {
+        self.storage
+            .find_all_archetype_with_comps(&imaginary_archetype_comps)
+            .into_iter()
+            .filter_map(|found_id| {
+                let found_archetype = self.storage.get_archetype_by_id(found_id).unwrap();
+
+                if found_archetype.component_cnt() < imaginary_archetype_comps.len() + 1 {
+                    return None;
+                }
+
+                let unique_comp_id = found_archetype
+                    .component_ids_sorted()
+                    .find(|id| {
+                        !imaginary_archetype_comps
+                            .iter()
+                            .any(|comp_metadata| comp_metadata.id == *id)
+                    })
+                    .expect("Oh noooo");
+
+                let mut add_edge_comp_ids = imaginary_archetype_comps
+                    .iter()
+                    .map(|comp_metadata| comp_metadata.id)
+                    .collect::<Vec<_>>();
+
+                add_edge_comp_ids
+                    .insert_at_partition_point_by_key(unique_comp_id, |id| *id);
+
+                let add_edge = ArchetypeId::new(&add_edge_comp_ids);
+
+                Some((
+                    unique_comp_id,
+                    ArchetypeEdges { add: Some(add_edge), remove: None },
+                ))
+            })
+            .collect::<Vec<_>>()
     }
 }
 
@@ -349,6 +543,13 @@ pub enum Error
     },
 }
 
+#[derive(Debug)]
+struct ImaginaryArchetype
+{
+    id: ArchetypeId,
+    component_ids: Vec<Uid>,
+}
+
 #[cfg(test)]
 mod tests
 {
diff --git a/ecs/src/component/storage/archetype.rs b/ecs/src/component/storage/archetype.rs
index ee3d7f8..9bf0feb 100644
--- a/ecs/src/component/storage/archetype.rs
+++ b/ecs/src/component/storage/archetype.rs
@@ -1,4 +1,3 @@
-use std::borrow::Borrow;
 use std::hash::{DefaultHasher, Hash, Hasher};
 use std::slice::Iter as SliceIter;
 
@@ -16,21 +15,24 @@ pub struct Archetype
     entities: Vec<ArchetypeEntity>,
     entity_index_lookup: HashMap<Uid, usize>,
     component_index_lookup: HashMap<Uid, usize>,
+    component_ids: Vec<Uid>,
 }
 
 impl Archetype
 {
-    pub fn new(id: Id, component_ids: impl IntoIterator<Item: Borrow<Uid>>) -> Self
+    pub fn new(id: Id, component_ids: impl AsRef<[Uid]>) -> Self
     {
         Self {
             id,
             entities: Vec::new(),
             entity_index_lookup: HashMap::new(),
             component_index_lookup: component_ids
-                .into_iter()
+                .as_ref()
+                .iter()
                 .enumerate()
-                .map(|(index, id)| (*id.borrow(), index))
+                .map(|(index, id)| (*id, index))
                 .collect(),
+            component_ids: component_ids.as_ref().to_vec(),
         }
     }
 
@@ -45,6 +47,12 @@ impl Archetype
             .keys_is_superset(&other.component_index_lookup)
     }
 
+    pub fn is_subset(&self, other: &Self) -> bool
+    {
+        self.component_index_lookup
+            .keys_is_subset(&other.component_index_lookup)
+    }
+
     pub fn get_entity_by_id(&self, entity_uid: Uid) -> Option<&ArchetypeEntity>
     {
         let index = *self.entity_index_lookup.get(&entity_uid)?;
@@ -118,6 +126,11 @@ impl Archetype
         self.component_index_lookup.keys().copied()
     }
 
+    pub fn component_ids_sorted(&self) -> impl Iterator<Item = Uid> + '_
+    {
+        self.component_ids.iter().copied()
+    }
+
     pub fn has_component_with_id(&self, component_id: Uid) -> bool
     {
         debug_assert_eq!(component_id.kind(), UidKind::Component);
@@ -178,37 +191,36 @@ impl Id
         Self { hash: hasher.finish() }
     }
 
-    pub fn from_components_metadata(
-        components_metadata: &impl AsRef<[ComponentMetadata]>,
+    pub fn from_components_metadata<'a>(
+        components_metadata: impl IntoIterator<Item = &'a ComponentMetadata>,
     ) -> Self
     {
-        if components_metadata.as_ref().len() == 0 {
+        let mut hasher = DefaultHasher::new();
+
+        let mut prev_component_id: Option<Uid> = None;
+
+        let mut comp_metadata_iter = components_metadata.into_iter().peekable();
+
+        if comp_metadata_iter.peek().is_none() {
             return Self { hash: 0 };
         }
 
-        debug_assert!(
-            components_metadata
-                .as_ref()
-                .is_sorted_by_key(|comp_metadata| comp_metadata.id),
-            "Cannot create archetype ID from a unsorted component metadata list"
-        );
-
-        let component_ids =
-            components_metadata
-                .as_ref()
-                .iter()
-                .filter_map(|component_metadata| {
-                    if component_metadata.is_optional {
-                        return None;
-                    }
+        for comp_metadata in comp_metadata_iter {
+            if prev_component_id
+                .is_some_and(|prev_comp_id| comp_metadata.id < prev_comp_id)
+            {
+                panic!(
+                    "Cannot create archetype ID from a unsorted component metadata list"
+                );
+            }
 
-                    Some(component_metadata.id)
-                });
+            prev_component_id = Some(comp_metadata.id);
 
-        let mut hasher = DefaultHasher::new();
+            if comp_metadata.is_optional {
+                continue;
+            }
 
-        for component_id in component_ids {
-            component_id.hash(&mut hasher);
+            comp_metadata.id.hash(&mut hasher);
         }
 
         Self { hash: hasher.finish() }
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,
+    },
+}
diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index dc31daf..e2ebb21 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -212,6 +212,9 @@ impl World
 
         self.perform_phases();
 
+        self.lock_component_storage_rw()
+            .create_imaginary_archetypes();
+
         self.perform_queued_actions();
 
         if self.stop.load(Ordering::Relaxed) {
diff --git a/ecs/src/query/flexible.rs b/ecs/src/query/flexible.rs
index 129ee66..3bb8fd6 100644
--- a/ecs/src/query/flexible.rs
+++ b/ecs/src/query/flexible.rs
@@ -35,7 +35,7 @@ where
         Iter {
             iter: self
                 .component_storage
-                .iter_archetypes_with_comps(&self.comp_metadata)
+                .search_archetypes(&self.comp_metadata.as_ref())
                 .map(
                     (|archetype| {
                         repeat_n(archetype, archetype.entity_cnt())
diff --git a/ecs/src/util.rs b/ecs/src/util.rs
index 71021cd..fbd33fa 100644
--- a/ecs/src/util.rs
+++ b/ecs/src/util.rs
@@ -1,8 +1,138 @@
 use std::hash::Hash;
-use std::ops::BitAnd;
+use std::mem::transmute;
+use std::ops::{BitAnd, Deref};
 
 use hashbrown::HashMap;
 
+pub trait VecExt<Item>
+{
+    fn insert_at_partition_point_by_key<Key>(
+        &mut self,
+        item: Item,
+        func: impl FnMut(&Item) -> Key,
+    ) where
+        Key: Ord;
+}
+
+impl<Item> VecExt<Item> for Vec<Item>
+{
+    fn insert_at_partition_point_by_key<Key>(
+        &mut self,
+        item: Item,
+        mut func: impl FnMut(&Item) -> Key,
+    ) where
+        Key: Ord,
+    {
+        let key = func(&item);
+
+        let insert_index = self.partition_point(|other_item| func(other_item) <= key);
+
+        self.insert(insert_index, item);
+    }
+}
+
+pub trait StreamingIterator
+{
+    type Item<'a>
+    where
+        Self: 'a;
+
+    fn streaming_next(&mut self) -> Option<Self::Item<'_>>;
+
+    fn streaming_map<NewItem, Func>(self, func: Func) -> StreamingMap<Self, Func>
+    where
+        Self: Sized,
+        Func: FnMut(Self::Item<'_>) -> NewItem,
+    {
+        StreamingMap { iter: self, func }
+    }
+
+    fn streaming_find<'this, Predicate>(
+        &'this mut self,
+        mut predicate: Predicate,
+    ) -> Option<Self::Item<'this>>
+    where
+        Self: Sized,
+        Predicate: FnMut(&Self::Item<'this>) -> bool,
+    {
+        while let Some(item) =
+            unsafe { transmute::<_, Option<Self::Item<'_>>>(self.streaming_next()) }
+        {
+            if predicate(&item) {
+                return Some(item);
+            }
+        }
+
+        None
+    }
+}
+
+pub struct StreamingMap<Iter, Func>
+{
+    iter: Iter,
+    func: Func,
+}
+
+impl<Iter, Func, Item> StreamingIterator for StreamingMap<Iter, Func>
+where
+    Iter: StreamingIterator,
+    Func: FnMut(Iter::Item<'_>) -> Item,
+{
+    type Item<'a>
+        = Item
+    where
+        Iter: 'a,
+        Func: 'a;
+
+    fn streaming_next(&mut self) -> Option<Self::Item<'_>>
+    {
+        Some((self.func)(self.iter.streaming_next()?))
+    }
+}
+
+#[derive(Debug)]
+pub enum BorrowedOrOwned<'a, Value>
+{
+    Borrowned(&'a Value),
+    Owned(Value),
+}
+
+impl<'a, Value> Deref for BorrowedOrOwned<'a, Value>
+{
+    type Target = Value;
+
+    fn deref(&self) -> &Self::Target
+    {
+        match self {
+            Self::Borrowned(value) => value,
+            Self::Owned(value) => value,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum Either<A, B>
+{
+    A(A),
+    B(B),
+}
+
+impl<A, B> Iterator for Either<A, B>
+where
+    A: Iterator,
+    B: Iterator<Item = A::Item>,
+{
+    type Item = A::Item;
+
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        match self {
+            Self::A(a) => a.next(),
+            Self::B(b) => b.next(),
+        }
+    }
+}
+
 pub trait HashMapExt<Key, Value>
 {
     /// Returns true if the keys are a subset of another [`HashMap`]'s keys, i.e., `other`
author	HampusM <hampus@hampusmat.com>	2025-03-17 22:08:45 +0100
committer	HampusM <hampus@hampusmat.com>	2025-03-17 23:08:53 +0100
commit	ee7c0cb713891ae480407f56823289f3fe3b9807 (patch)
tree	8c0fa0bd98050ccd74ab0c9100d12c057bdea6cd /ecs/src
parent	64eddc633cea0f4bc5603cc2d4c4c6eafac5c177 (diff)