Recalculate Aabbs when meshes change

mlodato517 · alice-i-cecile · IceSentry · mlodato517 · commit 9b36e77c8ab5 · 2022-07-18T12:30:19.000-04:00
**This Commit** Adds an `update_bounds` system that ensures entities have updated `Aabb`s after they're assigned new mesh handles or if their mesh is directly mutated. **Why?** The `calculate_bounds` system is responsible for taking entities with meshes, calculating their Aabbs, and attaching it to the entities. If an entity is assigned a new mesh, however, nothing updates the entity's Aabb to match the new mesh. Additionally, if an entity's mesh is mutated to be a different shape, the attached Aabb isn't updated. Fixes #4294. **Notes** - While it doesn't seem possible now there remains the possibility that systems are updated such that an entity is added to its mesh's list in the mapping twice. This could result in some duplicate Aabb assignment but is considered worth the risk since it should be very rare. - Updating the Aabb for entities assigned to a mesh could introduce performance problems if meshes are frequently mutated or there are lots of entities assigned with a mesh that's mutated. To help with this, there is a new component `NoAabbUpdate` that can be assigned to entities so that their Aabbs are not updated as their mesh changes. - We may decide to have mesh-update-events contain the list of entities affected in the future. That may simplify these systems. See [this comment][0] for more details. [0]: #4944 (comment) Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com> Co-authored-by: Charles <c.giguere42@gmail.com> Co-authored-by: François <mockersf@gmail.com> Co-authored-by: Nicola Papale <nico@nicopap.ch> Co-authored-by: Robert Swain <robert.swain@gmail.com>
diff --git a/crates/bevy_render/src/view/visibility/mod.rs b/crates/bevy_render/src/view/visibility/mod.rs
@@ -1,15 +1,18 @@
 mod render_layers;
 
+use smallvec::SmallVec;
+
 pub use render_layers::*;
 
 use bevy_app::{CoreStage, Plugin};
-use bevy_asset::{Assets, Handle};
+use bevy_asset::{AssetEvent, Assets, Handle};
 use bevy_ecs::prelude::*;
 use bevy_hierarchy::{Children, Parent};
 use bevy_reflect::std_traits::ReflectDefault;
 use bevy_reflect::Reflect;
 use bevy_transform::components::GlobalTransform;
 use bevy_transform::TransformSystem;
+use bevy_utils::HashMap;
 use std::cell::Cell;
 use thread_local::ThreadLocal;
 
@@ -106,6 +109,11 @@ pub struct VisibilityBundle {
 #[derive(Component)]
 pub struct NoFrustumCulling;
 
+/// Use this component to opt-out of built-in [`Aabb`] updating for Mesh entities (i.e. to opt out
+/// of [`update_bounds`]).
+#[derive(Component)]
+pub struct NoAabbUpdate;
+
 /// Collection of entities visible from the current view.
 ///
 /// This component contains all entities which are visible from the currently
@@ -139,9 +147,59 @@ impl VisibleEntities {
     }
 }
 
