Skip to content

Commit

Permalink
RemoveChildren command
Browse files Browse the repository at this point in the history
  • Loading branch information
jihiggins authored and cart committed Dec 22, 2021
1 parent 6c47964 commit 91a8c60
Showing 1 changed file with 130 additions and 7 deletions.
137 changes: 130 additions & 7 deletions crates/bevy_transform/src/hierarchy/child_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -40,11 +40,6 @@ pub struct PushChildren {
children: SmallVec<[Entity; 8]>,
}

pub struct ChildBuilder<'w, 's, 'a> {
commands: &'a mut Commands<'w, 's>,
push_children: PushChildren,
}

impl Command for PushChildren {
fn write(self, world: &mut World) {
for child in self.children.iter() {
Expand All @@ -71,6 +66,46 @@ impl Command for PushChildren {
}
}

pub struct RemoveChildren {
parent: Entity,
children: SmallVec<[Entity; 8]>,
}

fn remove_children(parent: Entity, children: &[Entity], world: &mut World) {
for child in children.iter() {
let mut child = world.entity_mut(*child);
let mut remove_parent = false;
if let Some(child_parent) = child.get_mut::<Parent>() {
if child_parent.0 == parent {
remove_parent = true;
}
}
if remove_parent {
if let Some(parent) = child.remove::<Parent>() {
child.insert(PreviousParent(parent.0));
}
}
}
// Remove the children from the parents.
if let Some(mut parent_children) = world.get_mut::<Children>(parent) {
parent_children
.0
.retain(|parent_child| !children.contains(parent_child));
}
}

impl Command for RemoveChildren {
fn write(self, world: &mut World) {
// Remove any matching Parent components from the children
remove_children(self.parent, &self.children, world);
}
}

pub struct ChildBuilder<'w, 's, 'a> {
commands: &'a mut Commands<'w, 's>,
push_children: PushChildren,
}

impl<'w, 's, 'a> ChildBuilder<'w, 's, 'a> {
pub fn spawn_bundle(&mut self, bundle: impl Bundle) -> EntityCommands<'w, 's, '_> {
let e = self.commands.spawn_bundle(bundle);
Expand Down Expand Up @@ -98,6 +133,7 @@ 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;
fn remove_children(&mut self, children: &[Entity]) -> &mut Self;
}

impl<'w, 's, 'a> BuildChildren for EntityCommands<'w, 's, 'a> {
Expand Down Expand Up @@ -137,6 +173,15 @@ impl<'w, 's, 'a> BuildChildren for EntityCommands<'w, 's, 'a> {
});
self
}

fn remove_children(&mut self, children: &[Entity]) -> &mut Self {
let parent = self.id();
self.commands().add(RemoveChildren {
children: SmallVec::from(children),
parent,
});
self
}
}

#[derive(Debug)]
Expand Down Expand Up @@ -196,6 +241,7 @@ pub trait BuildWorldChildren {
fn with_children(&mut self, spawn_children: impl FnOnce(&mut WorldChildBuilder)) -> &mut Self;
fn push_children(&mut self, children: &[Entity]) -> &mut Self;
fn insert_children(&mut self, index: usize, children: &[Entity]) -> &mut Self;
fn remove_children(&mut self, children: &[Entity]) -> &mut Self;
}

impl<'w> BuildWorldChildren for EntityMut<'w> {
Expand Down Expand Up @@ -260,6 +306,33 @@ impl<'w> BuildWorldChildren for EntityMut<'w> {
}
self
}

fn remove_children(&mut self, children: &[Entity]) -> &mut Self {
let parent = self.id();
// SAFE: This doesn't change the parent's location
let world = unsafe { self.world_mut() };
for child in children.iter() {
let mut child = world.entity_mut(*child);
let mut remove_parent = false;
if let Some(child_parent) = child.get_mut::<Parent>() {
if child_parent.0 == parent {
remove_parent = true;
}
}
if remove_parent {
if let Some(parent) = child.remove::<Parent>() {
child.insert(PreviousParent(parent.0));
}
}
}
// Remove the children from the parents.
if let Some(mut parent_children) = world.get_mut::<Children>(parent) {
parent_children
.0
.retain(|parent_child| !children.contains(parent_child));
}
self
}
}

impl<'w> BuildWorldChildren for WorldChildBuilder<'w> {
Expand Down Expand Up @@ -319,6 +392,15 @@ impl<'w> BuildWorldChildren for WorldChildBuilder<'w> {
}
self
}

fn remove_children(&mut self, children: &[Entity]) -> &mut Self {
let parent = self
.current_entity
.expect("Cannot remove children without a parent. Try creating an entity first.");

remove_children(parent, children, self.world);
self
}
}

#[cfg(test)]
Expand Down Expand Up @@ -369,7 +451,7 @@ mod tests {
}

#[test]
fn push_and_insert_children_commands() {
fn push_and_insert_and_remove_children_commands() {
let mut world = World::default();

let entities = world
Expand Down Expand Up @@ -427,10 +509,33 @@ mod tests {
*world.get::<PreviousParent>(child4).unwrap(),
PreviousParent(parent)
);

let remove_children = [child1, child4];
{
let mut commands = Commands::new(&mut queue, &world);
commands.entity(parent).remove_children(&remove_children);
}
queue.apply(&mut world);

let expected_children: SmallVec<[Entity; 8]> = smallvec![child3, child2];
assert_eq!(
world.get::<Children>(parent).unwrap().0.clone(),
expected_children
);
assert!(world.get::<Parent>(child1).is_none());
assert!(world.get::<Parent>(child4).is_none());
assert_eq!(
*world.get::<PreviousParent>(child1).unwrap(),
PreviousParent(parent)
);
assert_eq!(
*world.get::<PreviousParent>(child4).unwrap(),
PreviousParent(parent)
);
}

#[test]
fn push_and_insert_children_world() {
fn push_and_insert_and_remove_children_world() {
let mut world = World::default();

let entities = world
Expand Down Expand Up @@ -478,6 +583,24 @@ mod tests {
*world.get::<PreviousParent>(child4).unwrap(),
PreviousParent(parent)
);

let remove_children = [child1, child4];
world.entity_mut(parent).remove_children(&remove_children);
let expected_children: SmallVec<[Entity; 8]> = smallvec![child3, child2];
assert_eq!(
world.get::<Children>(parent).unwrap().0.clone(),
expected_children
);
assert!(world.get::<Parent>(child1).is_none());
assert!(world.get::<Parent>(child4).is_none());
assert_eq!(
*world.get::<PreviousParent>(child1).unwrap(),
PreviousParent(parent)
);
assert_eq!(
*world.get::<PreviousParent>(child4).unwrap(),
PreviousParent(parent)
);
}

#[test]
Expand Down

0 comments on commit 91a8c60

Please sign in to comment.