Fix partition change tracking (#66)

Fix an issue with partition change tracking for stationary entities, add
regression tests.
Snuck in a passhash perf improvement I noticed while I was in there.

Author: aevyrie

src/grid/propagation.rs

@@ -174,12 +174,16 @@ impl Grid {
     ) {
         let start = bevy_platform::time::Instant::now();
         let update_transforms = |low_precision_root, parent_transform: Ref<GlobalTransform>| {
-            // High precision global transforms are change-detected, and are only updated if that
+            // High precision global transforms are change-detected and are only updated if that
             // entity has moved relative to the floating origin's grid cell.
             let changed = parent_transform.is_changed();
 
+            #[expect(
+                unsafe_code,
+                reason = "`propagate_recursive()` is unsafe due to its use of `Query::get_unchecked()`."
+            )]
             // SAFETY:
-            // - Unlike the bevy version of this, we do not iterate over all children of the root,
+            // - Unlike the bevy version of this, we do not iterate over all children of the root
             //   and manually verify each child has a parent component that points back to the same
             //   entity. Instead, we query the roots directly, so we know they are unique.
             // - We may operate as if all descendants are consistent, since `propagate_recursive`
@@ -189,10 +193,6 @@ impl Grid {
             //   other root entities' `propagate_recursive` calls will not conflict with this one.
             // - Since this is the only place where `transform_query` gets used, there will be no
             //   conflicting fetches elsewhere.
-            #[expect(
-                unsafe_code,
-                reason = "`propagate_recursive()` is unsafe due to its use of `Query::get_unchecked()`."
-            )]
             unsafe {
                 Self::propagate_recursive(
                     &parent_transform,

src/partition/change_tracking.rs

@@ -8,7 +8,6 @@ use crate::hash::SpatialHashFilter;
 use crate::hash::SpatialHashSystems;
 use crate::partition::map::PartitionLookup;
 use crate::partition::PartitionId;
-use alloc::vec::Vec;
 use bevy_app::prelude::*;
 use bevy_ecs::entity::EntityHashMap;
 use bevy_ecs::prelude::*;
@@ -20,7 +19,7 @@ use core::marker::PhantomData;
 pub struct PartitionChangePlugin<F: SpatialHashFilter = ()>(PhantomData<F>);
 
 impl<F: SpatialHashFilter> PartitionChangePlugin<F> {
-    /// Create a new instance of [`crate::prelude::PartitionChangePlugin`].
+    /// Create a new instance of [`PartitionChangePlugin`].
     pub fn new() -> Self {
         Self(PhantomData)
     }
@@ -68,88 +67,87 @@ impl<F: SpatialHashFilter> FromWorld for PartitionEntities<F> {
 
 impl<F: SpatialHashFilter> PartitionEntities<F> {
     fn update(
-        mut entity_partitions: ResMut<Self>,
+        mut this: ResMut<Self>,
         cells: Res<CellLookup<F>>,
         changed_cells: Res<ChangedCells<F>>,
         all_hashes: Query<(Entity, &CellId), F>,
         mut old_reverse: Local<CellHashMap<PartitionId>>,
         partitions: Res<PartitionLookup<F>>,
     ) {
-        entity_partitions.changed.clear();
+        // 1. Clear the list of entities that have changed partitions
+        this.changed.clear();
 
-        // Compute cell-level partition changes: cells that remained occupied but changed partitions.
-        for (cell_hash, entry) in cells.all_entries() {
-            if let Some(&new_pid) = partitions.get(cell_hash) {
-                // Only mark entities whose previous recorded partition differs from the new one
-                for entity in entry.entities.iter().copied() {
-                    if let Some(prev_pid) = entity_partitions.map.get(&entity).copied() {
-                        if prev_pid != new_pid {
-                            if let Some((_from, to)) = entity_partitions.changed.get_mut(&entity) {
-                                *to = Some(new_pid);
-                            } else {
-                                entity_partitions
-                                    .changed
-                                    .insert(entity, (Some(prev_pid), Some(new_pid)));
-                            }
-                        }
-                    }
-                }
-            }
-        }
-
-        // Compute entity-level partition changes for entities that moved cells this frame,
-        // were newly spawned (added cell), or had cell removed (despawned/removed component).
+        // 2. Iterate through all entities that have moved between cells, using the already
+        // optimized `ChangedCells` resource used for grid cells. Check these moved entities to see
+        // if they have also changed partitions. This should also include spawned/despawned.
         for entity in changed_cells.iter() {
             match all_hashes.get(*entity) {
-                // Entity currently has a CellId
                 Ok((entity_id, cell_hash)) => {
                     let new_pid = partitions.get(cell_hash).copied();
-                    let old_pid = entity_partitions.map.get(&entity_id).copied();
-                    let record = match (old_pid, new_pid) {
+                    let old_pid = this.map.get(&entity_id).copied();
+                    let partition_change = match (old_pid, new_pid) {
                         (Some(o), Some(n)) if o == n => None, // Partition unchanged
                         (None, None) => None,                 // Nonsensical
-                        other => Some(other),
+                        other => Some(other),                 // Valid change
                     };
-                    if let Some((from, to)) = record {
-                        if let Some((existing_from, existing_to)) =
-                            entity_partitions.changed.get_mut(entity)
-                        {
+                    if let Some((from, to)) = partition_change {
+                        if let Some((existing_from, existing_to)) = this.changed.get_mut(entity) {
                             // Preserve the earliest known source if we already have one
                             if existing_from.is_none() {
                                 *existing_from = from;
                             }
                             *existing_to = to;
                         } else {
-                            entity_partitions.changed.insert(*entity, (from, to));
+                            this.changed.insert(*entity, (from, to));
                         }
                     }
                 }
-                // Entity no longer has a CellId -> removed/despawned
+                // If the query fails, the entity no longer has a `CellId` because it was removed
+                // or the entity was despawned.
                 Err(_) => {
-                    if let Some(prev_pid) = entity_partitions.map.get(entity).copied() {
-                        entity_partitions
-                            .changed
-                            .insert(*entity, (Some(prev_pid), None));
+                    if let Some(prev_pid) = this.map.get(entity).copied() {
+                        this.changed.insert(*entity, (Some(prev_pid), None));
                     }
                 }
             }
         }
 
-        // Apply the delta only after all changes have been collected.
-        let mut apply_list: Vec<(Entity, Option<PartitionId>)> =
-            Vec::with_capacity(entity_partitions.changed.len());
-        for (entity, (_from, to)) in entity_partitions.changed.iter() {
-            apply_list.push((*entity, *to));
+        // 3. Consider entities that have not moved, but their partition has changed out from
+        // underneath them. This can happen when partitions merge and split - the entity did not
+        // move but is now in a new partition.
+        //
+        // Check these changes at the cell level so we scale with the number of cells, not the
+        // number of entities, and additionally avoid many lookups.
+        //
+        // It's important this is run after the moved-entity checks in step 2, because the entities
+        // found from cell changes may *also* have moved, and we would miss that information if we
+        // only used cell-level tracking.
+        for (cell_id, entry) in cells.all_entries() {
+            if let Some(&new_pid) = partitions.get(cell_id) {
+                if let Some(&old_pid) = old_reverse.get(cell_id) {
+                    if new_pid != old_pid {
+                        for entity in entry.entities.iter().copied() {
+                            this.changed
+                                .entry(entity)
+                                // Don't overwrite entities that moved cells, they have already been tracked.
+                                .or_insert((Some(old_pid), Some(new_pid)));
+                        }
+                    }
+                }
+            }
         }
-        for (entity, to) in apply_list {
-            match to {
+
+        // 4. Apply the changes to the entity partition map after all changes have been collected.
+        let PartitionEntities { map, changed, .. } = this.as_mut();
+        for (entity, (_source, destination)) in changed.iter() {
+            match destination {
                 Some(pid) => {
-                    entity_partitions.map.insert(entity, pid);
+                    map.insert(*entity, *pid);
                 }
                 None => {
-                    entity_partitions.map.remove(&entity);
+                    map.remove(entity);
                 }
-            }
+            };
         }
 
         // Snapshot the cell->partition mapping to compute deltas next update

src/partition/map.rs

@@ -10,6 +10,7 @@ use alloc::vec;
 use alloc::vec::Vec;
 use bevy_ecs::prelude::*;
 use bevy_platform::collections::HashMap;
+use bevy_platform::hash::PassHash;
 use bevy_platform::time::Instant;
 use bevy_tasks::{ComputeTaskPool, ParallelSliceMut};
 use core::marker::PhantomData;
@@ -31,7 +32,7 @@ pub struct PartitionLookup<F = ()>
 where
     F: SpatialHashFilter,
 {
-    partitions: HashMap<PartitionId, Partition>,
+    partitions: HashMap<PartitionId, Partition, PassHash>,
     pub(crate) reverse_map: CellHashMap<PartitionId>,
     next_partition: u64,
     spooky: PhantomData<F>,
@@ -55,7 +56,7 @@ impl<F> Deref for PartitionLookup<F>
 where
     F: SpatialHashFilter,
 {
-    type Target = HashMap<PartitionId, Partition>;
+    type Target = HashMap<PartitionId, Partition, PassHash>;
 
     fn deref(&self) -> &Self::Target {
         &self.partitions
@@ -187,7 +188,7 @@ where
         cells: Res<CellLookup<F>>,
         // Scratch space allocations
         mut added_neighbors: Local<Vec<PartitionId>>,
-        mut split_candidates_map: Local<HashMap<PartitionId, CellHashSet>>,
+        mut split_candidates_map: Local<HashMap<PartitionId, CellHashSet, PassHash>>,
         mut split_candidates: Local<Vec<(PartitionId, CellHashSet)>>,
         mut split_results: Local<Vec<Vec<SplitResult>>>,
     ) {

src/partition/tests.rs

@@ -295,3 +295,178 @@ fn split_then_merge_back_same_frame_no_false_positive() {
     assert!(ep.changed.get(&app.world().resource::<R>().left).is_none());
     assert!(ep.changed.get(&app.world().resource::<R>().right).is_none());
 }
+
+#[test]
+fn partition_split_records_changes_for_stationary_entities() {
+    let mut app = App::new();
+    app.add_plugins((
+        BigSpaceMinimalPlugins,
+        CellHashingPlugin::default(),
+        PartitionPlugin::default(),
+        PartitionChangePlugin::default(),
+    ));
+
+    #[derive(Resource, Clone, Copy)]
+    struct R {
+        left: Entity,
+        mid: Entity,
+        right: Entity,
+    }
+
+    let setup = |mut commands: Commands| {
+        commands.spawn_big_space_default(|root| {
+            let left = root.spawn_spatial(CellCoord::new(0, 0, 0)).id();
+            let mid = root.spawn_spatial(CellCoord::new(1, 0, 0)).id();
+            let right = root.spawn_spatial(CellCoord::new(2, 0, 0)).id();
+            root.commands().insert_resource(R { left, mid, right });
+        });
+    };
+    app.add_systems(Startup, setup);
+
+    for _ in 0..10 {
+        run_app_once(&mut app);
+    }
+
+    let (left, right, mid) = {
+        let r = app.world().resource::<R>();
+        (r.left, r.right, r.mid)
+    };
+
+    let (pid_left, pid_mid, pid_right) = {
+        let parts = app.world().resource::<PartitionLookup>();
+        let cell_left = *app.world().get::<CellId>(left).unwrap();
+        let cell_mid = *app.world().get::<CellId>(mid).unwrap();
+        let cell_right = *app.world().get::<CellId>(right).unwrap();
+        (
+            parts.get(&cell_left).copied().unwrap(),
+            parts.get(&cell_mid).copied().unwrap(),
+            parts.get(&cell_right).copied().unwrap(),
+        )
+    };
+    assert_eq!(pid_left, pid_mid);
+    assert_eq!(pid_mid, pid_right);
+
+    // Remove the middle cell to cause a split
+    app.world_mut().commands().entity(mid).despawn();
+
+    let mut left_changed = false;
+    let mut right_changed = false;
+    let mut left_pid_entities = None;
+    let mut right_pid_entities = None;
+
+    for _ in 0..100 {
+        run_app_once(&mut app);
+        let ep = app.world().resource::<PartitionEntities>();
+        if ep.changed.contains_key(&left) {
+            left_changed = true;
+        }
+        if ep.changed.contains_key(&right) {
+            right_changed = true;
+        }
+        left_pid_entities = ep.map.get(&left).copied();
+        right_pid_entities = ep.map.get(&right).copied();
+    }
+
+    let parts = app.world().resource::<PartitionLookup>();
+    let cell_left = *app.world().get::<CellId>(left).unwrap();
+    let cell_right = *app.world().get::<CellId>(right).unwrap();
+    let left_pid_actual = parts.get(&cell_left).copied().unwrap();
+    let right_pid_actual = parts.get(&cell_right).copied().unwrap();
+
+    assert_ne!(
+        left_pid_actual, right_pid_actual,
+        "Partitions should have split in PartitionLookup"
+    );
+    assert_ne!(
+        left_pid_entities, right_pid_entities,
+        "Partitions should have split in PartitionEntities"
+    );
+
+    assert!(
+        left_changed || right_changed,
+        "At least one entity should have recorded a partition change"
+    );
+}
+
+#[test]
+fn partition_merge_records_changes_for_stationary_entities() {
+    let mut app = App::new();
+    app.add_plugins((
+        BigSpaceMinimalPlugins,
+        CellHashingPlugin::default(),
+        PartitionPlugin::default(),
+        PartitionChangePlugin::default(),
+    ));
+
+    #[derive(Resource, Clone, Copy)]
+    struct R {
+        root: Entity,
+        left: Entity,
+        right: Entity,
+    }
+
+    let setup = |mut commands: Commands| {
+        commands.spawn_big_space_default(|root| {
+            let left = root.spawn_spatial(CellCoord::new(0, 0, 0)).id();
+            let right = root.spawn_spatial(CellCoord::new(2, 0, 0)).id();
+            let root_id = root.id();
+            root.commands().insert_resource(R {
+                root: root_id,
+                left,
+                right,
+            });
+        });
+    };
+    app.add_systems(Startup, setup);
+
+    for _ in 0..10 {
+        run_app_once(&mut app);
+    }
+
+    let (left, right) = {
+        let r = app.world().resource::<R>();
+        (r.left, r.right)
+    };
+
+    let ep = app.world().resource::<PartitionEntities>();
+    let left_pid_initial = *ep.map.get(&left).expect("left not mapped");
+    let right_pid_initial = *ep.map.get(&right).expect("right not mapped");
+    assert_ne!(left_pid_initial, right_pid_initial);
+
+    // Add a connector to merge them
+    let root_id = app.world().resource::<R>().root;
+    app.world_mut()
+        .commands()
+        .grid(root_id, Grid::default())
+        .spawn_spatial(CellCoord::new(1, 0, 0));
+
+    let mut left_changed = false;
+    let mut right_changed = false;
+    let mut left_pid_entities = None;
+    let mut right_pid_entities = None;
+
+    for _ in 0..100 {
+        run_app_once(&mut app);
+        let ep = app.world().resource::<PartitionEntities>();
+        if ep.changed.contains_key(&left) {
+            left_changed = true;
+        }
+        if ep.changed.contains_key(&right) {
+            right_changed = true;
+        }
+        left_pid_entities = ep.map.get(&left).copied();
+        right_pid_entities = ep.map.get(&right).copied();
+    }
+
+    // After merge, both should be in the same partition
+    assert_eq!(
+        left_pid_entities, right_pid_entities,
+        "Partitions should have merged in PartitionEntities"
+    );
+
+    // At least one of them MUST have changed its partition ID to match the other
+    assert!(
+        left_changed || right_changed,
+        "At least one entity should have recorded a partition change during merge"
+    );
+}