1use crate::{
2 bundle::Bundle,
3 entity::{Entity, EntityIndexSet},
4 prelude::Children,
5 relationship::{
6 Relationship, RelationshipHookMode, RelationshipSourceCollection, RelationshipTarget,
7 },
8 system::{Commands, EntityCommands},
9 world::{DeferredWorld, EntityWorldMut, World},
10};
11use bevy_platform::prelude::{Box, Vec};
12use core::{marker::PhantomData, mem};
13
14use super::OrderedRelationshipSourceCollection;
15
16impl<'w> EntityWorldMut<'w> {
17 pub fn with_related<R: Relationship>(&mut self, bundle: impl Bundle) -> &mut Self {
19 let parent = self.id();
20 self.world_scope(|world| {
21 world.spawn((bundle, R::from(parent)));
22 });
23 self
24 }
25
26 pub fn with_related_entities<R: Relationship>(
28 &mut self,
29 func: impl FnOnce(&mut RelatedSpawner<R>),
30 ) -> &mut Self {
31 let parent = self.id();
32 self.world_scope(|world| {
33 func(&mut RelatedSpawner::new(world, parent));
34 });
35 self
36 }
37
38 pub fn add_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
42 let id = self.id();
43 self.world_scope(|world| {
44 for related in related {
45 world
46 .entity_mut(*related)
47 .modify_or_insert_relation_with_relationship_hook_mode::<R>(
48 id,
49 RelationshipHookMode::Run,
50 );
51 }
52 });
53 self
54 }
55
56 pub fn detach_all_related<R: Relationship>(&mut self) -> &mut Self {
58 self.remove::<R::RelationshipTarget>()
59 }
60
61 pub fn insert_related<R: Relationship>(&mut self, index: usize, related: &[Entity]) -> &mut Self
88 where
89 <R::RelationshipTarget as RelationshipTarget>::Collection:
90 OrderedRelationshipSourceCollection,
91 {
92 let id = self.id();
93 self.world_scope(|world| {
94 for (offset, related) in related.iter().enumerate() {
95 let index = index.saturating_add(offset);
96 if world
97 .get::<R>(*related)
98 .is_some_and(|relationship| relationship.get() == id)
99 {
100 world
101 .get_mut::<R::RelationshipTarget>(id)
102 .expect("hooks should have added relationship target")
103 .collection_mut_risky()
104 .place(*related, index);
105 } else {
106 world
107 .entity_mut(*related)
108 .modify_or_insert_relation_with_relationship_hook_mode::<R>(
109 id,
110 RelationshipHookMode::Run,
111 );
112 world
113 .get_mut::<R::RelationshipTarget>(id)
114 .expect("hooks should have added relationship target")
115 .collection_mut_risky()
116 .place_most_recent(index);
117 }
118 }
119 });
120
121 self
122 }
123
124 pub fn remove_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
126 let id = self.id();
127 self.world_scope(|world| {
128 for related in related {
129 if world
130 .get::<R>(*related)
131 .is_some_and(|relationship| relationship.get() == id)
132 {
133 world.entity_mut(*related).remove::<R>();
134 }
135 }
136 });
137
138 self
139 }
140
141 pub fn replace_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
145 type Collection<R> =
146 <<R as Relationship>::RelationshipTarget as RelationshipTarget>::Collection;
147
148 if related.is_empty() {
149 self.remove::<R::RelationshipTarget>();
150
151 return self;
152 }
153
154 let Some(existing_relations) = self.get_mut::<R::RelationshipTarget>() else {
155 return self.add_related::<R>(related);
156 };
157
158 let mut relations = mem::replace(
161 existing_relations.into_inner(),
162 <R as Relationship>::RelationshipTarget::from_collection_risky(
163 Collection::<R>::with_capacity(0),
164 ),
165 );
166
167 let collection = relations.collection_mut_risky();
168
169 let existing_relations = EntityIndexSet::from_iter(collection.iter());
170 let final_relations = EntityIndexSet::from_iter(related.iter().copied());
171
172 let id = self.id();
173 self.world_scope(|world| {
174 for &related in existing_relations.difference(&final_relations) {
176 world.entity_mut(related).remove::<R>();
177 }
178
179 for &related in final_relations.difference(&existing_relations) {
181 world
183 .entity_mut(related)
184 .modify_or_insert_relation_with_relationship_hook_mode::<R>(
185 id,
186 RelationshipHookMode::Skip,
187 );
188 }
189 });
190
191 collection.clear();
193 collection.extend_from_iter(final_relations);
194 self.insert(relations);
195
196 self
197 }
198
199 pub fn replace_related_with_difference<R: Relationship>(
220 &mut self,
221 entities_to_unrelate: &[Entity],
222 entities_to_relate: &[Entity],
223 newly_related_entities: &[Entity],
224 ) -> &mut Self {
225 #[cfg(debug_assertions)]
226 {
227 use crate::entity::hash_set::EntityHashSet;
228 let entities_to_relate = EntityHashSet::from_iter(entities_to_relate.iter().copied());
229 let entities_to_unrelate =
230 EntityHashSet::from_iter(entities_to_unrelate.iter().copied());
231 let mut newly_related_entities =
232 EntityHashSet::from_iter(newly_related_entities.iter().copied());
233 assert!(
234 entities_to_relate.is_disjoint(&entities_to_unrelate),
235 "`entities_to_relate` ({entities_to_relate:?}) shared entities with `entities_to_unrelate` ({entities_to_unrelate:?})"
236 );
237 assert!(
238 newly_related_entities.is_disjoint(&entities_to_unrelate),
239 "`newly_related_entities` ({newly_related_entities:?}) shared entities with `entities_to_unrelate ({entities_to_unrelate:?})`"
240 );
241 assert!(
242 newly_related_entities.is_subset(&entities_to_relate),
243 "`newly_related_entities` ({newly_related_entities:?}) wasn't a subset of `entities_to_relate` ({entities_to_relate:?})"
244 );
245
246 if let Some(target) = self.get::<R::RelationshipTarget>() {
247 let existing_relationships: EntityHashSet = target.collection().iter().collect();
248
249 assert!(
250 existing_relationships.is_disjoint(&newly_related_entities),
251 "`newly_related_entities` contains an entity that wouldn't be newly related"
252 );
253
254 newly_related_entities.extend(existing_relationships);
255 newly_related_entities -= &entities_to_unrelate;
256 }
257
258 assert_eq!(newly_related_entities, entities_to_relate, "`entities_to_relate` ({entities_to_relate:?}) didn't contain all entities that would end up related");
259 };
260
261 match self.get_mut::<R::RelationshipTarget>() {
262 None => {
263 self.add_related::<R>(entities_to_relate);
264
265 return self;
266 }
267 Some(mut target) => {
268 let collection = target.collection_mut_risky();
270 collection.clear();
271
272 collection.extend_from_iter(entities_to_relate.iter().copied());
273 }
274 }
275
276 let this = self.id();
277 self.world_scope(|world| {
278 for unrelate in entities_to_unrelate {
279 world.entity_mut(*unrelate).remove::<R>();
280 }
281
282 for new_relation in newly_related_entities {
283 world
285 .entity_mut(*new_relation)
286 .modify_or_insert_relation_with_relationship_hook_mode::<R>(
287 this,
288 RelationshipHookMode::Skip,
289 );
290 }
291 });
292
293 self
294 }
295
296 pub fn add_one_related<R: Relationship>(&mut self, entity: Entity) -> &mut Self {
300 self.add_related::<R>(&[entity])
301 }
302
303 pub fn despawn_related<S: RelationshipTarget>(&mut self) -> &mut Self {
306 if let Some(sources) = self.get::<S>() {
307 let sources = sources.iter().collect::<Vec<_>>();
310 self.world_scope(|world| {
311 for entity in sources {
312 if let Ok(entity_mut) = world.get_entity_mut(entity) {
313 entity_mut.despawn();
314 };
315 }
316 });
317 }
318 self
319 }
320
321 pub fn despawn_children(&mut self) -> &mut Self {
326 self.despawn_related::<Children>();
327 self
328 }
329
330 pub fn insert_recursive<S: RelationshipTarget>(
340 &mut self,
341 bundle: impl Bundle + Clone,
342 ) -> &mut Self {
343 self.insert(bundle.clone());
344 if let Some(relationship_target) = self.get::<S>() {
345 let related_vec: Vec<Entity> = relationship_target.iter().collect();
346 for related in related_vec {
347 self.world_scope(|world| {
348 world
349 .entity_mut(related)
350 .insert_recursive::<S>(bundle.clone());
351 });
352 }
353 }
354
355 self
356 }
357
358 pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
366 self.remove::<B>();
367 if let Some(relationship_target) = self.get::<S>() {
368 let related_vec: Vec<Entity> = relationship_target.iter().collect();
369 for related in related_vec {
370 self.world_scope(|world| {
371 world.entity_mut(related).remove_recursive::<S, B>();
372 });
373 }
374 }
375
376 self
377 }
378
379 fn modify_or_insert_relation_with_relationship_hook_mode<R: Relationship>(
380 &mut self,
381 entity: Entity,
382 relationship_hook_mode: RelationshipHookMode,
383 ) {
384 if size_of::<R>() > size_of::<Entity>() {
386 self.assert_not_despawned();
387
388 let this = self.id();
389
390 let modified = self.world_scope(|world| {
391 let modified = DeferredWorld::from(&mut *world)
392 .modify_component_with_relationship_hook_mode::<R, _>(
393 this,
394 relationship_hook_mode,
395 |r| r.set_risky(entity),
396 )
397 .expect("entity access must be valid")
398 .is_some();
399
400 world.flush();
401
402 modified
403 });
404
405 if modified {
406 return;
407 }
408 }
409
410 self.insert_with_relationship_hook_mode(R::from(entity), relationship_hook_mode);
411 }
412}
413
414impl<'a> EntityCommands<'a> {
415 pub fn with_related<R: Relationship>(&mut self, bundle: impl Bundle) -> &mut Self {
417 let parent = self.id();
418 self.commands.spawn((bundle, R::from(parent)));
419 self
420 }
421
422 pub fn with_related_entities<R: Relationship>(
424 &mut self,
425 func: impl FnOnce(&mut RelatedSpawnerCommands<R>),
426 ) -> &mut Self {
427 let id = self.id();
428 func(&mut RelatedSpawnerCommands::new(self.commands(), id));
429 self
430 }
431
432 pub fn add_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
436 let related: Box<[Entity]> = related.into();
437
438 self.queue(move |mut entity: EntityWorldMut| {
439 entity.add_related::<R>(&related);
440 })
441 }
442
443 pub fn detach_all_related<R: Relationship>(&mut self) -> &mut Self {
445 self.queue(|mut entity: EntityWorldMut| {
446 entity.detach_all_related::<R>();
447 })
448 }
449
450 pub fn insert_related<R: Relationship>(&mut self, index: usize, related: &[Entity]) -> &mut Self
456 where
457 <R::RelationshipTarget as RelationshipTarget>::Collection:
458 OrderedRelationshipSourceCollection,
459 {
460 let related: Box<[Entity]> = related.into();
461
462 self.queue(move |mut entity: EntityWorldMut| {
463 entity.insert_related::<R>(index, &related);
464 })
465 }
466
467 pub fn add_one_related<R: Relationship>(&mut self, entity: Entity) -> &mut Self {
471 self.add_related::<R>(&[entity])
472 }
473
474 pub fn remove_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
476 let related: Box<[Entity]> = related.into();
477
478 self.queue(move |mut entity: EntityWorldMut| {
479 entity.remove_related::<R>(&related);
480 })
481 }
482
483 pub fn replace_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
487 let related: Box<[Entity]> = related.into();
488
489 self.queue(move |mut entity: EntityWorldMut| {
490 entity.replace_related::<R>(&related);
491 })
492 }
493
494 pub fn replace_related_with_difference<R: Relationship>(
505 &mut self,
506 entities_to_unrelate: &[Entity],
507 entities_to_relate: &[Entity],
508 newly_related_entities: &[Entity],
509 ) -> &mut Self {
510 let entities_to_unrelate: Box<[Entity]> = entities_to_unrelate.into();
511 let entities_to_relate: Box<[Entity]> = entities_to_relate.into();
512 let newly_related_entities: Box<[Entity]> = newly_related_entities.into();
513
514 self.queue(move |mut entity: EntityWorldMut| {
515 entity.replace_related_with_difference::<R>(
516 &entities_to_unrelate,
517 &entities_to_relate,
518 &newly_related_entities,
519 );
520 })
521 }
522
523 pub fn despawn_related<S: RelationshipTarget>(&mut self) -> &mut Self {
526 self.queue(move |mut entity: EntityWorldMut| {
527 entity.despawn_related::<S>();
528 })
529 }
530
531 pub fn despawn_children(&mut self) -> &mut Self {
536 self.despawn_related::<Children>()
537 }
538
539 pub fn insert_recursive<S: RelationshipTarget>(
547 &mut self,
548 bundle: impl Bundle + Clone,
549 ) -> &mut Self {
550 self.queue(move |mut entity: EntityWorldMut| {
551 entity.insert_recursive::<S>(bundle);
552 })
553 }
554
555 pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
563 self.queue(move |mut entity: EntityWorldMut| {
564 entity.remove_recursive::<S, B>();
565 })
566 }
567}
568
569pub struct RelatedSpawner<'w, R: Relationship> {
572 target: Entity,
573 world: &'w mut World,
574 _marker: PhantomData<R>,
575}
576
577impl<'w, R: Relationship> RelatedSpawner<'w, R> {
578 pub fn new(world: &'w mut World, target: Entity) -> Self {
580 Self {
581 world,
582 target,
583 _marker: PhantomData,
584 }
585 }
586
587 pub fn spawn(&mut self, bundle: impl Bundle) -> EntityWorldMut<'_> {
590 self.world.spawn((R::from(self.target), bundle))
591 }
592
593 pub fn spawn_empty(&mut self) -> EntityWorldMut<'_> {
596 self.world.spawn(R::from(self.target))
597 }
598
599 pub fn target_entity(&self) -> Entity {
601 self.target
602 }
603
604 pub fn world(&self) -> &World {
606 self.world
607 }
608
609 pub fn world_mut(&mut self) -> &mut World {
611 self.world
612 }
613}
614
615pub struct RelatedSpawnerCommands<'w, R: Relationship> {
618 target: Entity,
619 commands: Commands<'w, 'w>,
620 _marker: PhantomData<R>,
621}
622
623impl<'w, R: Relationship> RelatedSpawnerCommands<'w, R> {
624 pub fn new(commands: Commands<'w, 'w>, target: Entity) -> Self {
626 Self {
627 commands,
628 target,
629 _marker: PhantomData,
630 }
631 }
632
633 pub fn reborrow(&mut self) -> RelatedSpawnerCommands<'_, R> {
637 RelatedSpawnerCommands {
638 target: self.target,
639 commands: self.commands.reborrow(),
640 _marker: PhantomData,
641 }
642 }
643
644 pub fn spawn(&mut self, bundle: impl Bundle) -> EntityCommands<'_> {
647 self.commands.spawn((R::from(self.target), bundle))
648 }
649
650 pub fn spawn_empty(&mut self) -> EntityCommands<'_> {
653 self.commands.spawn(R::from(self.target))
654 }
655
656 pub fn target_entity(&self) -> Entity {
658 self.target
659 }
660
661 pub fn commands(&mut self) -> Commands<'_, '_> {
663 self.commands.reborrow()
664 }
665
666 pub fn commands_mut(&mut self) -> &mut Commands<'w, 'w> {
668 &mut self.commands
669 }
670}
671
672#[cfg(test)]
673mod tests {
674 use super::*;
675 use crate::prelude::{ChildOf, Children, Component};
676
677 #[derive(Component, Clone, Copy)]
678 struct TestComponent;
679
680 #[test]
681 fn insert_and_remove_recursive() {
682 let mut world = World::new();
683
684 let a = world.spawn_empty().id();
685 let b = world.spawn(ChildOf(a)).id();
686 let c = world.spawn(ChildOf(a)).id();
687 let d = world.spawn(ChildOf(b)).id();
688
689 world
690 .entity_mut(a)
691 .insert_recursive::<Children>(TestComponent);
692
693 for entity in [a, b, c, d] {
694 assert!(world.entity(entity).contains::<TestComponent>());
695 }
696
697 world
698 .entity_mut(b)
699 .remove_recursive::<Children, TestComponent>();
700
701 assert!(world.entity(a).contains::<TestComponent>());
703 assert!(!world.entity(b).contains::<TestComponent>());
705 assert!(world.entity(c).contains::<TestComponent>());
707 assert!(!world.entity(d).contains::<TestComponent>());
709
710 world
711 .entity_mut(a)
712 .remove_recursive::<Children, TestComponent>();
713
714 for entity in [a, b, c, d] {
715 assert!(!world.entity(entity).contains::<TestComponent>());
716 }
717 }
718
719 #[test]
720 fn remove_all_related() {
721 let mut world = World::new();
722
723 let a = world.spawn_empty().id();
724 let b = world.spawn(ChildOf(a)).id();
725 let c = world.spawn(ChildOf(a)).id();
726
727 world.entity_mut(a).detach_all_related::<ChildOf>();
728
729 assert_eq!(world.entity(a).get::<Children>(), None);
730 assert_eq!(world.entity(b).get::<ChildOf>(), None);
731 assert_eq!(world.entity(c).get::<ChildOf>(), None);
732 }
733
734 #[test]
735 fn replace_related_works() {
736 let mut world = World::new();
737 let child1 = world.spawn_empty().id();
738 let child2 = world.spawn_empty().id();
739 let child3 = world.spawn_empty().id();
740
741 let mut parent = world.spawn_empty();
742 parent.add_children(&[child1, child2]);
743 let child_value = ChildOf(parent.id());
744 let some_child = Some(&child_value);
745
746 parent.replace_children(&[child2, child3]);
747 let children = parent.get::<Children>().unwrap().collection();
748 assert_eq!(children, &[child2, child3]);
749 assert_eq!(parent.world().get::<ChildOf>(child1), None);
750 assert_eq!(parent.world().get::<ChildOf>(child2), some_child);
751 assert_eq!(parent.world().get::<ChildOf>(child3), some_child);
752
753 parent.replace_children_with_difference(&[child3], &[child1, child2], &[child1]);
754 let children = parent.get::<Children>().unwrap().collection();
755 assert_eq!(children, &[child1, child2]);
756 assert_eq!(parent.world().get::<ChildOf>(child1), some_child);
757 assert_eq!(parent.world().get::<ChildOf>(child2), some_child);
758 assert_eq!(parent.world().get::<ChildOf>(child3), None);
759 }
760
761 #[test]
762 fn add_related_keeps_relationship_data() {
763 #[derive(Component, PartialEq, Debug)]
764 #[relationship(relationship_target = Parent)]
765 struct Child {
766 #[relationship]
767 parent: Entity,
768 data: u8,
769 }
770
771 #[derive(Component)]
772 #[relationship_target(relationship = Child)]
773 struct Parent(Vec<Entity>);
774
775 let mut world = World::new();
776 let parent1 = world.spawn_empty().id();
777 let parent2 = world.spawn_empty().id();
778 let child = world
779 .spawn(Child {
780 parent: parent1,
781 data: 42,
782 })
783 .id();
784
785 world.entity_mut(parent2).add_related::<Child>(&[child]);
786 assert_eq!(
787 world.get::<Child>(child),
788 Some(&Child {
789 parent: parent2,
790 data: 42
791 })
792 );
793 }
794
795 #[test]
796 fn insert_related_keeps_relationship_data() {
797 #[derive(Component, PartialEq, Debug)]
798 #[relationship(relationship_target = Parent)]
799 struct Child {
800 #[relationship]
801 parent: Entity,
802 data: u8,
803 }
804
805 #[derive(Component)]
806 #[relationship_target(relationship = Child)]
807 struct Parent(Vec<Entity>);
808
809 let mut world = World::new();
810 let parent1 = world.spawn_empty().id();
811 let parent2 = world.spawn_empty().id();
812 let child = world
813 .spawn(Child {
814 parent: parent1,
815 data: 42,
816 })
817 .id();
818
819 world
820 .entity_mut(parent2)
821 .insert_related::<Child>(0, &[child]);
822 assert_eq!(
823 world.get::<Child>(child),
824 Some(&Child {
825 parent: parent2,
826 data: 42
827 })
828 );
829 }
830
831 #[test]
832 fn replace_related_keeps_relationship_data() {
833 #[derive(Component, PartialEq, Debug)]
834 #[relationship(relationship_target = Parent)]
835 struct Child {
836 #[relationship]
837 parent: Entity,
838 data: u8,
839 }
840
841 #[derive(Component)]
842 #[relationship_target(relationship = Child)]
843 struct Parent(Vec<Entity>);
844
845 let mut world = World::new();
846 let parent1 = world.spawn_empty().id();
847 let parent2 = world.spawn_empty().id();
848 let child = world
849 .spawn(Child {
850 parent: parent1,
851 data: 42,
852 })
853 .id();
854
855 world
856 .entity_mut(parent2)
857 .replace_related_with_difference::<Child>(&[], &[child], &[child]);
858 assert_eq!(
859 world.get::<Child>(child),
860 Some(&Child {
861 parent: parent2,
862 data: 42
863 })
864 );
865
866 world.entity_mut(parent1).replace_related::<Child>(&[child]);
867 assert_eq!(
868 world.get::<Child>(child),
869 Some(&Child {
870 parent: parent1,
871 data: 42
872 })
873 );
874 }
875
876 #[test]
877 fn replace_related_keeps_relationship_target_data() {
878 #[derive(Component)]
879 #[relationship(relationship_target = Parent)]
880 struct Child(Entity);
881
882 #[derive(Component)]
883 #[relationship_target(relationship = Child)]
884 struct Parent {
885 #[relationship]
886 children: Vec<Entity>,
887 data: u8,
888 }
889
890 let mut world = World::new();
891 let child1 = world.spawn_empty().id();
892 let child2 = world.spawn_empty().id();
893 let mut parent = world.spawn_empty();
894 parent.add_related::<Child>(&[child1]);
895 parent.get_mut::<Parent>().unwrap().data = 42;
896
897 parent.replace_related_with_difference::<Child>(&[child1], &[child2], &[child2]);
898 let data = parent.get::<Parent>().unwrap().data;
899 assert_eq!(data, 42);
900
901 parent.replace_related::<Child>(&[child1]);
902 let data = parent.get::<Parent>().unwrap().data;
903 assert_eq!(data, 42);
904 }
905
906 #[test]
907 fn despawn_related_observers_can_access_relationship_data() {
908 use crate::lifecycle::Discard;
909 use crate::observer::On;
910 use crate::prelude::Has;
911 use crate::system::Query;
912
913 #[derive(Component)]
914 struct MyComponent;
915
916 #[derive(Component, Default)]
917 struct ObserverResult {
918 success: bool,
919 }
920
921 let mut world = World::new();
922 let result_entity = world.spawn(ObserverResult::default()).id();
923
924 world.add_observer(
925 move |replace: On<Discard<MyComponent>>,
926 has_relationship: Query<Has<ChildOf>>,
927 mut results: Query<&mut ObserverResult>| {
928 if has_relationship.get(replace.entity).unwrap_or(false) {
929 results.get_mut(result_entity).unwrap().success = true;
930 }
931 },
932 );
933
934 let parent = world.spawn_empty().id();
935 let _child = world.spawn((MyComponent, ChildOf(parent))).id();
936
937 world.entity_mut(parent).despawn_related::<Children>();
938
939 assert!(world.get::<ObserverResult>(result_entity).unwrap().success);
940 }
941}