1 files changed, 24 insertions, 12 deletions

diff --git a/ecs/src/component/storage/graph.rs b/ecs/src/component/storage/graph.rs
index 11160e7..29fa937 100644
--- a/ecs/src/component/storage/graph.rs
+++ b/ecs/src/component/storage/graph.rs
@@ -140,19 +140,31 @@ impl Graph
     }
 
     fn create_missing_subset_node_edges(
-        target_node: &ArchetypeNode,
+        target_node: &mut 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))
+            .find(|id| {
+                !subset_node
+                    .archetype()
+                    .contains_component_with_exact_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);
+
+        if target_node.archetype().component_cnt()
+            == subset_node.archetype().component_cnt() + 1
+        {
+            target_node
+                .get_or_insert_edges(uniq_comp_id, ArchetypeEdges::default)
+                .remove = Some(subset_node.archetype().id());
+        }
     }
 
     fn create_missing_superset_node_edges(
@@ -169,7 +181,7 @@ impl Graph
                 .find(|other_archetype_comp_id| {
                     !target_node
                         .archetype()
-                        .has_component_with_id(*other_archetype_comp_id)
+                        .contains_component_with_exact_id(*other_archetype_comp_id)
                 })
                 .or_else(|| {
                     if target_node.archetype().component_cnt() != 0 {
@@ -196,7 +208,11 @@ impl Graph
         let extra_comp_id = superset_node
             .archetype()
             .component_ids_unsorted()
-            .find(|comp_id| !target_node.archetype().has_component_with_id(*comp_id))
+            .find(|comp_id| {
+                !target_node
+                    .archetype()
+                    .contains_component_with_exact_id(*comp_id)
+            })
             .expect("Archetype should contain one extra component ID");
 
         superset_node
@@ -234,7 +250,10 @@ impl ArchetypeNode
         insert_fn: impl FnOnce() -> ArchetypeEdges,
     ) -> &mut ArchetypeEdges
     {
-        debug_assert_eq!(component_id.kind(), UidKind::Component);
+        debug_assert!(matches!(
+            component_id.kind(),
+            UidKind::Component | UidKind::Pair
+        ));
 
         self.edges.entry(component_id).or_insert_with(insert_fn)
     }
@@ -245,13 +264,6 @@ impl ArchetypeNode
         self.edges.iter()
     }
 
-    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