Fix a crash that occurs when an off-screen entity's material type changes and the material then becomes visible.

crates/bevy_pbr/src/material.rs

@@ -387,9 +387,16 @@ where
                         early_sweep_material_instances::<M>
                             .after(MaterialExtractionSystems)
                             .before(late_sweep_material_instances),
+                        // See the comments in
+                        // `sweep_entities_needing_specialization` for an
+                        // explanation of why the systems are ordered this way.
                         extract_entities_needs_specialization::<M>
+                            .in_set(MaterialExtractEntitiesNeedingSpecializationSystems),
+                        sweep_entities_needing_specialization::<M>
+                            .after(MaterialExtractEntitiesNeedingSpecializationSystems)
+                            .after(MaterialExtractionSystems)
                             .after(extract_cameras)
-                            .after(MaterialExtractionSystems),
+                            .before(late_sweep_material_instances),
                     ),
                 );
         }
@@ -604,6 +611,11 @@ pub struct RenderMaterialInstance {
 #[derive(SystemSet, Clone, PartialEq, Eq, Debug, Hash)]
 pub struct MaterialExtractionSystems;
 
+/// A [`SystemSet`] that contains all `extract_entities_needs_specialization`
+/// systems.
+#[derive(SystemSet, Clone, PartialEq, Eq, Debug, Hash)]
+pub struct MaterialExtractEntitiesNeedingSpecializationSystems;
+
 /// Deprecated alias for [`MaterialExtractionSystems`].
 #[deprecated(since = "0.17.0", note = "Renamed to `MaterialExtractionSystems`.")]
 pub type ExtractMaterialsSet = MaterialExtractionSystems;
@@ -750,10 +762,10 @@ fn early_sweep_material_instances<M>(
 /// Removes mesh materials from [`RenderMaterialInstances`] when their
 /// [`ViewVisibility`] components are removed.
 ///
-/// This runs after all invocations of [`early_sweep_material_instances`] and is
+/// This runs after all invocations of `early_sweep_material_instances` and is
 /// responsible for bumping [`RenderMaterialInstances::current_change_tick`] in
 /// preparation for a new frame.
-pub(crate) fn late_sweep_material_instances(
+pub fn late_sweep_material_instances(
     mut material_instances: ResMut<RenderMaterialInstances>,
     mut removed_meshes_query: Extract<RemovedComponents<Mesh3d>>,
 ) {
@@ -777,7 +789,39 @@ pub(crate) fn late_sweep_material_instances(
 
 pub fn extract_entities_needs_specialization<M>(
     entities_needing_specialization: Extract<Res<EntitiesNeedingSpecialization<M>>>,
-    material_instances: Res<RenderMaterialInstances>,
+    mut entity_specialization_ticks: ResMut<EntitySpecializationTicks>,
+    render_material_instances: Res<RenderMaterialInstances>,
+    ticks: SystemChangeTick,
+) where
+    M: Material,
+{
+    for entity in entities_needing_specialization.iter() {
+        // Update the entity's specialization tick with this run's tick
+        entity_specialization_ticks.insert(
+            (*entity).into(),
+            EntitySpecializationTickPair {
+                system_tick: ticks.this_run(),
+                material_instances_tick: render_material_instances.current_change_tick,
+            },
+        );
+    }
+}
+
+/// A system that runs after all instances of
+/// [`extract_entities_needs_specialization`] in order to delete specialization
+/// ticks for entities that are no longer renderable.
+///
+/// We delete entities from the [`EntitySpecializationTicks`] table *after*
+/// updating it with newly-discovered renderable entities in order to handle the
+/// case in which a single entity changes material types. If we naïvely removed
+/// entities from that table when their [`MeshMaterial3d<M>`] components were
+/// removed, and an entity changed material types, we might end up adding a new
+/// set of [`EntitySpecializationTickPair`] for the new material and then
+/// deleting it upon detecting the removed component for the old material.
+/// Deferring [`sweep_entities_needing_specialization`] to the end allows us to
+/// detect the case in which another material type updated the entity
+/// specialization ticks this frame and avoid deleting it if so.
+pub fn sweep_entities_needing_specialization<M>(
     mut entity_specialization_ticks: ResMut<EntitySpecializationTicks>,
     mut removed_mesh_material_components: Extract<RemovedComponents<MeshMaterial3d<M>>>,
     mut specialized_material_pipeline_cache: ResMut<SpecializedMaterialPipelineCache>,
@@ -787,24 +831,31 @@ pub fn extract_entities_needs_specialization<M>(
     mut specialized_shadow_material_pipeline_cache: Option<
         ResMut<SpecializedShadowMaterialPipelineCache>,
     >,
+    render_material_instances: Res<RenderMaterialInstances>,
     views: Query<&ExtractedView>,
-    ticks: SystemChangeTick,
 ) where
     M: Material,
 {
     // Clean up any despawned entities, we do this first in case the removed material was re-added
     // the same frame, thus will appear both in the removed components list and have been added to
     // the `EntitiesNeedingSpecialization` collection by triggering the `Changed` filter
     //
-    // Additionally, we need to make sure that we are careful about materials that could have changed
-    // type, e.g. from a `StandardMaterial` to a `CustomMaterial`, as this will also appear in the
-    // removed components list. As such, we make sure that this system runs after `MaterialExtractionSystems`
-    // so that the `RenderMaterialInstances` bookkeeping has already been done, and we can check if the entity
-    // still has a valid material instance.
+    // Additionally, we need to make sure that we are careful about materials
+    // that could have changed type, e.g. from a `StandardMaterial` to a
+    // `CustomMaterial`, as this will also appear in the removed components
+    // list. As such, we make sure that this system runs after
+    // `extract_entities_needs_specialization` so that the entity specialization
+    // tick bookkeeping has already been done, and we can check if the entity's
+    // tick was updated this frame.
     for entity in removed_mesh_material_components.read() {
-        if material_instances
-            .instances
-            .contains_key(&MainEntity::from(entity))
+        // If the entity's specialization tick was updated this frame, that
+        // means that that entity changed materials this frame. Don't remove the
+        // entity from the table in that case.
+        if entity_specialization_ticks
+            .get(&MainEntity::from(entity))
+            .is_some_and(|ticks| {
+                ticks.material_instances_tick == render_material_instances.current_change_tick
+            })
         {
             continue;
         }
@@ -830,11 +881,6 @@ pub fn extract_entities_needs_specialization<M>(
             }
         }
     }
-
-    for entity in entities_needing_specialization.iter() {
-        // Update the entity's specialization tick with this run's tick
-        entity_specialization_ticks.insert((*entity).into(), ticks.this_run());
-    }
 }
 
 #[derive(Resource, Deref, DerefMut, Clone, Debug)]
@@ -853,10 +899,58 @@ impl<M> Default for EntitiesNeedingSpecialization<M> {
     }
 }
 
+/// Stores ticks specifying the last time Bevy specialized the pipelines of each
+/// entity.
+///
+/// Every entity that has a mesh and material must be present in this table,
+/// even if that mesh isn't visible.
 #[derive(Resource, Deref, DerefMut, Default, Clone, Debug)]
 pub struct EntitySpecializationTicks {
+    /// A mapping from each main entity to ticks that specify the last time this
+    /// entity's pipeline was specialized.
+    ///
+    /// Every entity that has a mesh and material must be present in this table,
+    /// even if that mesh isn't visible.
     #[deref]
-    pub entities: MainEntityHashMap<Tick>,
+    pub entities: MainEntityHashMap<EntitySpecializationTickPair>,
+}
+
+/// Ticks that specify the last time an entity's pipeline was specialized.
+///
+/// We need two different types of ticks here for a subtle reason. First, we
+/// need the [`Self::system_tick`], which maps to Bevy's [`SystemChangeTick`],
+/// because that's what we use in [`specialize_material_meshes`] to check
+/// whether pipelines need specialization. But we also need
+/// [`Self::material_instances_tick`], which maps to the
+/// [`RenderMaterialInstances::current_change_tick`]. That's because the latter
+/// only changes once per frame, which is a guarantee we need to handle the
+/// following case:
+///
+/// 1. The app removes material A from a mesh and replaces it with material B.
+///    Both A and B are of different [`Material`] types entirely.
+///
+/// 2. [`extract_entities_needs_specialization`] runs for material B and marks
+///    the mesh as up to date by recording the current tick.
+///
+/// 3. [`sweep_entities_needing_specialization`] runs for material A and checks
+///    to ensure it's safe to remove the [`EntitySpecializationTickPair`] for the mesh
+///    from the [`EntitySpecializationTicks`]. To do this, it needs to know
+///    whether [`extract_entities_needs_specialization`] for some *different*
+///    material (in this case, material B) ran earlier in the frame and updated the
+///    change tick, and to skip removing the [`EntitySpecializationTickPair`] if so.
+///    It can't reliably use the [`Self::system_tick`] to determine this because
+///    the [`SystemChangeTick`] can be updated multiple times in the same frame.
+///    Instead, it needs a type of tick that's updated only once per frame, after
+///    all materials' versions of [`sweep_entities_needing_specialization`] have
+///    run. The [`RenderMaterialInstances`] tick satisfies this criterion, and so
+///    that's what [`sweep_entities_needing_specialization`] uses.
+#[derive(Clone, Copy, Debug)]
+pub struct EntitySpecializationTickPair {
+    /// The standard Bevy system tick.
+    pub system_tick: Tick,
+    /// The tick in [`RenderMaterialInstances`], which is updated in
+    /// `late_sweep_material_instances`.
+    pub material_instances_tick: Tick,
 }
 
 /// Stores the [`SpecializedMaterialViewPipelineCache`] for each view.
@@ -966,7 +1060,10 @@ pub fn specialize_material_meshes(
             else {
                 continue;
             };
-            let entity_tick = entity_specialization_ticks.get(visible_entity).unwrap();
+            let entity_tick = entity_specialization_ticks
+                .get(visible_entity)
+                .unwrap()
+                .system_tick;
             let last_specialized_tick = view_specialized_material_pipeline_cache
                 .get(visible_entity)
                 .map(|(tick, _)| *tick);

crates/bevy_pbr/src/prepass/mod.rs

@@ -877,7 +877,10 @@ pub fn specialize_prepass_material_meshes(
             else {
                 continue;
             };
-            let entity_tick = entity_specialization_ticks.get(visible_entity).unwrap();
+            let entity_tick = entity_specialization_ticks
+                .get(visible_entity)
+                .unwrap()
+                .system_tick;
             let last_specialized_tick = view_specialized_material_pipeline_cache
                 .get(visible_entity)
                 .map(|(tick, _)| *tick);

crates/bevy_pbr/src/render/light.rs

@@ -1844,7 +1844,9 @@ pub fn specialize_shadows(
                     .map(|(tick, _)| *tick);
                 let needs_specialization = last_specialized_tick.is_none_or(|tick| {
                     view_tick.is_newer_than(tick, ticks.this_run())
-                        || entity_tick.is_newer_than(tick, ticks.this_run())
+                        || entity_tick
+                            .system_tick
+                            .is_newer_than(tick, ticks.this_run())
                 });
                 if !needs_specialization {
                     continue;

crates/bevy_sprite_render/src/mesh2d/material.rs

@@ -283,7 +283,7 @@ where
 
         if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
             render_app
-                .init_resource::<EntitySpecializationTicks<M>>()
+                .init_resource::<EntitySpecializationTickPair<M>>()
                 .init_resource::<SpecializedMaterial2dPipelineCache<M>>()
                 .add_render_command::<Opaque2d, DrawMaterial2d<M>>()
                 .add_render_command::<AlphaMask2d, DrawMaterial2d<M>>()
@@ -566,7 +566,7 @@ pub const fn tonemapping_pipeline_key(tonemapping: Tonemapping) -> Mesh2dPipelin
 
 pub fn extract_entities_needs_specialization<M>(
     entities_needing_specialization: Extract<Res<EntitiesNeedingSpecialization<M>>>,
-    mut entity_specialization_ticks: ResMut<EntitySpecializationTicks<M>>,
+    mut entity_specialization_ticks: ResMut<EntitySpecializationTickPair<M>>,
     mut removed_mesh_material_components: Extract<RemovedComponents<MeshMaterial2d<M>>>,
     mut specialized_material2d_pipeline_cache: ResMut<SpecializedMaterial2dPipelineCache<M>>,
     views: Query<&MainEntity, With<ExtractedView>>,
@@ -608,13 +608,13 @@ impl<M> Default for EntitiesNeedingSpecialization<M> {
 }
 
 #[derive(Clone, Resource, Deref, DerefMut, Debug)]
-pub struct EntitySpecializationTicks<M> {
+pub struct EntitySpecializationTickPair<M> {
     #[deref]
     pub entities: MainEntityHashMap<Tick>,
     _marker: PhantomData<M>,
 }
 
-impl<M> Default for EntitySpecializationTicks<M> {
+impl<M> Default for EntitySpecializationTickPair<M> {
     fn default() -> Self {
         Self {
             entities: MainEntityHashMap::default(),
@@ -702,7 +702,7 @@ pub fn specialize_material2d_meshes<M: Material2d>(
     alpha_mask_render_phases: Res<ViewBinnedRenderPhases<AlphaMask2d>>,
     views: Query<(&MainEntity, &ExtractedView, &RenderVisibleEntities)>,
     view_key_cache: Res<ViewKeyCache>,
-    entity_specialization_ticks: Res<EntitySpecializationTicks<M>>,
+    entity_specialization_ticks: Res<EntitySpecializationTickPair<M>>,
     view_specialization_ticks: Res<ViewSpecializationTicks>,
     ticks: SystemChangeTick,
     mut specialized_material_pipeline_cache: ResMut<SpecializedMaterial2dPipelineCache<M>>,

examples/3d/manual_material.rs

@@ -8,8 +8,10 @@ use bevy::{
         SystemChangeTick, SystemParamItem,
     },
     pbr::{
-        DrawMaterial, EntitiesNeedingSpecialization, EntitySpecializationTicks,
-        MaterialBindGroupAllocator, MaterialBindGroupAllocators, MaterialDrawFunction,
+        late_sweep_material_instances, DrawMaterial, EntitiesNeedingSpecialization,
+        EntitySpecializationTickPair, EntitySpecializationTicks, MaterialBindGroupAllocator,
+        MaterialBindGroupAllocators, MaterialDrawFunction,
+        MaterialExtractEntitiesNeedingSpecializationSystems, MaterialExtractionSystems,
         MaterialFragmentShader, MaterialProperties, PreparedMaterial, RenderMaterialBindings,
         RenderMaterialInstance, RenderMaterialInstances, SpecializedMaterialPipelineCache,
     },
@@ -66,7 +68,12 @@ impl Plugin for ImageMaterialPlugin {
                 ExtractSchedule,
                 (
                     extract_image_materials,
-                    extract_image_materials_needing_specialization,
+                    extract_image_materials_needing_specialization
+                        .in_set(MaterialExtractEntitiesNeedingSpecializationSystems),
+                    sweep_image_materials_needing_specialization
+                        .after(MaterialExtractEntitiesNeedingSpecializationSystems)
+                        .after(MaterialExtractionSystems)
+                        .before(late_sweep_material_instances),
                 ),
             );
     }
@@ -285,6 +292,7 @@ fn extract_image_materials_needing_specialization(
     mut entity_specialization_ticks: ResMut<EntitySpecializationTicks>,
     mut removed_mesh_material_components: Extract<RemovedComponents<ImageMaterial3d>>,
     mut specialized_material_pipeline_cache: ResMut<SpecializedMaterialPipelineCache>,
+    render_material_instances: Res<RenderMaterialInstances>,
     views: Query<&ExtractedView>,
     ticks: SystemChangeTick,
 ) {
@@ -304,6 +312,44 @@ fn extract_image_materials_needing_specialization(
 
     for entity in entities_needing_specialization.iter() {
         // Update the entity's specialization tick with this run's tick
-        entity_specialization_ticks.insert((*entity).into(), ticks.this_run());
+        entity_specialization_ticks.insert(
+            (*entity).into(),
+            EntitySpecializationTickPair {
+                system_tick: ticks.this_run(),
+                material_instances_tick: render_material_instances.current_change_tick,
+            },
+        );
+    }
+}
+
+fn sweep_image_materials_needing_specialization(
+    mut entity_specialization_ticks: ResMut<EntitySpecializationTicks>,
+    mut removed_mesh_material_components: Extract<RemovedComponents<ImageMaterial3d>>,
+    mut specialized_material_pipeline_cache: ResMut<SpecializedMaterialPipelineCache>,
+    render_material_instances: Res<RenderMaterialInstances>,
+    views: Query<&ExtractedView>,
+) {
+    // Clean up any despawned entities, we do this first in case the removed material was re-added
+    // the same frame, thus will appear both in the removed components list and have been added to
+    // the `EntitiesNeedingSpecialization` collection by triggering the `Changed` filter
+    for entity in removed_mesh_material_components.read() {
+        if entity_specialization_ticks
+            .get(&MainEntity::from(entity))
+            .is_some_and(|ticks| {
+                ticks.material_instances_tick == render_material_instances.current_change_tick
+            })
+        {
+            continue;
+        }
+
+        entity_specialization_ticks.remove(&MainEntity::from(entity));
+
+        for view in views {
+            if let Some(cache) =
+                specialized_material_pipeline_cache.get_mut(&view.retained_view_entity)
+            {
+                cache.remove(&MainEntity::from(entity));
+            }
+        }
     }
 }