use crate::prelude::{Children, Parent, PreviousParent}; use bevy_ecs::{ bundle::Bundle, entity::Entity, system::{Command, Commands, EntityCommands}, world::{EntityMut, World}, }; use smallvec::SmallVec; #[derive(Debug)] pub struct InsertChildren { parent: Entity, children: SmallVec<[Entity; 8]>, index: usize, } impl Command for InsertChildren { fn write(self: Box, world: &mut World) { for child in self.children.iter() { world .entity_mut(*child) .insert_bundle((Parent(self.parent), PreviousParent(self.parent))); } { let mut added = false; if let Some(mut children) = world.get_mut::(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 { world .entity_mut(self.parent) .insert(Children(self.children)); } } } } #[derive(Debug)] pub struct PushChildren { parent: Entity, children: SmallVec<[Entity; 8]>, } pub struct ChildBuilder<'a, 'b> { commands: &'b mut Commands<'a>, push_children: PushChildren, } impl Command for PushChildren { fn write(self: Box, world: &mut World) { for child in self.children.iter() { world .entity_mut(*child) .insert_bundle((Parent(self.parent), PreviousParent(self.parent))); } { let mut added = false; if let Some(mut children) = world.get_mut::(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 { world .entity_mut(self.parent) .insert(Children(self.children)); } } } } impl<'a, 'b> ChildBuilder<'a, 'b> { pub fn spawn_bundle(&mut self, bundle: impl Bundle) -> EntityCommands<'a, '_> { let e = self.commands.spawn_bundle(bundle); self.push_children.children.push(e.id()); e } pub fn spawn(&mut self) -> EntityCommands<'a, '_> { let e = self.commands.spawn(); self.push_children.children.push(e.id()); e } pub fn parent_entity(&self) -> Entity { self.push_children.parent } pub fn add_command(&mut self, command: C) -> &mut Self { self.commands.add(command); self } } pub trait BuildChildren { fn with_children(&mut self, f: impl FnOnce(&mut ChildBuilder)) -> &mut Self; fn push_children(&mut self, children: &[Entity]) -> &mut Self; fn insert_children(&mut self, index: usize, children: &[Entity]) -> &mut Self; } impl<'a, 'b> BuildChildren for EntityCommands<'a, 'b> { fn with_children(&mut self, spawn_children: impl FnOnce(&mut ChildBuilder)) -> &mut Self { let parent = self.id(); let push_children = { let mut builder = ChildBuilder { commands: self.commands(), push_children: PushChildren { children: SmallVec::default(), parent, }, }; spawn_children(&mut builder); builder.push_children }; self.commands().add(push_children); self } fn push_children(&mut self, children: &[Entity]) -> &mut Self { let parent = self.id(); self.commands().add(PushChildren { children: SmallVec::from(children), parent, }); self } fn insert_children(&mut self, index: usize, children: &[Entity]) -> &mut Self { let parent = self.id(); self.commands().add(InsertChildren { children: SmallVec::from(children), index, parent, }); self } } #[derive(Debug)] pub struct WorldChildBuilder<'w> { world: &'w mut World, current_entity: Option, parent_entities: Vec, } impl<'w> WorldChildBuilder<'w> { pub fn spawn_bundle(&mut self, bundle: impl Bundle + Send + Sync + 'static) -> EntityMut<'_> { let parent_entity = self.parent_entity(); let entity = self .world .spawn() .insert_bundle(bundle) .insert_bundle((Parent(parent_entity), PreviousParent(parent_entity))) .id(); self.current_entity = Some(entity); if let Some(mut parent) = self.world.get_entity_mut(parent_entity) { if let Some(mut children) = parent.get_mut::() { children.0.push(entity); } else { parent.insert(Children(smallvec::smallvec![entity])); } } self.world.entity_mut(entity) } pub fn spawn(&mut self) -> EntityMut<'_> { let parent_entity = self.parent_entity(); let entity = self .world .spawn() .insert_bundle((Parent(parent_entity), PreviousParent(parent_entity))) .id(); self.current_entity = Some(entity); if let Some(mut parent) = self.world.get_entity_mut(parent_entity) { if let Some(mut children) = parent.get_mut::() { children.0.push(entity); } else { parent.insert(Children(smallvec::smallvec![entity])); } } self.world.entity_mut(entity) } pub fn parent_entity(&self) -> Entity { self.parent_entities .last() .cloned() .expect("There should always be a parent at this point.") } } pub trait BuildWorldChildren { fn with_children(&mut self, spawn_children: impl FnOnce(&mut WorldChildBuilder)) -> &mut Self; } impl<'w> BuildWorldChildren for EntityMut<'w> { fn with_children(&mut self, spawn_children: impl FnOnce(&mut WorldChildBuilder)) -> &mut Self { { let entity = self.id(); let mut builder = WorldChildBuilder { current_entity: None, parent_entities: vec![entity], // SAFE: self.update_location() is called below. It is impossible to make EntityMut // function calls on `self` within the scope defined here world: unsafe { self.world_mut() }, }; spawn_children(&mut builder); } self.update_location(); self } } impl<'w> BuildWorldChildren for WorldChildBuilder<'w> { fn with_children( &mut self, spawn_children: impl FnOnce(&mut WorldChildBuilder<'w>), ) -> &mut Self { let current_entity = self .current_entity .expect("Cannot add children without a parent. Try creating an entity first."); self.parent_entities.push(current_entity); self.current_entity = None; spawn_children(self); self.current_entity = self.parent_entities.pop(); self } } #[cfg(test)] mod tests { use super::BuildChildren; use crate::prelude::{Children, Parent, PreviousParent}; use bevy_ecs::{ entity::Entity, system::{CommandQueue, Commands}, world::World, }; use smallvec::{smallvec, SmallVec}; #[test] fn build_children() { let mut world = World::default(); let mut queue = CommandQueue::default(); let mut commands = Commands::new(&mut queue, &world); let mut children = Vec::new(); let parent = commands.spawn().insert(1).id(); commands.entity(parent).with_children(|parent| { children.push(parent.spawn().insert(2).id()); children.push(parent.spawn().insert(3).id()); children.push(parent.spawn().insert(4).id()); }); queue.apply(&mut world); assert_eq!( world.get::(parent).unwrap().0.as_slice(), children.as_slice(), ); assert_eq!(*world.get::(children[0]).unwrap(), Parent(parent)); assert_eq!(*world.get::(children[1]).unwrap(), Parent(parent)); assert_eq!( *world.get::(children[0]).unwrap(), PreviousParent(parent) ); assert_eq!( *world.get::(children[1]).unwrap(), PreviousParent(parent) ); } #[test] fn push_and_insert_children() { let mut world = World::default(); let entities = world .spawn_batch(vec![(1,), (2,), (3,), (4,), (5,)]) .collect::>(); let mut queue = CommandQueue::default(); { let mut commands = Commands::new(&mut queue, &world); commands.entity(entities[0]).push_children(&entities[1..3]); } queue.apply(&mut world); let parent = entities[0]; let child1 = entities[1]; let child2 = entities[2]; let child3 = entities[3]; let child4 = entities[4]; let expected_children: SmallVec<[Entity; 8]> = smallvec![child1, child2]; assert_eq!( world.get::(parent).unwrap().0.clone(), expected_children ); assert_eq!(*world.get::(child1).unwrap(), Parent(parent)); assert_eq!(*world.get::(child2).unwrap(), Parent(parent)); assert_eq!( *world.get::(child1).unwrap(), PreviousParent(parent) ); assert_eq!( *world.get::(child2).unwrap(), PreviousParent(parent) ); { let mut commands = Commands::new(&mut queue, &world); commands.entity(parent).insert_children(1, &entities[3..]); } queue.apply(&mut world); let expected_children: SmallVec<[Entity; 8]> = smallvec![child1, child3, child4, child2]; assert_eq!( world.get::(parent).unwrap().0.clone(), expected_children ); assert_eq!(*world.get::(child3).unwrap(), Parent(parent)); assert_eq!(*world.get::(child4).unwrap(), Parent(parent)); assert_eq!( *world.get::(child3).unwrap(), PreviousParent(parent) ); assert_eq!( *world.get::(child4).unwrap(), PreviousParent(parent) ); } }