Synchronize removed components with the render world

crates/bevy_pbr/src/lib.rs

@@ -124,6 +124,7 @@ use bevy_render::{
     render_asset::prepare_assets,
     render_graph::RenderGraph,
     render_resource::Shader,
+    sync_component::SyncComponentPlugin,
     texture::{GpuImage, Image},
     view::{check_visibility, VisibilitySystems},
     ExtractSchedule, Render, RenderApp, RenderSet,
@@ -308,7 +309,6 @@ impl Plugin for PbrPlugin {
             .register_type::<PointLight>()
             .register_type::<PointLightShadowMap>()
             .register_type::<SpotLight>()
-            .register_type::<DistanceFog>()
             .register_type::<ShadowFilteringMethod>()
             .init_resource::<AmbientLight>()
             .init_resource::<GlobalVisibleClusterableObjects>()
@@ -340,6 +340,11 @@ impl Plugin for PbrPlugin {
                 VolumetricFogPlugin,
                 ScreenSpaceReflectionsPlugin,
             ))
+            .add_plugins((
+                SyncComponentPlugin::<DirectionalLight>::default(),
+                SyncComponentPlugin::<PointLight>::default(),
+                SyncComponentPlugin::<SpotLight>::default(),
+            ))
             .configure_sets(
                 PostUpdate,
                 (

crates/bevy_pbr/src/light/directional_light.rs

@@ -1,4 +1,4 @@
-use bevy_render::{view::Visibility, world_sync::SyncToRenderWorld};
+use bevy_render::view::Visibility;
 
 use super::*;
 
@@ -57,8 +57,7 @@ use super::*;
     CascadeShadowConfig,
     CascadesVisibleEntities,
     Transform,
-    Visibility,
-    SyncToRenderWorld
+    Visibility
 )]
 pub struct DirectionalLight {
     /// The color of the light.

crates/bevy_pbr/src/light/point_light.rs

@@ -1,4 +1,4 @@
-use bevy_render::{view::Visibility, world_sync::SyncToRenderWorld};
+use bevy_render::view::Visibility;
 
 use super::*;
 
@@ -21,13 +21,7 @@ use super::*;
 /// Source: [Wikipedia](https://en.wikipedia.org/wiki/Lumen_(unit)#Lighting)
 #[derive(Component, Debug, Clone, Copy, Reflect)]
 #[reflect(Component, Default, Debug)]
-#[require(
-    CubemapFrusta,
-    CubemapVisibleEntities,
-    Transform,
-    Visibility,
-    SyncToRenderWorld
-)]
+#[require(CubemapFrusta, CubemapVisibleEntities, Transform, Visibility)]
 pub struct PointLight {
     /// The color of this light source.
     pub color: Color,

crates/bevy_pbr/src/light/spot_light.rs

@@ -1,4 +1,4 @@
-use bevy_render::{view::Visibility, world_sync::SyncToRenderWorld};
+use bevy_render::view::Visibility;
 
 use super::*;
 
@@ -9,7 +9,7 @@ use super::*;
 /// the transform, and can be specified with [`Transform::looking_at`](Transform::looking_at).
 #[derive(Component, Debug, Clone, Copy, Reflect)]
 #[reflect(Component, Default, Debug)]
-#[require(Frustum, VisibleMeshEntities, Transform, Visibility, SyncToRenderWorld)]
+#[require(Frustum, VisibleMeshEntities, Transform, Visibility)]
 pub struct SpotLight {
     /// The color of the light.
     ///

crates/bevy_pbr/src/render/mesh.rs

@@ -37,6 +37,7 @@ use bevy_render::{
         prepare_view_targets, GpuCulling, RenderVisibilityRanges, ViewTarget, ViewUniformOffset,
         ViewVisibility, VisibilityRange,
     },
+    world_sync::RenderEntity,
     Extract,
 };
 use bevy_transform::components::GlobalTransform;
@@ -549,11 +550,11 @@ pub enum RenderMeshInstanceGpuQueue {
     /// The version of [`RenderMeshInstanceGpuQueue`] that omits the
     /// [`MeshCullingData`], so that we don't waste space when GPU
     /// culling is disabled.
-    CpuCulling(Vec<(Entity, RenderMeshInstanceGpuBuilder)>),
+    CpuCulling(Vec<(RenderEntity, RenderMeshInstanceGpuBuilder)>),
     /// The version of [`RenderMeshInstanceGpuQueue`] that contains the
     /// [`MeshCullingData`], used when any view has GPU culling
     /// enabled.
-    GpuCulling(Vec<(Entity, RenderMeshInstanceGpuBuilder, MeshCullingData)>),
+    GpuCulling(Vec<(RenderEntity, RenderMeshInstanceGpuBuilder, MeshCullingData)>),
 }
 
 /// The per-thread queues containing mesh instances, populated during the
@@ -732,7 +733,7 @@ impl RenderMeshInstanceGpuQueue {
     /// Adds a new mesh to this queue.
     fn push(
         &mut self,
-        entity: Entity,
+        entity: RenderEntity,
         instance_builder: RenderMeshInstanceGpuBuilder,
         culling_data_builder: Option<MeshCullingData>,
     ) {
@@ -964,6 +965,7 @@ pub fn extract_meshes_for_gpu_building(
     meshes_query: Extract<
         Query<(
             Entity,
+            &RenderEntity,
             &ViewVisibility,
             &GlobalTransform,
             Option<&PreviousGlobalTransform>,
@@ -998,6 +1000,7 @@ pub fn extract_meshes_for_gpu_building(
         |queue,
          (
             entity,
+            render_entity,
             view_visibility,
             transform,
             previous_transform,
@@ -1056,7 +1059,11 @@ pub fn extract_meshes_for_gpu_building(
                 previous_input_index,
             };
 
-            queue.push(entity, gpu_mesh_instance_builder, gpu_mesh_culling_data);
+            queue.push(
+                *render_entity,
+                gpu_mesh_instance_builder,
+                gpu_mesh_culling_data,
+            );
         },
     );
 }
@@ -1129,7 +1136,7 @@ pub fn collect_meshes_for_gpu_building(
             RenderMeshInstanceGpuQueue::CpuCulling(ref mut queue) => {
                 for (entity, mesh_instance_builder) in queue.drain(..) {
                     mesh_instance_builder.add_to(
-                        entity,
+                        entity.id(),
                         &mut *render_mesh_instances,
                         current_input_buffer,
                         &mesh_allocator,
@@ -1139,7 +1146,7 @@ pub fn collect_meshes_for_gpu_building(
             RenderMeshInstanceGpuQueue::GpuCulling(ref mut queue) => {
                 for (entity, mesh_instance_builder, mesh_culling_builder) in queue.drain(..) {
                     let instance_data_index = mesh_instance_builder.add_to(
-                        entity,
+                        entity.id(),
                         &mut *render_mesh_instances,
                         current_input_buffer,
                         &mesh_allocator,

crates/bevy_pbr/src/volumetric_fog/mod.rs

@@ -51,6 +51,7 @@ use bevy_render::{
     mesh::{Mesh, Meshable},
     render_graph::{RenderGraphApp, ViewNodeRunner},
     render_resource::{Shader, SpecializedRenderPipelines},
+    sync_component::SyncComponentPlugin,
     texture::Image,
     view::{InheritedVisibility, ViewVisibility, Visibility},
     ExtractSchedule, Render, RenderApp, RenderSet,
@@ -231,6 +232,8 @@ impl Plugin for VolumetricFogPlugin {
         app.register_type::<VolumetricFog>()
             .register_type::<VolumetricLight>();
 
+        app.add_plugins(SyncComponentPlugin::<FogVolume>::default());
+
         let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
             return;
         };

crates/bevy_render/src/extract_component.rs

@@ -1,8 +1,9 @@
 use crate::{
     render_resource::{encase::internal::WriteInto, DynamicUniformBuffer, ShaderType},
     renderer::{RenderDevice, RenderQueue},
+    sync_component::SyncComponentPlugin,
     view::ViewVisibility,
-    world_sync::{RenderEntity, SyncToRenderWorld},
+    world_sync::RenderEntity,
     Extract, ExtractSchedule, Render, RenderApp, RenderSet,
 };
 use bevy_app::{App, Plugin};
@@ -12,7 +13,6 @@ use bevy_ecs::{
     prelude::*,
     query::{QueryFilter, QueryItem, ReadOnlyQueryData},
     system::lifetimeless::Read,
-    world::OnAdd,
 };
 use core::{marker::PhantomData, ops::Deref};
 
@@ -160,15 +160,6 @@ fn prepare_uniform_components<C>(
 /// This plugin extracts the components into the render world for synced entities.
 ///
 /// To do so, it sets up the [`ExtractSchedule`] step for the specified [`ExtractComponent`].
-///
-/// # Warning
-///
-/// Be careful when removing the [`ExtractComponent`] from an entity. When an [`ExtractComponent`]
-/// is added to an entity, that entity is automatically synced with the render world (see also
-/// [`WorldSyncPlugin`](crate::world_sync::WorldSyncPlugin)). When removing the entity in the main
-/// world, the synced entity also gets removed. However, if only the [`ExtractComponent`] is removed
-/// this *doesn't* happen, and the synced entity stays around with the old extracted data.
-/// We recommend despawning the entire entity, instead of only removing [`ExtractComponent`].
 pub struct ExtractComponentPlugin<C, F = ()> {
     only_extract_visible: bool,
     marker: PhantomData<fn() -> (C, F)>,
@@ -194,10 +185,8 @@ impl<C, F> ExtractComponentPlugin<C, F> {
 
 impl<C: ExtractComponent> Plugin for ExtractComponentPlugin<C> {
     fn build(&self, app: &mut App) {
-        // TODO: use required components
-        app.observe(|trigger: Trigger<OnAdd, C>, mut commands: Commands| {
-            commands.entity(trigger.entity()).insert(SyncToRenderWorld);
-        });
+        app.add_plugins(SyncComponentPlugin::<C>::default());
+
         if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
             if self.only_extract_visible {
                 render_app.add_systems(ExtractSchedule, extract_visible_components::<C>);
@@ -219,7 +208,7 @@ impl<T: Asset> ExtractComponent for Handle<T> {
     }
 }
 
-/// This system extracts all components of the corresponding [`ExtractComponent`], for entities that are synced via [`SyncToRenderWorld`].
+/// This system extracts all components of the corresponding [`ExtractComponent`], for entities that are synced via [`crate::world_sync::SyncToRenderWorld`].
 fn extract_components<C: ExtractComponent>(
     mut commands: Commands,
     mut previous_len: Local<usize>,
@@ -235,7 +224,7 @@ fn extract_components<C: ExtractComponent>(
     commands.insert_or_spawn_batch(values);
 }
 
-/// This system extracts all components of the corresponding [`ExtractComponent`], for entities that are visible and synced via [`SyncToRenderWorld`].
+/// This system extracts all components of the corresponding [`ExtractComponent`], for entities that are visible and synced via [`crate::world_sync::SyncToRenderWorld`].
 fn extract_visible_components<C: ExtractComponent>(
     mut commands: Commands,
     mut previous_len: Local<usize>,

crates/bevy_render/src/gpu_readback.rs

@@ -6,6 +6,7 @@ use crate::{
     renderer::{render_system, RenderDevice},
     storage::{GpuShaderStorageBuffer, ShaderStorageBuffer},
     texture::{GpuImage, TextureFormatPixelInfo},
+    world_sync::MainEntity,
     ExtractSchedule, MainWorld, Render, RenderApp, RenderSet,
 };
 use async_channel::{Receiver, Sender};
@@ -232,7 +233,7 @@ fn prepare_buffers(
     mut buffer_pool: ResMut<GpuReadbackBufferPool>,
     gpu_images: Res<RenderAssets<GpuImage>>,
     ssbos: Res<RenderAssets<GpuShaderStorageBuffer>>,
-    handles: Query<(Entity, &Readback)>,
+    handles: Query<(&MainEntity, &Readback)>,
 ) {
     for (entity, readback) in handles.iter() {
         match readback {
@@ -254,7 +255,7 @@ fn prepare_buffers(
                     );
                     let (tx, rx) = async_channel::bounded(1);
                     readbacks.requested.push(GpuReadback {
-                        entity,
+                        entity: entity.id(),
                         src: ReadbackSource::Texture {
                             texture: gpu_image.texture.clone(),
                             layout,
@@ -272,7 +273,7 @@ fn prepare_buffers(
                     let buffer = buffer_pool.get(&render_device, size);
                     let (tx, rx) = async_channel::bounded(1);
                     readbacks.requested.push(GpuReadback {
-                        entity,
+                        entity: entity.id(),
                         src: ReadbackSource::Buffer {
                             src_start: 0,
                             dst_start: 0,

crates/bevy_render/src/lib.rs

@@ -38,6 +38,7 @@ pub mod renderer;
 pub mod settings;
 mod spatial_bundle;
 pub mod storage;
+pub mod sync_component;
 pub mod texture;
 pub mod view;
 pub mod world_sync;

crates/bevy_render/src/sync_component.rs

@@ -0,0 +1,27 @@
+use core::marker::PhantomData;
+
+use bevy_app::{App, Plugin};
+use bevy_ecs::component::Component;
+
+use crate::world_sync::{EntityRecord, PendingSyncEntity, SyncToRenderWorld};
+
+pub struct SyncComponentPlugin<C: Component>(PhantomData<C>);
+
+impl<C: Component> Default for SyncComponentPlugin<C> {
+    fn default() -> Self {
+        Self(PhantomData)
+    }
+}
+
+impl<C: Component> Plugin for SyncComponentPlugin<C> {
+    fn build(&self, app: &mut App) {
+        app.register_required_components::<C, SyncToRenderWorld>();
+
+        app.world_mut().register_component_hooks::<C>().on_remove(
+            |mut world, entity, _component_id| {
+                let mut pending = world.resource_mut::<PendingSyncEntity>();
+                pending.push(EntityRecord::ComponentRemoved(entity));
+            },
+        );
+    }
+}

crates/bevy_render/src/world_sync.rs

@@ -86,7 +86,7 @@ impl Plugin for WorldSyncPlugin {
              mut pending: ResMut<PendingSyncEntity>,
              query: Query<&RenderEntity>| {
                 if let Ok(e) = query.get(trigger.entity()) {
-                    pending.push(EntityRecord::Removed(e.id()));
+                    pending.push(EntityRecord::Removed(*e));
                 };
             },
         );
@@ -127,13 +127,17 @@ impl MainEntity {
 pub struct TemporaryRenderEntity;
 
 /// A record enum to what entities with [`SyncToRenderWorld`] have been added or removed.
+#[derive(Debug)]
 pub(crate) enum EntityRecord {
     /// When an entity is spawned on the main world, notify the render world so that it can spawn a corresponding
     /// entity. This contains the main world entity.
     Added(Entity),
     /// When an entity is despawned on the main world, notify the render world so that the corresponding entity can be
     /// despawned. This contains the render world entity.
-    Removed(Entity),
+    Removed(RenderEntity),
+    /// When a component is removed from an entity, notify the render world so that the corresponding component can be
+    /// removed. This contains the main world entity.
+    ComponentRemoved(Entity),
 }
 
 // Entity Record in MainWorld pending to Sync
@@ -148,8 +152,8 @@ pub(crate) fn entity_sync_system(main_world: &mut World, render_world: &mut Worl
         for record in pending.drain(..) {
             match record {
                 EntityRecord::Added(e) => {
-                    if let Some(mut entity) = world.get_entity_mut(e) {
-                        match entity.entry::<RenderEntity>() {
+                    if let Some(mut main_entity) = world.get_entity_mut(e) {
+                        match main_entity.entry::<RenderEntity>() {
                             bevy_ecs::world::Entry::Occupied(_) => {
                                 panic!("Attempting to synchronize an entity that has already been synchronized!");
                             }
@@ -161,11 +165,24 @@ pub(crate) fn entity_sync_system(main_world: &mut World, render_world: &mut Worl
                         };
                     }
                 }
-                EntityRecord::Removed(e) => {
-                    if let Some(ec) = render_world.get_entity_mut(e) {
+                EntityRecord::Removed(render_entity) => {
+                    if let Some(ec) = render_world.get_entity_mut(render_entity.id()) {
                         ec.despawn();
                     };
                 }
+                EntityRecord::ComponentRemoved(main_entity) => {
+                    // It's difficult to remove only the relevant component because component ids aren't stable across worlds,
+                    // so we just clear the entire render world entity.
+                    let Some(mut render_entity) = world.get_mut::<RenderEntity>(main_entity) else {
+                        continue;
+                    };
+                    if let Some(render_world_entity) = render_world.get_entity_mut(render_entity.id()) {
+                        render_world_entity.despawn();
+
+                        let id = render_world.spawn(MainEntity(main_entity)).id();
+                        render_entity.0 = id;
+                    }
+                },
             }
         }
     });
@@ -222,7 +239,7 @@ mod tests {
              mut pending: ResMut<PendingSyncEntity>,
              query: Query<&RenderEntity>| {
                 if let Ok(e) = query.get(trigger.entity()) {
-                    pending.push(EntityRecord::Removed(e.id()));
+                    pending.push(EntityRecord::Removed(*e));
                 };
             },
         );