Enforce uniqueness in Children

crates/bevy_transform/src/components/children.rs

@@ -1,18 +1,22 @@
 use bevy_ecs::{Entity, MapEntities};
 use bevy_reflect::{Reflect, ReflectComponent, ReflectMapEntities};
+use bevy_utils::HashMap;
 use smallvec::SmallVec;
 use std::ops::Deref;
 
-#[derive(Default, Clone, Debug, Reflect)]
+#[derive(Default, Clone, Debug, Reflect, PartialEq, Eq)]
 #[reflect(Component, MapEntities)]
-pub struct Children(pub(crate) SmallVec<[Entity; 8]>);
+pub struct Children {
+    order: SmallVec<[Entity; 8]>,
+    uniqueness: HashMap<Entity, usize>,
+}
 
 impl MapEntities for Children {
     fn map_entities(
         &mut self,
         entity_map: &bevy_ecs::EntityMap,
     ) -> Result<(), bevy_ecs::MapEntitiesError> {
-        for entity in self.0.iter_mut() {
+        for entity in self.order.iter_mut() {
             *entity = entity_map.get(*entity)?;
         }
 
@@ -21,20 +25,91 @@ impl MapEntities for Children {
 }
 
 impl Children {
-    pub fn with(entity: &[Entity]) -> Self {
-        Self(SmallVec::from_slice(entity))
+    pub fn with(entities: &[Entity]) -> Self {
+        let mut children = Self::default();
+        for entity in entities {
+            children.push(*entity);
+        }
+        children
     }
 
     /// Swaps the child at `a_index` with the child at `b_index`
     pub fn swap(&mut self, a_index: usize, b_index: usize) {
-        self.0.swap(a_index, b_index);
+        let a_entity = self.order[a_index];
+        let b_entity = self.order[b_index];
+        self.order.swap(a_index, b_index);
+        self.uniqueness.insert(a_entity, b_index);
+        self.uniqueness.insert(b_entity, a_index);
+    }
+
+    pub(crate) fn push(&mut self, entity: Entity) {
+        let order = &mut self.order;
+        let uniqueness = &mut self.uniqueness;
+
+        uniqueness.entry(entity).or_insert_with(|| {
+            order.push(entity);
+            order.len() - 1
+        });
+
+        let desired_index = order.len() - 1;
+        let current_index = uniqueness[&entity];
+        if current_index != desired_index {
+            self.swap(current_index, desired_index);
+        }
+    }
+
+    pub(crate) fn retain<F: FnMut(&Entity) -> bool>(&mut self, mut f: F) {
+        let order = &mut self.order;
+        let uniqueness = &mut self.uniqueness;
+
+        let mut offset = 0;
+        order.retain(|e| {
+            if f(e) {
+                *uniqueness.get_mut(e).unwrap() -= offset;
+                true
+            } else {
+                offset += 1;
+                uniqueness.remove(e);
+                false
+            }
+        });
+    }
+
+    pub fn contains(&self, entity: &Entity) -> bool {
+        self.uniqueness.contains_key(entity)
+    }
+
+    pub(crate) fn extend<T: IntoIterator<Item = Entity>>(&mut self, iter: T) {
+        for entity in iter {
+            self.push(entity);
+        }
+    }
+
+    pub(crate) fn insert(&mut self, index: usize, entities: &[Entity]) {
+        let initial_count = self.order.len();
+        self.extend(entities.iter().cloned());
+        let actually_inserted = self.order.len() - initial_count;
+        let mut desired_index =
+            (index as i32 - (entities.len() as i32 - actually_inserted as i32)).max(0) as usize;
+        for entity in entities {
+            let current_index = self.uniqueness[entity];
+            if current_index != desired_index {
+                self.swap(current_index, desired_index);
+            }
+            desired_index += 1;
+        }
+    }
+
+    pub(crate) fn take(&mut self) -> SmallVec<[Entity; 8]> {
+        self.uniqueness.clear();
+        std::mem::take(&mut self.order)
     }
 }
 
 impl Deref for Children {
     type Target = [Entity];
 
     fn deref(&self) -> &Self::Target {
-        &self.0[..]
+        &self.order[..]
     }
 }

crates/bevy_transform/src/hierarchy/child_builder.rs

@@ -17,16 +17,11 @@ impl Command for InsertChildren {
                 .unwrap();
         }
         {
-            let mut added = false;
             if let Ok(mut children) = world.get_mut::<Children>(self.parent) {
-                children.0.insert_from_slice(self.index, &self.children);
-                added = true;
-            }
-
-            // NOTE: ideally this is just an else statement, but currently that _incorrectly_ fails borrow-checking
-            if !added {
+                children.insert(self.index, &self.children);
+            } else {
                 world
-                    .insert_one(self.parent, Children(self.children))
+                    .insert_one(self.parent, Children::with(&self.children))
                     .unwrap();
             }
         }
@@ -52,16 +47,11 @@ impl Command for PushChildren {
                 .unwrap();
         }
         {
-            let mut added = false;
             if let Ok(mut children) = world.get_mut::<Children>(self.parent) {
-                children.0.extend(self.children.iter().cloned());
-                added = true;
-            }
-
-            // NOTE: ideally this is just an else statement, but currently that _incorrectly_ fails borrow-checking
-            if !added {
+                children.extend(self.children.iter().cloned());
+            } else {
                 world
-                    .insert_one(self.parent, Children(self.children))
+                    .insert_one(self.parent, Children::with(&self.children))
                     .unwrap();
             }
         }
@@ -203,7 +193,6 @@ mod tests {
     use super::BuildChildren;
     use crate::prelude::{Children, Parent, PreviousParent};
     use bevy_ecs::{Commands, Entity, Resources, World};
-    use smallvec::{smallvec, SmallVec};
 
     #[test]
     fn build_children() {
@@ -236,10 +225,15 @@ mod tests {
         let child1 = child1.expect("child1 should exist");
         let child2 = child2.expect("child2 should exist");
         let child3 = child3.expect("child3 should exist");
-        let expected_children: SmallVec<[Entity; 8]> = smallvec![child1, child2, child3];
+        let expected_children = vec![child1, child2, child3];
 
         assert_eq!(
-            world.get::<Children>(parent).unwrap().0.clone(),
+            world
+                .get::<Children>(parent)
+                .unwrap()
+                .iter()
+                .copied()
+                .collect::<Vec<_>>(),
             expected_children
         );
         assert_eq!(*world.get::<Parent>(child1).unwrap(), Parent(parent));
@@ -255,58 +249,117 @@ mod tests {
         );
     }
 
-    #[test]
-    fn push_and_insert_children() {
+    fn setup() -> (World, Resources, Commands, Vec<Entity>, Entity) {
         let mut world = World::default();
-        let mut resources = Resources::default();
-        let mut commands = Commands::default();
+        let resources = Resources::default();
+        let commands = Commands::default();
         let entities = world
-            .spawn_batch(vec![(1,), (2,), (3,), (4,), (5,)])
+            .spawn_batch(vec![(0,), (1,), (2,), (3,), (4,)])
             .collect::<Vec<Entity>>();
+        let parent = entities[0];
+        (world, resources, commands, entities, parent)
+    }
 
-        commands.push_children(entities[0], &entities[1..3]);
+    #[test]
+    fn push_children_adds_parent_component() {
+        let (mut world, mut resources, mut commands, child, parent) = setup();
+        commands.push_children(parent, &child[1..=2]);
         commands.apply(&mut world, &mut resources);
+        assert_eq!(world.get::<Parent>(child[2]).unwrap(), &Parent(parent));
+    }
 
-        let parent = entities[0];
-        let child1 = entities[1];
-        let child2 = entities[2];
-        let child3 = entities[3];
-        let child4 = entities[4];
+    #[test]
+    fn push_children_adds_previous_parent_component() {
+        let (mut world, mut resources, mut commands, child, parent) = setup();
+        commands.push_children(parent, &child[1..=2]);
+        commands.apply(&mut world, &mut resources);
+        assert_eq!(
+            world.get::<PreviousParent>(child[2]).unwrap(),
+            &PreviousParent(parent)
+        );
+    }
 
-        let expected_children: SmallVec<[Entity; 8]> = smallvec![child1, child2];
+    #[test]
+    fn push_children_adds_children_component() {
+        let (mut world, mut resources, mut commands, child, parent) = setup();
+        commands.push_children(parent, &child[1..=2]);
+        commands.apply(&mut world, &mut resources);
         assert_eq!(
-            world.get::<Children>(parent).unwrap().0.clone(),
-            expected_children
+            world.get::<Children>(parent).unwrap(),
+            &Children::with(&child[1..=2])
         );
-        assert_eq!(*world.get::<Parent>(child1).unwrap(), Parent(parent));
-        assert_eq!(*world.get::<Parent>(child2).unwrap(), Parent(parent));
+    }
 
+    #[test]
+    fn push_children_keeps_children_unique() {
+        let (mut world, mut resources, mut commands, child, parent) = setup();
+        commands.push_children(parent, &child[1..=2]);
+        commands.apply(&mut world, &mut resources);
+        commands.push_children(parent, &child[1..=2]);
+        commands.apply(&mut world, &mut resources);
         assert_eq!(
-            *world.get::<PreviousParent>(child1).unwrap(),
-            PreviousParent(parent)
+            world.get::<Children>(parent).unwrap(),
+            &Children::with(&child[1..=2])
         );
+    }
+
+    #[test]
+    fn insert_children_adds_parent_component() {
+        let (mut world, mut resources, mut commands, child, parent) = setup();
+        commands.insert_children(parent, 0, &child[1..=2]);
+        commands.apply(&mut world, &mut resources);
+        assert_eq!(world.get::<Parent>(child[2]).unwrap(), &Parent(parent));
+    }
+
+    #[test]
+    fn insert_children_adds_previous_parent_component() {
+        let (mut world, mut resources, mut commands, child, parent) = setup();
+        commands.insert_children(parent, 0, &child[1..=2]);
+        commands.apply(&mut world, &mut resources);
         assert_eq!(
-            *world.get::<PreviousParent>(child2).unwrap(),
-            PreviousParent(parent)
+            world.get::<PreviousParent>(child[2]).unwrap(),
+            &PreviousParent(parent)
         );
+    }
 
-        commands.insert_children(parent, 1, &entities[3..]);
+    #[test]
+    fn insert_children_adds_children_component() {
+        let (mut world, mut resources, mut commands, child, parent) = setup();
+        commands.insert_children(parent, 0, &child[1..=2]);
         commands.apply(&mut world, &mut resources);
+        assert_eq!(
+            world.get::<Children>(parent).unwrap(),
+            &Children::with(&child[1..=2])
+        );
+    }
 
-        let expected_children: SmallVec<[Entity; 8]> = smallvec![child1, child3, child4, child2];
+    #[test]
+    fn insert_children_keeps_children_unique() {
+        let (mut world, mut resources, mut commands, child, parent) = setup();
+        commands.insert_children(parent, 0, &child[1..=2]);
+        commands.apply(&mut world, &mut resources);
+        commands.insert_children(parent, 1, &child[1..=2]);
+        commands.apply(&mut world, &mut resources);
         assert_eq!(
-            world.get::<Children>(parent).unwrap().0.clone(),
-            expected_children
+            world.get::<Children>(parent).unwrap(),
+            &Children::with(&child[1..=2])
         );
-        assert_eq!(*world.get::<Parent>(child3).unwrap(), Parent(parent));
-        assert_eq!(*world.get::<Parent>(child4).unwrap(), Parent(parent));
+    }
+
+    #[test]
+    fn insert_children_keeps_children_order() {
+        let (mut world, mut resources, mut commands, child, parent) = setup();
+        commands.insert_children(parent, 0, &child[1..=2]);
+        commands.apply(&mut world, &mut resources);
         assert_eq!(
-            *world.get::<PreviousParent>(child3).unwrap(),
-            PreviousParent(parent)
+            world.get::<Children>(parent).unwrap(),
+            &Children::with(&[child[1], child[2]])
         );
+        commands.insert_children(parent, 1, &child[3..=4]);
+        commands.apply(&mut world, &mut resources);
         assert_eq!(
-            *world.get::<PreviousParent>(child4).unwrap(),
-            PreviousParent(parent)
+            world.get::<Children>(parent).unwrap(),
+            &Children::with(&[child[1], child[3], child[4], child[2]])
         );
     }
 }

crates/bevy_transform/src/hierarchy/hierarchy.rs

@@ -11,7 +11,7 @@ pub fn despawn_with_children_recursive(world: &mut World, entity: Entity) {
     // first, make the entity's own parent forget about it
     if let Ok(parent) = world.get::<Parent>(entity).map(|parent| parent.0) {
         if let Ok(mut children) = world.get_mut::<Children>(parent) {
-            children.0.retain(|c| *c != entity);
+            children.retain(|c| *c != entity);
         }
     }
 
@@ -22,7 +22,7 @@ pub fn despawn_with_children_recursive(world: &mut World, entity: Entity) {
 // Should only be called by `despawn_with_children_recursive`!
 fn despawn_with_children_recursive_inner(world: &mut World, entity: Entity) {
     if let Ok(mut children) = world.get_mut::<Children>(entity) {
-        for e in std::mem::take(&mut children.0) {
+        for e in children.take() {
             despawn_with_children_recursive(world, e);
         }
     }

crates/bevy_transform/src/hierarchy/hierarchy_maintenance_system.rs

@@ -15,7 +15,7 @@ pub fn parent_update_system(
     // them from the `Children` of the `PreviousParent`.
     for (entity, previous_parent) in removed_parent_query.iter() {
         if let Ok(mut previous_parent_children) = children_query.get_mut(previous_parent.0) {
-            previous_parent_children.0.retain(|e| *e != entity);
+            previous_parent_children.retain(|e| *e != entity);
             commands.remove_one::<PreviousParent>(entity);
         }
     }
@@ -33,7 +33,7 @@ pub fn parent_update_system(
 
             // Remove from `PreviousParent.Children`.
             if let Ok(mut previous_parent_children) = children_query.get_mut(previous_parent.0) {
-                (*previous_parent_children).0.retain(|e| *e != entity);
+                (*previous_parent_children).retain(|e| *e != entity);
             }
 
             // Set `PreviousParent = Parent`.
@@ -47,10 +47,10 @@ pub fn parent_update_system(
         if let Ok(mut new_parent_children) = children_query.get_mut(parent.0) {
             // This is the parent
             debug_assert!(
-                !(*new_parent_children).0.contains(&entity),
+                !(*new_parent_children).contains(&entity),
                 "children already added"
             );
-            (*new_parent_children).0.push(entity);
+            (*new_parent_children).push(entity);
         } else {
             // The parent doesn't have a children entity, lets add it
             children_additions
@@ -109,7 +109,6 @@ mod test {
             world
                 .get::<Children>(parent)
                 .unwrap()
-                .0
                 .iter()
                 .cloned()
                 .collect::<Vec<_>>(),