+/// Tracks which [`Entities`](Entity) have which meshes for entities whose [`Aabb`]s are managed by
+/// the [`calculate_bounds`] and [`update_bounds`] systems. This is needed because `update_bounds`
+/// recomputes `Aabb`s for entities whose mesh has been mutated. These mutations are visible via
+/// [`AssetEvent<Mesh>`](AssetEvent) which tells us which mesh was changed but not which entities
+/// have that mesh.
+#[derive(Debug, Default, Clone)]
+pub struct EntityMeshMap {
+    entities_with_mesh: HashMap<Handle<Mesh>, SmallVec<[Entity; 1]>>,
+    mesh_for_entity: HashMap<Entity, Handle<Mesh>>,
+}
+
+impl EntityMeshMap {
+    /// Register the passed `entity` as having the passed `mesh_handle`.
+    fn register(&mut self, entity: Entity, mesh_handle: &Handle<Mesh>) {
+        // Note that this list can have duplicates if an entity is registered for a mesh multiple
+        // times. This should be rare and only cause an additional `Aabb.clone()` in
+        // `update_bounds` so it is preferable to a `HashSet` for now.
+        self.entities_with_mesh
+            .entry(mesh_handle.clone_weak())
+            .or_default()
+            .push(entity);
+        self.mesh_for_entity
+            .insert(entity, mesh_handle.clone_weak());
+    }
+
+    /// Deregisters the mapping between an `Entity` and `Mesh`. Used so [`update_bounds`] can
+    /// track which mappings are still active so `Aabb`s are updated correctly.
+    fn deregister(&mut self, entity: Entity) {
+        let mut inner = || {
+            let mesh = self.mesh_for_entity.remove(&entity)?;
+
+            // This lookup failing is _probably_ an error.
+            let entities = self.entities_with_mesh.get_mut(&mesh)?;
+
+            // There could be duplicate entries in here if an entity was registered with a mesh
+            // multiple times. It's important to remove all references so that if an entity gets a
+            // new mesh and its old mesh is mutated, the entity doesn't get its old mesh's new
+            // `Aabb`. Note that there _should_ only be one entity.
+            for i in (0..entities.len()).rev() {
+                if entities[i] == entity {
+                    entities.swap_remove(i);
+                }
+            }
+            Some(())
+        };
+        inner();
+    }
+}
+
 #[derive(Debug, Hash, PartialEq, Eq, Clone, SystemLabel)]
 pub enum VisibilitySystems {
     CalculateBounds,
+    UpdateBounds,
     UpdateOrthographicFrusta,
     UpdatePerspectiveFrusta,
     UpdateProjectionFrusta,
@@ -157,60 +215,121 @@ impl Plugin for VisibilityPlugin {
     fn build(&self, app: &mut bevy_app::App) {
         use VisibilitySystems::*;
 
-        app.add_system_to_stage(
-            CoreStage::PostUpdate,
-            calculate_bounds.label(CalculateBounds),
-        )
-        .add_system_to_stage(
-            CoreStage::PostUpdate,
-            update_frusta::<OrthographicProjection>
-                .label(UpdateOrthographicFrusta)
-                .after(TransformSystem::TransformPropagate),
-        )
-        .add_system_to_stage(
-            CoreStage::PostUpdate,
-            update_frusta::<PerspectiveProjection>
-                .label(UpdatePerspectiveFrusta)
-                .after(TransformSystem::TransformPropagate),
-        )
-        .add_system_to_stage(
-            CoreStage::PostUpdate,
-            update_frusta::<Projection>
-                .label(UpdateProjectionFrusta)
-                .after(TransformSystem::TransformPropagate),
-        )
-        .add_system_to_stage(
-            CoreStage::PostUpdate,
-            visibility_propagate_system.label(VisibilityPropagate),
-        )
-        .add_system_to_stage(
-            CoreStage::PostUpdate,
-            check_visibility
-                .label(CheckVisibility)
-                .after(CalculateBounds)
-                .after(UpdateOrthographicFrusta)
-                .after(UpdatePerspectiveFrusta)
-                .after(UpdateProjectionFrusta)
-                .after(VisibilityPropagate)
-                .after(TransformSystem::TransformPropagate),
-        );
+        app.init_resource::<EntityMeshMap>()
+            .add_system_to_stage(
+                CoreStage::PostUpdate,
+                calculate_bounds.label(CalculateBounds),
+            )
+            .add_system_to_stage(CoreStage::PostUpdate, update_bounds.label(UpdateBounds))
+            .add_system_to_stage(
+                CoreStage::PostUpdate,
+                update_frusta::<OrthographicProjection>
+                    .label(UpdateOrthographicFrusta)
+                    .after(TransformSystem::TransformPropagate),
+            )
+            .add_system_to_stage(
+                CoreStage::PostUpdate,
+                update_frusta::<PerspectiveProjection>
+                    .label(UpdatePerspectiveFrusta)
+                    .after(TransformSystem::TransformPropagate),
+            )
+            .add_system_to_stage(
+                CoreStage::PostUpdate,
+                update_frusta::<Projection>
+                    .label(UpdateProjectionFrusta)
+                    .after(TransformSystem::TransformPropagate),
+            )
+            .add_system_to_stage(
+                CoreStage::PostUpdate,
+                visibility_propagate_system.label(VisibilityPropagate),
+            )
+            .add_system_to_stage(
+                CoreStage::PostUpdate,
+                check_visibility
+                    .label(CheckVisibility)
+                    .after(CalculateBounds)
+                    .after(UpdateBounds)
+                    .after(UpdateOrthographicFrusta)
+                    .after(UpdatePerspectiveFrusta)
+                    .after(UpdateProjectionFrusta)
+                    .after(VisibilityPropagate)
+                    .after(TransformSystem::TransformPropagate),
+            );
     }
 }
 
+/// Calculates [`Aabb`]s for [`Entities`](Entity) with [`Mesh`]es.
 pub fn calculate_bounds(
     mut commands: Commands,
     meshes: Res<Assets<Mesh>>,
     without_aabb: Query<(Entity, &Handle<Mesh>), (Without<Aabb>, Without<NoFrustumCulling>)>,
+    mut entity_mesh_map: ResMut<EntityMeshMap>,
 ) {
     for (entity, mesh_handle) in &without_aabb {
         if let Some(mesh) = meshes.get(mesh_handle) {
             if let Some(aabb) = mesh.compute_aabb() {
+                entity_mesh_map.register(entity, mesh_handle);
                 commands.entity(entity).insert(aabb);
             }
         }
     }
 }
 
+/// Updates [`Aabb`]s for [`Entities`](Entity) with [`Mesh`]es. This includes `Entities` that have
+/// been assigned new `Mesh`es as well as `Entities` whose `Mesh` has been directly mutated.
+///
+/// To opt out of bound calculation for an `Entity`, give it the [`NoAabbUpdate`] component.
+///
+/// NOTE: This system needs to remove entities from their collection in
+/// [`EntityMeshMap`] whenever a mesh handle is reassigned or an entity's mesh handle is
+/// removed. This may impact performance if meshes with many entities are frequently
+/// reassigned/removed.
+pub fn update_bounds(
+    mut commands: Commands,
+    meshes: Res<Assets<Mesh>>,
+    mut mesh_reassigned: Query<
+        (Entity, &Handle<Mesh>, &mut Aabb),
+        (
+            Changed<Handle<Mesh>>,
+            Without<NoFrustumCulling>,
+            Without<NoAabbUpdate>,
+        ),
+    >,
+    mut entity_mesh_map: ResMut<EntityMeshMap>,
+    mut mesh_events: EventReader<AssetEvent<Mesh>>,
+    entities_lost_mesh: RemovedComponents<Handle<Mesh>>,
+) {
+    for entity in entities_lost_mesh.iter() {
+        entity_mesh_map.deregister(entity);
+    }
+
+    for (entity, mesh_handle, mut aabb) in mesh_reassigned.iter_mut() {
+        entity_mesh_map.deregister(entity);
+        if let Some(mesh) = meshes.get(mesh_handle) {
+            if let Some(new_aabb) = mesh.compute_aabb() {
+                entity_mesh_map.register(entity, mesh_handle);
+                *aabb = new_aabb;
+            }
+        }
+    }
+
+    let to_update = |event: &AssetEvent<Mesh>| {
+        let handle = match event {
+            AssetEvent::Modified { handle } => handle,
+            _ => return None,
+        };
+        let mesh = meshes.get(handle)?;
+        let entities_with_handle = entity_mesh_map.entities_with_mesh.get(handle)?;
+        let aabb = mesh.compute_aabb()?;
+        Some((aabb, entities_with_handle))
+    };
+    for (aabb, entities_with_handle) in mesh_events.iter().filter_map(to_update) {
+        for entity in entities_with_handle {
+            commands.entity(*entity).insert(aabb.clone());
+        }
+    }
+}
+
 pub fn update_frusta<T: Component + CameraProjection + Send + Sync + 'static>(
     mut views: Query<(&GlobalTransform, &T, &mut Frustum)>,
 ) {
