1use crate::{
2 bundle::Bundle,
3 entity::{hash_set::EntityHashSet, Entity},
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 clear_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 {
143 type Collection<R> =
144 <<R as Relationship>::RelationshipTarget as RelationshipTarget>::Collection;
145
146 if related.is_empty() {
147 self.remove::<R::RelationshipTarget>();
148
149 return self;
150 }
151
152 let Some(existing_relations) = self.get_mut::<R::RelationshipTarget>() else {
153 return self.add_related::<R>(related);
154 };
155
156 let mut relations = mem::replace(
159 existing_relations.into_inner(),
160 <R as Relationship>::RelationshipTarget::from_collection_risky(
161 Collection::<R>::with_capacity(0),
162 ),
163 );
164
165 let collection = relations.collection_mut_risky();
166
167 let mut potential_relations = EntityHashSet::from_iter(related.iter().copied());
168
169 let id = self.id();
170 self.world_scope(|world| {
171 for related in collection.iter() {
172 if !potential_relations.remove(related) {
173 world.entity_mut(related).remove::<R>();
174 }
175 }
176
177 for related in potential_relations {
178 world
180 .entity_mut(related)
181 .modify_or_insert_relation_with_relationship_hook_mode::<R>(
182 id,
183 RelationshipHookMode::Skip,
184 );
185 }
186 });
187
188 collection.clear();
190 collection.extend_from_iter(related.iter().copied());
191 self.insert(relations);
192
193 self
194 }
195
196 pub fn replace_related_with_difference<R: Relationship>(
217 &mut self,
218 entities_to_unrelate: &[Entity],
219 entities_to_relate: &[Entity],
220 newly_related_entities: &[Entity],
221 ) -> &mut Self {
222 #[cfg(debug_assertions)]
223 {
224 let entities_to_relate = EntityHashSet::from_iter(entities_to_relate.iter().copied());
225 let entities_to_unrelate =
226 EntityHashSet::from_iter(entities_to_unrelate.iter().copied());
227 let mut newly_related_entities =
228 EntityHashSet::from_iter(newly_related_entities.iter().copied());
229 assert!(
230 entities_to_relate.is_disjoint(&entities_to_unrelate),
231 "`entities_to_relate` ({entities_to_relate:?}) shared entities with `entities_to_unrelate` ({entities_to_unrelate:?})"
232 );
233 assert!(
234 newly_related_entities.is_disjoint(&entities_to_unrelate),
235 "`newly_related_entities` ({newly_related_entities:?}) shared entities with `entities_to_unrelate ({entities_to_unrelate:?})`"
236 );
237 assert!(
238 newly_related_entities.is_subset(&entities_to_relate),
239 "`newly_related_entities` ({newly_related_entities:?}) wasn't a subset of `entities_to_relate` ({entities_to_relate:?})"
240 );
241
242 if let Some(target) = self.get::<R::RelationshipTarget>() {
243 let existing_relationships: EntityHashSet = target.collection().iter().collect();
244
245 assert!(
246 existing_relationships.is_disjoint(&newly_related_entities),
247 "`newly_related_entities` contains an entity that wouldn't be newly related"
248 );
249
250 newly_related_entities.extend(existing_relationships);
251 newly_related_entities -= &entities_to_unrelate;
252 }
253
254 assert_eq!(newly_related_entities, entities_to_relate, "`entities_to_relate` ({entities_to_relate:?}) didn't contain all entities that would end up related");
255 };
256
257 match self.get_mut::<R::RelationshipTarget>() {
258 None => {
259 self.add_related::<R>(entities_to_relate);
260
261 return self;
262 }
263 Some(mut target) => {
264 let collection = target.collection_mut_risky();
266 collection.clear();
267
268 collection.extend_from_iter(entities_to_relate.iter().copied());
269 }
270 }
271
272 let this = self.id();
273 self.world_scope(|world| {
274 for unrelate in entities_to_unrelate {
275 world.entity_mut(*unrelate).remove::<R>();
276 }
277
278 for new_relation in newly_related_entities {
279 world
281 .entity_mut(*new_relation)
282 .modify_or_insert_relation_with_relationship_hook_mode::<R>(
283 this,
284 RelationshipHookMode::Skip,
285 );
286 }
287 });
288
289 self
290 }
291
292 pub fn add_one_related<R: Relationship>(&mut self, entity: Entity) -> &mut Self {
296 self.add_related::<R>(&[entity])
297 }
298
299 pub fn despawn_related<S: RelationshipTarget>(&mut self) -> &mut Self {
302 if let Some(sources) = self.get::<S>() {
303 let sources = sources.iter().collect::<Vec<_>>();
306 self.world_scope(|world| {
307 for entity in sources {
308 if let Ok(entity_mut) = world.get_entity_mut(entity) {
309 entity_mut.despawn();
310 };
311 }
312 });
313 }
314 self
315 }
316
317 pub fn despawn_children(&mut self) -> &mut Self {
322 self.despawn_related::<Children>();
323 self
324 }
325
326 pub fn insert_recursive<S: RelationshipTarget>(
336 &mut self,
337 bundle: impl Bundle + Clone,
338 ) -> &mut Self {
339 self.insert(bundle.clone());
340 if let Some(relationship_target) = self.get::<S>() {
341 let related_vec: Vec<Entity> = relationship_target.iter().collect();
342 for related in related_vec {
343 self.world_scope(|world| {
344 world
345 .entity_mut(related)
346 .insert_recursive::<S>(bundle.clone());
347 });
348 }
349 }
350
351 self
352 }
353
354 pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
362 self.remove::<B>();
363 if let Some(relationship_target) = self.get::<S>() {
364 let related_vec: Vec<Entity> = relationship_target.iter().collect();
365 for related in related_vec {
366 self.world_scope(|world| {
367 world.entity_mut(related).remove_recursive::<S, B>();
368 });
369 }
370 }
371
372 self
373 }
374
375 fn modify_or_insert_relation_with_relationship_hook_mode<R: Relationship>(
376 &mut self,
377 entity: Entity,
378 relationship_hook_mode: RelationshipHookMode,
379 ) {
380 if size_of::<R>() > size_of::<Entity>() {
382 self.assert_not_despawned();
383
384 let this = self.id();
385
386 let modified = self.world_scope(|world| {
387 let modified = DeferredWorld::from(&mut *world)
388 .modify_component_with_relationship_hook_mode::<R, _>(
389 this,
390 relationship_hook_mode,
391 |r| r.set_risky(entity),
392 )
393 .expect("entity access must be valid")
394 .is_some();
395
396 world.flush();
397
398 modified
399 });
400
401 if modified {
402 return;
403 }
404 }
405
406 self.insert_with_relationship_hook_mode(R::from(entity), relationship_hook_mode);
407 }
408}
409
410impl<'a> EntityCommands<'a> {
411 pub fn with_related<R: Relationship>(&mut self, bundle: impl Bundle) -> &mut Self {
413 let parent = self.id();
414 self.commands.spawn((bundle, R::from(parent)));
415 self
416 }
417
418 pub fn with_related_entities<R: Relationship>(
420 &mut self,
421 func: impl FnOnce(&mut RelatedSpawnerCommands<R>),
422 ) -> &mut Self {
423 let id = self.id();
424 func(&mut RelatedSpawnerCommands::new(self.commands(), id));
425 self
426 }
427
428 pub fn add_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
432 let related: Box<[Entity]> = related.into();
433
434 self.queue(move |mut entity: EntityWorldMut| {
435 entity.add_related::<R>(&related);
436 })
437 }
438
439 pub fn clear_related<R: Relationship>(&mut self) -> &mut Self {
441 self.queue(|mut entity: EntityWorldMut| {
442 entity.clear_related::<R>();
443 })
444 }
445
446 pub fn insert_related<R: Relationship>(&mut self, index: usize, related: &[Entity]) -> &mut Self
452 where
453 <R::RelationshipTarget as RelationshipTarget>::Collection:
454 OrderedRelationshipSourceCollection,
455 {
456 let related: Box<[Entity]> = related.into();
457
458 self.queue(move |mut entity: EntityWorldMut| {
459 entity.insert_related::<R>(index, &related);
460 })
461 }
462
463 pub fn add_one_related<R: Relationship>(&mut self, entity: Entity) -> &mut Self {
467 self.add_related::<R>(&[entity])
468 }
469
470 pub fn remove_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
472 let related: Box<[Entity]> = related.into();
473
474 self.queue(move |mut entity: EntityWorldMut| {
475 entity.remove_related::<R>(&related);
476 })
477 }
478
479 pub fn replace_related<R: Relationship>(&mut self, related: &[Entity]) -> &mut Self {
481 let related: Box<[Entity]> = related.into();
482
483 self.queue(move |mut entity: EntityWorldMut| {
484 entity.replace_related::<R>(&related);
485 })
486 }
487
488 pub fn replace_related_with_difference<R: Relationship>(
499 &mut self,
500 entities_to_unrelate: &[Entity],
501 entities_to_relate: &[Entity],
502 newly_related_entities: &[Entity],
503 ) -> &mut Self {
504 let entities_to_unrelate: Box<[Entity]> = entities_to_unrelate.into();
505 let entities_to_relate: Box<[Entity]> = entities_to_relate.into();
506 let newly_related_entities: Box<[Entity]> = newly_related_entities.into();
507
508 self.queue(move |mut entity: EntityWorldMut| {
509 entity.replace_related_with_difference::<R>(
510 &entities_to_unrelate,
511 &entities_to_relate,
512 &newly_related_entities,
513 );
514 })
515 }
516
517 pub fn despawn_related<S: RelationshipTarget>(&mut self) -> &mut Self {
520 self.queue(move |mut entity: EntityWorldMut| {
521 entity.despawn_related::<S>();
522 })
523 }
524
525 pub fn despawn_children(&mut self) -> &mut Self {
530 self.despawn_related::<Children>()
531 }
532
533 pub fn insert_recursive<S: RelationshipTarget>(
541 &mut self,
542 bundle: impl Bundle + Clone,
543 ) -> &mut Self {
544 self.queue(move |mut entity: EntityWorldMut| {
545 entity.insert_recursive::<S>(bundle);
546 })
547 }
548
549 pub fn remove_recursive<S: RelationshipTarget, B: Bundle>(&mut self) -> &mut Self {
557 self.queue(move |mut entity: EntityWorldMut| {
558 entity.remove_recursive::<S, B>();
559 })
560 }
561}
562
563pub struct RelatedSpawner<'w, R: Relationship> {
566 target: Entity,
567 world: &'w mut World,
568 _marker: PhantomData<R>,
569}
570
571impl<'w, R: Relationship> RelatedSpawner<'w, R> {
572 pub fn new(world: &'w mut World, target: Entity) -> Self {
574 Self {
575 world,
576 target,
577 _marker: PhantomData,
578 }
579 }
580
581 pub fn spawn(&mut self, bundle: impl Bundle) -> EntityWorldMut<'_> {
584 self.world.spawn((R::from(self.target), bundle))
585 }
586
587 pub fn spawn_empty(&mut self) -> EntityWorldMut<'_> {
590 self.world.spawn(R::from(self.target))
591 }
592
593 pub fn target_entity(&self) -> Entity {
595 self.target
596 }
597
598 pub fn world(&self) -> &World {
600 self.world
601 }
602
603 pub fn world_mut(&mut self) -> &mut World {
605 self.world
606 }
607}
608
609pub struct RelatedSpawnerCommands<'w, R: Relationship> {
612 target: Entity,
613 commands: Commands<'w, 'w>,
614 _marker: PhantomData<R>,
615}
616
617impl<'w, R: Relationship> RelatedSpawnerCommands<'w, R> {
618 pub fn new(commands: Commands<'w, 'w>, target: Entity) -> Self {
620 Self {
621 commands,
622 target,
623 _marker: PhantomData,
624 }
625 }
626
627 pub fn spawn(&mut self, bundle: impl Bundle) -> EntityCommands<'_> {
630 self.commands.spawn((R::from(self.target), bundle))
631 }
632
633 pub fn spawn_empty(&mut self) -> EntityCommands<'_> {
636 self.commands.spawn(R::from(self.target))
637 }
638
639 pub fn target_entity(&self) -> Entity {
641 self.target
642 }
643
644 pub fn commands(&mut self) -> Commands<'_, '_> {
646 self.commands.reborrow()
647 }
648
649 pub fn commands_mut(&mut self) -> &mut Commands<'w, 'w> {
651 &mut self.commands
652 }
653}
654
655#[cfg(test)]
656mod tests {
657 use super::*;
658 use crate::prelude::{ChildOf, Children, Component};
659
660 #[derive(Component, Clone, Copy)]
661 struct TestComponent;
662
663 #[test]
664 fn insert_and_remove_recursive() {
665 let mut world = World::new();
666
667 let a = world.spawn_empty().id();
668 let b = world.spawn(ChildOf(a)).id();
669 let c = world.spawn(ChildOf(a)).id();
670 let d = world.spawn(ChildOf(b)).id();
671
672 world
673 .entity_mut(a)
674 .insert_recursive::<Children>(TestComponent);
675
676 for entity in [a, b, c, d] {
677 assert!(world.entity(entity).contains::<TestComponent>());
678 }
679
680 world
681 .entity_mut(b)
682 .remove_recursive::<Children, TestComponent>();
683
684 assert!(world.entity(a).contains::<TestComponent>());
686 assert!(!world.entity(b).contains::<TestComponent>());
688 assert!(world.entity(c).contains::<TestComponent>());
690 assert!(!world.entity(d).contains::<TestComponent>());
692
693 world
694 .entity_mut(a)
695 .remove_recursive::<Children, TestComponent>();
696
697 for entity in [a, b, c, d] {
698 assert!(!world.entity(entity).contains::<TestComponent>());
699 }
700 }
701
702 #[test]
703 fn remove_all_related() {
704 let mut world = World::new();
705
706 let a = world.spawn_empty().id();
707 let b = world.spawn(ChildOf(a)).id();
708 let c = world.spawn(ChildOf(a)).id();
709
710 world.entity_mut(a).clear_related::<ChildOf>();
711
712 assert_eq!(world.entity(a).get::<Children>(), None);
713 assert_eq!(world.entity(b).get::<ChildOf>(), None);
714 assert_eq!(world.entity(c).get::<ChildOf>(), None);
715 }
716
717 #[test]
718 fn replace_related_works() {
719 let mut world = World::new();
720 let child1 = world.spawn_empty().id();
721 let child2 = world.spawn_empty().id();
722 let child3 = world.spawn_empty().id();
723
724 let mut parent = world.spawn_empty();
725 parent.add_children(&[child1, child2]);
726 let child_value = ChildOf(parent.id());
727 let some_child = Some(&child_value);
728
729 parent.replace_children(&[child2, child3]);
730 let children = parent.get::<Children>().unwrap().collection();
731 assert_eq!(children, &[child2, child3]);
732 assert_eq!(parent.world().get::<ChildOf>(child1), None);
733 assert_eq!(parent.world().get::<ChildOf>(child2), some_child);
734 assert_eq!(parent.world().get::<ChildOf>(child3), some_child);
735
736 parent.replace_children_with_difference(&[child3], &[child1, child2], &[child1]);
737 let children = parent.get::<Children>().unwrap().collection();
738 assert_eq!(children, &[child1, child2]);
739 assert_eq!(parent.world().get::<ChildOf>(child1), some_child);
740 assert_eq!(parent.world().get::<ChildOf>(child2), some_child);
741 assert_eq!(parent.world().get::<ChildOf>(child3), None);
742 }
743
744 #[test]
745 fn add_related_keeps_relationship_data() {
746 #[derive(Component, PartialEq, Debug)]
747 #[relationship(relationship_target = Parent)]
748 struct Child {
749 #[relationship]
750 parent: Entity,
751 data: u8,
752 }
753
754 #[derive(Component)]
755 #[relationship_target(relationship = Child)]
756 struct Parent(Vec<Entity>);
757
758 let mut world = World::new();
759 let parent1 = world.spawn_empty().id();
760 let parent2 = world.spawn_empty().id();
761 let child = world
762 .spawn(Child {
763 parent: parent1,
764 data: 42,
765 })
766 .id();
767
768 world.entity_mut(parent2).add_related::<Child>(&[child]);
769 assert_eq!(
770 world.get::<Child>(child),
771 Some(&Child {
772 parent: parent2,
773 data: 42
774 })
775 );
776 }
777
778 #[test]
779 fn insert_related_keeps_relationship_data() {
780 #[derive(Component, PartialEq, Debug)]
781 #[relationship(relationship_target = Parent)]
782 struct Child {
783 #[relationship]
784 parent: Entity,
785 data: u8,
786 }
787
788 #[derive(Component)]
789 #[relationship_target(relationship = Child)]
790 struct Parent(Vec<Entity>);
791
792 let mut world = World::new();
793 let parent1 = world.spawn_empty().id();
794 let parent2 = world.spawn_empty().id();
795 let child = world
796 .spawn(Child {
797 parent: parent1,
798 data: 42,
799 })
800 .id();
801
802 world
803 .entity_mut(parent2)
804 .insert_related::<Child>(0, &[child]);
805 assert_eq!(
806 world.get::<Child>(child),
807 Some(&Child {
808 parent: parent2,
809 data: 42
810 })
811 );
812 }
813
814 #[test]
815 fn replace_related_keeps_relationship_data() {
816 #[derive(Component, PartialEq, Debug)]
817 #[relationship(relationship_target = Parent)]
818 struct Child {
819 #[relationship]
820 parent: Entity,
821 data: u8,
822 }
823
824 #[derive(Component)]
825 #[relationship_target(relationship = Child)]
826 struct Parent(Vec<Entity>);
827
828 let mut world = World::new();
829 let parent1 = world.spawn_empty().id();
830 let parent2 = world.spawn_empty().id();
831 let child = world
832 .spawn(Child {
833 parent: parent1,
834 data: 42,
835 })
836 .id();
837
838 world
839 .entity_mut(parent2)
840 .replace_related_with_difference::<Child>(&[], &[child], &[child]);
841 assert_eq!(
842 world.get::<Child>(child),
843 Some(&Child {
844 parent: parent2,
845 data: 42
846 })
847 );
848
849 world.entity_mut(parent1).replace_related::<Child>(&[child]);
850 assert_eq!(
851 world.get::<Child>(child),
852 Some(&Child {
853 parent: parent1,
854 data: 42
855 })
856 );
857 }
858
859 #[test]
860 fn replace_related_keeps_relationship_target_data() {
861 #[derive(Component)]
862 #[relationship(relationship_target = Parent)]
863 struct Child(Entity);
864
865 #[derive(Component)]
866 #[relationship_target(relationship = Child)]
867 struct Parent {
868 #[relationship]
869 children: Vec<Entity>,
870 data: u8,
871 }
872
873 let mut world = World::new();
874 let child1 = world.spawn_empty().id();
875 let child2 = world.spawn_empty().id();
876 let mut parent = world.spawn_empty();
877 parent.add_related::<Child>(&[child1]);
878 parent.get_mut::<Parent>().unwrap().data = 42;
879
880 parent.replace_related_with_difference::<Child>(&[child1], &[child2], &[child2]);
881 let data = parent.get::<Parent>().unwrap().data;
882 assert_eq!(data, 42);
883
884 parent.replace_related::<Child>(&[child1]);
885 let data = parent.get::<Parent>().unwrap().data;
886 assert_eq!(data, 42);
887 }
888
889 #[test]
890 fn despawn_related_observers_can_access_relationship_data() {
891 use crate::lifecycle::Replace;
892 use crate::observer::On;
893 use crate::prelude::Has;
894 use crate::system::Query;
895
896 #[derive(Component)]
897 struct MyComponent;
898
899 #[derive(Component, Default)]
900 struct ObserverResult {
901 success: bool,
902 }
903
904 let mut world = World::new();
905 let result_entity = world.spawn(ObserverResult::default()).id();
906
907 world.add_observer(
908 move |replace: On<Replace, MyComponent>,
909 has_relationship: Query<Has<ChildOf>>,
910 mut results: Query<&mut ObserverResult>| {
911 if has_relationship.get(replace.entity).unwrap_or(false) {
912 results.get_mut(result_entity).unwrap().success = true;
913 }
914 },
915 );
916
917 let parent = world.spawn_empty().id();
918 let _child = world.spawn((MyComponent, ChildOf(parent))).id();
919
920 world.entity_mut(parent).despawn_related::<Children>();
921
922 assert!(world.get::<ObserverResult>(result_entity).unwrap().success);
923 }
924}