1use crate::{
2 prelude::Mut,
3 reflect::{AppTypeRegistry, ReflectBundle, ReflectComponent},
4 resource::Resource,
5 system::EntityCommands,
6 world::EntityWorldMut,
7};
8use alloc::{borrow::Cow, boxed::Box};
9use bevy_reflect::{PartialReflect, TypeRegistry};
10
11pub trait ReflectCommandExt {
13 fn insert_reflect(&mut self, component: Box<dyn PartialReflect>) -> &mut Self;
88
89 fn insert_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
100 &mut self,
101 component: Box<dyn PartialReflect>,
102 ) -> &mut Self;
103
104 fn remove_reflect(&mut self, component_type_path: impl Into<Cow<'static, str>>) -> &mut Self;
162 fn remove_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
165 &mut self,
166 component_type_path: impl Into<Cow<'static, str>>,
167 ) -> &mut Self;
168}
169
170impl ReflectCommandExt for EntityCommands<'_> {
171 fn insert_reflect(&mut self, component: Box<dyn PartialReflect>) -> &mut Self {
172 self.queue(move |mut entity: EntityWorldMut| {
173 entity.insert_reflect(component);
174 })
175 }
176
177 fn insert_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
178 &mut self,
179 component: Box<dyn PartialReflect>,
180 ) -> &mut Self {
181 self.queue(move |mut entity: EntityWorldMut| {
182 entity.insert_reflect_with_registry::<T>(component);
183 })
184 }
185
186 fn remove_reflect(&mut self, component_type_path: impl Into<Cow<'static, str>>) -> &mut Self {
187 let component_type_path: Cow<'static, str> = component_type_path.into();
188 self.queue(move |mut entity: EntityWorldMut| {
189 entity.remove_reflect(component_type_path);
190 })
191 }
192
193 fn remove_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
194 &mut self,
195 component_type_path: impl Into<Cow<'static, str>>,
196 ) -> &mut Self {
197 let component_type_path: Cow<'static, str> = component_type_path.into();
198 self.queue(move |mut entity: EntityWorldMut| {
199 entity.remove_reflect_with_registry::<T>(component_type_path);
200 })
201 }
202}
203
204impl<'w> EntityWorldMut<'w> {
205 pub fn insert_reflect(&mut self, component: Box<dyn PartialReflect>) -> &mut Self {
223 self.assert_not_despawned();
224 self.resource_scope(|entity, registry: Mut<AppTypeRegistry>| {
225 let type_registry = ®istry.as_ref().read();
226 insert_reflect_with_registry_ref(entity, type_registry, component);
227 });
228 self
229 }
230
231 pub fn insert_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
244 &mut self,
245 component: Box<dyn PartialReflect>,
246 ) -> &mut Self {
247 self.assert_not_despawned();
248 self.resource_scope(|entity, registry: Mut<T>| {
249 let type_registry = registry.as_ref().as_ref();
250 insert_reflect_with_registry_ref(entity, type_registry, component);
251 });
252 self
253 }
254
255 pub fn remove_reflect(&mut self, component_type_path: Cow<'static, str>) -> &mut Self {
274 self.assert_not_despawned();
275 self.resource_scope(|entity, registry: Mut<AppTypeRegistry>| {
276 let type_registry = ®istry.as_ref().read();
277 remove_reflect_with_registry_ref(entity, type_registry, component_type_path);
278 });
279 self
280 }
281
282 pub fn remove_reflect_with_registry<T: Resource + AsRef<TypeRegistry>>(
297 &mut self,
298 component_type_path: Cow<'static, str>,
299 ) -> &mut Self {
300 self.assert_not_despawned();
301 self.resource_scope(|entity, registry: Mut<T>| {
302 let type_registry = registry.as_ref().as_ref();
303 remove_reflect_with_registry_ref(entity, type_registry, component_type_path);
304 });
305 self
306 }
307}
308
309fn insert_reflect_with_registry_ref(
311 entity: &mut EntityWorldMut,
312 type_registry: &TypeRegistry,
313 component: Box<dyn PartialReflect>,
314) {
315 let type_info = component
316 .get_represented_type_info()
317 .expect("component should represent a type.");
318 let type_path = type_info.type_path();
319 let Some(type_registration) = type_registry.get(type_info.type_id()) else {
320 panic!("`{type_path}` should be registered in type registry via `App::register_type<{type_path}>`");
321 };
322
323 if let Some(reflect_component) = type_registration.data::<ReflectComponent>() {
324 reflect_component.insert(entity, component.as_partial_reflect(), type_registry);
325 } else if let Some(reflect_bundle) = type_registration.data::<ReflectBundle>() {
326 reflect_bundle.insert(entity, component.as_partial_reflect(), type_registry);
327 } else {
328 panic!("`{type_path}` should have #[reflect(Component)] or #[reflect(Bundle)]");
329 }
330}
331
332fn remove_reflect_with_registry_ref(
334 entity: &mut EntityWorldMut,
335 type_registry: &TypeRegistry,
336 component_type_path: Cow<'static, str>,
337) {
338 let Some(type_registration) = type_registry.get_with_type_path(&component_type_path) else {
339 return;
340 };
341 if let Some(reflect_component) = type_registration.data::<ReflectComponent>() {
342 reflect_component.remove(entity);
343 } else if let Some(reflect_bundle) = type_registration.data::<ReflectBundle>() {
344 reflect_bundle.remove(entity);
345 }
346}
347
348#[cfg(test)]
349mod tests {
350 use crate::{
351 bundle::Bundle,
352 component::Component,
353 prelude::{AppTypeRegistry, ReflectComponent},
354 reflect::{ReflectBundle, ReflectCommandExt},
355 system::{Commands, SystemState},
356 world::World,
357 };
358 use alloc::{borrow::ToOwned, boxed::Box};
359 use bevy_ecs_macros::Resource;
360 use bevy_reflect::{PartialReflect, Reflect, TypeRegistry};
361
362 #[derive(Resource)]
363 struct TypeRegistryResource {
364 type_registry: TypeRegistry,
365 }
366
367 impl AsRef<TypeRegistry> for TypeRegistryResource {
368 fn as_ref(&self) -> &TypeRegistry {
369 &self.type_registry
370 }
371 }
372
373 #[derive(Component, Reflect, Default, PartialEq, Eq, Debug)]
374 #[reflect(Component)]
375 struct ComponentA(u32);
376
377 #[derive(Component, Reflect, Default, PartialEq, Eq, Debug)]
378 #[reflect(Component)]
379 struct ComponentB(u32);
380
381 #[derive(Bundle, Reflect, Default, Debug, PartialEq)]
382 #[reflect(Bundle)]
383 struct BundleA {
384 a: ComponentA,
385 b: ComponentB,
386 }
387
388 #[test]
389 fn insert_reflected() {
390 let mut world = World::new();
391
392 let type_registry = AppTypeRegistry::default();
393 {
394 let mut registry = type_registry.write();
395 registry.register::<ComponentA>();
396 registry.register_type_data::<ComponentA, ReflectComponent>();
397 }
398 world.insert_resource(type_registry);
399
400 let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
401 let mut commands = system_state.get_mut(&mut world);
402
403 let entity = commands.spawn_empty().id();
404 let entity2 = commands.spawn_empty().id();
405 let entity3 = commands.spawn_empty().id();
406
407 let boxed_reflect_component_a = Box::new(ComponentA(916)) as Box<dyn PartialReflect>;
408 let boxed_reflect_component_a_clone = boxed_reflect_component_a.reflect_clone().unwrap();
409 let boxed_reflect_component_a_dynamic = boxed_reflect_component_a.to_dynamic();
410
411 commands
412 .entity(entity)
413 .insert_reflect(boxed_reflect_component_a);
414 commands
415 .entity(entity2)
416 .insert_reflect(boxed_reflect_component_a_clone.into_partial_reflect());
417 commands
418 .entity(entity3)
419 .insert_reflect(boxed_reflect_component_a_dynamic);
420 system_state.apply(&mut world);
421
422 assert_eq!(
423 world.entity(entity).get::<ComponentA>(),
424 world.entity(entity2).get::<ComponentA>(),
425 );
426 assert_eq!(
427 world.entity(entity).get::<ComponentA>(),
428 world.entity(entity3).get::<ComponentA>(),
429 );
430 }
431
432 #[test]
433 fn insert_reflected_with_registry() {
434 let mut world = World::new();
435
436 let mut type_registry = TypeRegistryResource {
437 type_registry: TypeRegistry::new(),
438 };
439
440 type_registry.type_registry.register::<ComponentA>();
441 type_registry
442 .type_registry
443 .register_type_data::<ComponentA, ReflectComponent>();
444 world.insert_resource(type_registry);
445
446 let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
447 let mut commands = system_state.get_mut(&mut world);
448
449 let entity = commands.spawn_empty().id();
450
451 let boxed_reflect_component_a = Box::new(ComponentA(916)) as Box<dyn PartialReflect>;
452
453 commands
454 .entity(entity)
455 .insert_reflect_with_registry::<TypeRegistryResource>(boxed_reflect_component_a);
456 system_state.apply(&mut world);
457
458 assert_eq!(
459 world.entity(entity).get::<ComponentA>(),
460 Some(&ComponentA(916))
461 );
462 }
463
464 #[test]
465 fn remove_reflected() {
466 let mut world = World::new();
467
468 let type_registry = AppTypeRegistry::default();
469 {
470 let mut registry = type_registry.write();
471 registry.register::<ComponentA>();
472 registry.register_type_data::<ComponentA, ReflectComponent>();
473 }
474 world.insert_resource(type_registry);
475
476 let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
477 let mut commands = system_state.get_mut(&mut world);
478
479 let entity = commands.spawn(ComponentA(0)).id();
480
481 let boxed_reflect_component_a = Box::new(ComponentA(916)) as Box<dyn Reflect>;
482
483 commands
484 .entity(entity)
485 .remove_reflect(boxed_reflect_component_a.reflect_type_path().to_owned());
486 system_state.apply(&mut world);
487
488 assert_eq!(world.entity(entity).get::<ComponentA>(), None);
489 }
490
491 #[test]
492 fn remove_reflected_with_registry() {
493 let mut world = World::new();
494
495 let mut type_registry = TypeRegistryResource {
496 type_registry: TypeRegistry::new(),
497 };
498
499 type_registry.type_registry.register::<ComponentA>();
500 type_registry
501 .type_registry
502 .register_type_data::<ComponentA, ReflectComponent>();
503 world.insert_resource(type_registry);
504
505 let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
506 let mut commands = system_state.get_mut(&mut world);
507
508 let entity = commands.spawn(ComponentA(0)).id();
509
510 let boxed_reflect_component_a = Box::new(ComponentA(916)) as Box<dyn Reflect>;
511
512 commands
513 .entity(entity)
514 .remove_reflect_with_registry::<TypeRegistryResource>(
515 boxed_reflect_component_a.reflect_type_path().to_owned(),
516 );
517 system_state.apply(&mut world);
518
519 assert_eq!(world.entity(entity).get::<ComponentA>(), None);
520 }
521
522 #[test]
523 fn insert_reflect_bundle() {
524 let mut world = World::new();
525
526 let type_registry = AppTypeRegistry::default();
527 {
528 let mut registry = type_registry.write();
529 registry.register::<BundleA>();
530 registry.register_type_data::<BundleA, ReflectBundle>();
531 }
532 world.insert_resource(type_registry);
533
534 let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
535 let mut commands = system_state.get_mut(&mut world);
536
537 let entity = commands.spawn_empty().id();
538 let bundle = Box::new(BundleA {
539 a: ComponentA(31),
540 b: ComponentB(20),
541 }) as Box<dyn PartialReflect>;
542 commands.entity(entity).insert_reflect(bundle);
543
544 system_state.apply(&mut world);
545
546 assert_eq!(world.get::<ComponentA>(entity), Some(&ComponentA(31)));
547 assert_eq!(world.get::<ComponentB>(entity), Some(&ComponentB(20)));
548 }
549
550 #[test]
551 fn insert_reflect_bundle_with_registry() {
552 let mut world = World::new();
553
554 let mut type_registry = TypeRegistryResource {
555 type_registry: TypeRegistry::new(),
556 };
557
558 type_registry.type_registry.register::<BundleA>();
559 type_registry
560 .type_registry
561 .register_type_data::<BundleA, ReflectBundle>();
562 world.insert_resource(type_registry);
563
564 let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
565 let mut commands = system_state.get_mut(&mut world);
566
567 let entity = commands.spawn_empty().id();
568 let bundle = Box::new(BundleA {
569 a: ComponentA(31),
570 b: ComponentB(20),
571 }) as Box<dyn PartialReflect>;
572
573 commands
574 .entity(entity)
575 .insert_reflect_with_registry::<TypeRegistryResource>(bundle);
576 system_state.apply(&mut world);
577
578 assert_eq!(world.get::<ComponentA>(entity), Some(&ComponentA(31)));
579 assert_eq!(world.get::<ComponentB>(entity), Some(&ComponentB(20)));
580 }
581
582 #[test]
583 fn remove_reflected_bundle() {
584 let mut world = World::new();
585
586 let type_registry = AppTypeRegistry::default();
587 {
588 let mut registry = type_registry.write();
589 registry.register::<BundleA>();
590 registry.register_type_data::<BundleA, ReflectBundle>();
591 }
592 world.insert_resource(type_registry);
593
594 let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
595 let mut commands = system_state.get_mut(&mut world);
596
597 let entity = commands
598 .spawn(BundleA {
599 a: ComponentA(31),
600 b: ComponentB(20),
601 })
602 .id();
603
604 let boxed_reflect_bundle_a = Box::new(BundleA {
605 a: ComponentA(1),
606 b: ComponentB(23),
607 }) as Box<dyn Reflect>;
608
609 commands
610 .entity(entity)
611 .remove_reflect(boxed_reflect_bundle_a.reflect_type_path().to_owned());
612 system_state.apply(&mut world);
613
614 assert_eq!(world.entity(entity).get::<ComponentA>(), None);
615 assert_eq!(world.entity(entity).get::<ComponentB>(), None);
616 }
617
618 #[test]
619 fn remove_reflected_bundle_with_registry() {
620 let mut world = World::new();
621
622 let mut type_registry = TypeRegistryResource {
623 type_registry: TypeRegistry::new(),
624 };
625
626 type_registry.type_registry.register::<BundleA>();
627 type_registry
628 .type_registry
629 .register_type_data::<BundleA, ReflectBundle>();
630 world.insert_resource(type_registry);
631
632 let mut system_state: SystemState<Commands> = SystemState::new(&mut world);
633 let mut commands = system_state.get_mut(&mut world);
634
635 let entity = commands
636 .spawn(BundleA {
637 a: ComponentA(31),
638 b: ComponentB(20),
639 })
640 .id();
641
642 let boxed_reflect_bundle_a = Box::new(BundleA {
643 a: ComponentA(1),
644 b: ComponentB(23),
645 }) as Box<dyn Reflect>;
646
647 commands
648 .entity(entity)
649 .remove_reflect_with_registry::<TypeRegistryResource>(
650 boxed_reflect_bundle_a.reflect_type_path().to_owned(),
651 );
652 system_state.apply(&mut world);
653
654 assert_eq!(world.entity(entity).get::<ComponentA>(), None);
655 assert_eq!(world.entity(entity).get::<ComponentB>(), None);
656 }
657}