From f2cf02408ff8766ef53ad7fa2e42fd30d2f0f8e5 Mon Sep 17 00:00:00 2001
From: Alix Bott <bott.alix@gmail.com>
Date: Fri, 30 Aug 2024 02:43:56 +0200
Subject: [PATCH] Fix observer unregistering unsetting archetype flags (#14963)

# Objective

- Fixes https://github.com/bevyengine/bevy/issues/14961

## Solution

- Check that the archetypes don't contain any other observed components
before unsetting their flags

## Testing

- I added a regression test: `observer_despawn_archetype_flags`
---
 crates/bevy_ecs/src/archetype.rs    |  2 +-
 crates/bevy_ecs/src/observer/mod.rs | 32 ++++++++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/crates/bevy_ecs/src/archetype.rs b/crates/bevy_ecs/src/archetype.rs
index b7cad4e389..15e962291c 100644
--- a/crates/bevy_ecs/src/archetype.rs
+++ b/crates/bevy_ecs/src/archetype.rs
@@ -367,7 +367,7 @@ pub struct Archetype {
     edges: Edges,
     entities: Vec<ArchetypeEntity>,
     components: ImmutableSparseSet<ComponentId, ArchetypeComponentInfo>,
-    flags: ArchetypeFlags,
+    pub(crate) flags: ArchetypeFlags,
 }
 
 impl Archetype {
diff --git a/crates/bevy_ecs/src/observer/mod.rs b/crates/bevy_ecs/src/observer/mod.rs
index 3ef2d6b28d..c6ace80ca5 100644
--- a/crates/bevy_ecs/src/observer/mod.rs
+++ b/crates/bevy_ecs/src/observer/mod.rs
@@ -429,7 +429,17 @@ impl World {
                     if observers.map.is_empty() && observers.entity_map.is_empty() {
                         cache.component_observers.remove(component);
                         if let Some(flag) = Observers::is_archetype_cached(event_type) {
-                            archetypes.update_flags(*component, flag, false);
+                            for archetype in &mut archetypes.archetypes {
+                                if archetype.contains(*component) {
+                                    let no_longer_observed = archetype
+                                        .components()
+                                        .all(|id| !cache.component_observers.contains_key(&id));
+
+                                    if no_longer_observed {
+                                        archetype.flags.set(flag, false);
+                                    }
+                                }
+                            }
                         }
                     }
                 }
@@ -656,6 +666,26 @@ mod tests {
         world.spawn(A).flush();
     }
 
+    // Regression test for https://github.com/bevyengine/bevy/issues/14961
+    #[test]
+    fn observer_despawn_archetype_flags() {
+        let mut world = World::new();
+        world.init_resource::<R>();
+
+        let entity = world.spawn((A, B)).flush();
+
+        world.observe(|_: Trigger<OnRemove, A>, mut res: ResMut<R>| res.0 += 1);
+
+        let observer = world
+            .observe(|_: Trigger<OnRemove, B>| panic!("Observer triggered after being despawned."))
+            .flush();
+        world.despawn(observer);
+
+        world.despawn(entity);
+
+        assert_eq!(1, world.resource::<R>().0);
+    }
+
     #[test]
     fn observer_multiple_matches() {
         let mut world = World::new();