Add the AssetChanged query filter

crates/bevy_asset/src/asset_changed.rs

@@ -0,0 +1,398 @@
+//! Filter query to limit iteration over [`Handle`]s which content was updated recently.
+use std::marker::PhantomData;
+
+use bevy_ecs::{
+    archetype::{Archetype, ArchetypeComponentId},
+    component::{ComponentId, Tick},
+    prelude::{Changed, Entity, Or, World},
+    query::{Access, FilteredAccess, QueryItem, ReadFetch, ReadOnlyWorldQuery, WorldQuery},
+    storage::{Table, TableRow},
+    world::unsafe_world_cell::UnsafeWorldCell,
+};
+use bevy_utils::HashMap;
+
+use crate::{asset_changes::AssetChanges, Asset, Handle, HandleId};
+
+#[derive(Clone, Copy)]
+struct AssetChangeCheck<'w> {
+    handle_change_map: &'w HashMap<HandleId, Tick>,
+    last_run: Tick,
+    this_run: Tick,
+}
+impl<'w> AssetChangeCheck<'w> {
+    fn new<T: Asset>(changes: &'w AssetChanges<T>, last_run: Tick, this_run: Tick) -> Self {
+        Self {
+            handle_change_map: &changes.change_ticks,
+            last_run,
+            this_run,
+        }
+    }
+    // TODO(perf): some sort of caching? Each check has two levels of indirection,
+    // which is not optimal.
+    fn has_changed<T: Asset>(&self, handle: &Handle<T>) -> bool {
+        if let Some(&handle_change_tick) = self.handle_change_map.get(&handle.id()) {
+            handle_change_tick.is_newer_than(self.last_run, self.this_run)
+        } else {
+            false
+        }
+    }
+}
+
+/// A shortcut for the commonly used `Or<(Changed<Handle<T>>, AssetChanged<T>)>`
+/// query.
+///
+/// If you want to react to changes to `Handle<T>` component of an entity, you have
+/// two cases to worry about:
+/// - The `Handle<T>` was changed through a `Query<&mut Handle<T>>`.
+/// - The `T` in `Assets<T>` pointed by `Handle<T>` was changed through a
+///   `ResMut<Assets<T>>::get_mut`.
+///
+/// To properly handle both those cases, you need to combine the `Changed` and `AssetChanged`
+/// filters. This query type does it for you.
+pub type AssetOrHandleChanged<T> = Or<(Changed<Handle<T>>, AssetChanged<T>)>;
+
+/// Query filter to get all entities which `Handle<T>` component underlying asset changed.
+///
+/// Unlike `Changed<Handle<T>>`, this filter returns entities for which **the underlying `T`
+/// asset** was changed. For example, when modifying a `Mesh`, you would do so through the
+/// [`Assets<Mesh>::get_mut`] method. This does not change the `Handle<T>` that entities have
+/// as component, so `Changed<Handle<T>>` **will not do anything**. While with `AssetChanged<T>`,
+/// all entities with the modified mesh will be queried.
+///
+/// # Notes
+/// **performance**: This reads a hashmap once per entity with a `Handle<T>` component, this might
+/// result in slowdown, particularly if there are a lot of such entities. Note that if there was no
+/// updates, it will simply skip everything.
+///
+/// **frame delay**: The list of changed assets only gets updated during the
+/// [`AssetStage::AssetEvents`] stage. Which means asset changes are not immediately visible to
+/// this query. Therefore, asset change detection will have one frame of delay.
+///
+/// Since oftentimes, it is necessary to check for both handle and asset changes, the
+/// [`AssetOrHandleChanged`] query is provided for convenience.
+///
+/// [`AssetStage::AssetEvents`]: crate::AssetStage::AssetEvents
+/// [`Assets<Mesh>::get_mut`]: crate::Assets::get_mut
+pub struct AssetChanged<T>(PhantomData<T>);
+
+/// Fetch for [`AssetChanged`].
+#[doc(hidden)]
+pub struct AssetChangedFetch<'w, T: Asset> {
+    inner: Option<ReadFetch<'w, Handle<T>>>,
+    check: AssetChangeCheck<'w>,
+}
+impl<'w, T: Asset> Clone for AssetChangedFetch<'w, T> {
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner,
+            check: self.check,
+        }
+    }
+}
+
+/// State for [`AssetChanged`].
+#[doc(hidden)]
+pub struct AssetChangedState<T: Asset> {
+    handle_id: ComponentId,
+    asset_changes_id: ComponentId,
+    asset_changes_archetype_id: ArchetypeComponentId,
+    _asset: PhantomData<fn(T)>,
+}
+
+// SAFETY: `ROQueryFetch<Self>` is the same as `QueryFetch<Self>`
+unsafe impl<T: Asset> WorldQuery for AssetChanged<T> {
+    type ReadOnly = Self;
+    type Fetch<'w> = AssetChangedFetch<'w, T>;
+    type State = AssetChangedState<T>;
+
+    type Item<'w> = bool;
+
+    fn shrink<'wlong: 'wshort, 'wshort>(item: QueryItem<'wlong, Self>) -> QueryItem<'wshort, Self> {
+        item
+    }
+
+    fn init_state(world: &mut World) -> AssetChangedState<T> {
+        let asset_changes_resource_id =
+            if let Some(id) = world.components().resource_id::<AssetChanges<T>>() {
+                id
+            } else {
+                world.init_resource::<AssetChanges<T>>()
+            };
+        let asset_changes_archetype_id = world
+            .storages()
+            .resources
+            .get(asset_changes_resource_id)
+            .unwrap()
+            .id();
+        AssetChangedState {
+            handle_id: world.init_component::<Handle<T>>(),
+            asset_changes_id: asset_changes_resource_id,
+            asset_changes_archetype_id,
+            _asset: PhantomData,
+        }
+    }
+
+    fn matches_component_set(
+        state: &Self::State,
+        set_contains_id: &impl Fn(ComponentId) -> bool,
+    ) -> bool {
+        set_contains_id(state.handle_id)
+    }
+    unsafe fn init_fetch<'w>(
+        world: UnsafeWorldCell<'w>,
+        state: &Self::State,
+        last_run: Tick,
+        this_run: Tick,
+    ) -> Self::Fetch<'w> {
+        let err_msg = || {
+            panic!(
+            "AssetChanges<{ty}> was removed, please do not remove AssetChanges<{ty}> when using AssetChanged<{ty}>",
+            ty = std::any::type_name::<T>()
+        )
+        };
+        let changes: &AssetChanges<_> = world.get_resource().unwrap_or_else(err_msg);
+        let has_updates = changes.last_change_tick.is_newer_than(last_run, this_run);
+        AssetChangedFetch {
+            inner: has_updates
+                .then(|| <&_>::init_fetch(world, &state.handle_id, last_run, this_run)),
+            check: AssetChangeCheck::new::<T>(changes, last_run, this_run),
+        }
+    }
+
+    const IS_DENSE: bool = <&Handle<T>>::IS_DENSE;
+
+    // This Fetch filters more than just the archetype.
+    const IS_ARCHETYPAL: bool = false;
+
+    unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, table: &'w Table) {
+        if let Some(inner) = &mut fetch.inner {
+            <&Handle<T>>::set_table(inner, &state.handle_id, table);
+        }
+    }
+
+    unsafe fn set_archetype<'w>(
+        fetch: &mut Self::Fetch<'w>,
+        state: &Self::State,
+        archetype: &'w Archetype,
+        table: &'w Table,
+    ) {
+        if let Some(inner) = &mut fetch.inner {
+            <&Handle<T>>::set_archetype(inner, &state.handle_id, archetype, table);
+        }
+    }
+
+    unsafe fn fetch<'w>(
+        fetch: &mut Self::Fetch<'w>,
+        entity: Entity,
+        table_row: TableRow,
+    ) -> Self::Item<'w> {
+        fetch.inner.as_mut().map_or(false, |inner| {
+            let handle = <&Handle<T>>::fetch(inner, entity, table_row);
+            fetch.check.has_changed(handle)
+        })
+    }
+
+    #[inline]
+    unsafe fn filter_fetch(
+        fetch: &mut Self::Fetch<'_>,
+        entity: Entity,
+        table_row: TableRow,
+    ) -> bool {
+        Self::fetch(fetch, entity, table_row)
+    }
+
+    #[inline]
+    fn update_component_access(state: &Self::State, access: &mut FilteredAccess<ComponentId>) {
+        <&Handle<T>>::update_component_access(&state.handle_id, access);
+        access.add_read(state.asset_changes_id);
+    }
+
+    #[inline]
+    fn update_archetype_component_access(
+        state: &Self::State,
+        archetype: &Archetype,
+        access: &mut Access<ArchetypeComponentId>,
+    ) {
+        access.add_read(state.asset_changes_archetype_id);
+        <&Handle<T>>::update_archetype_component_access(&state.handle_id, archetype, access);
+    }
+}
+
+/// SAFETY: read-only access
+unsafe impl<T: Asset> ReadOnlyWorldQuery for AssetChanged<T> {}
+
+#[cfg(test)]
+mod tests {
+    use crate::{AddAsset, Assets};
+    use bevy_app::{App, AppExit, Startup, Update};
+    use bevy_ecs::{
+        entity::Entity,
+        event::EventWriter,
+        schedule::IntoSystemConfigs,
+        system::{Commands, IntoSystem, Local, Query, Res, ResMut, Resource},
+    };
+    use bevy_reflect::TypePath;
+
+    use super::*;
+
+    #[derive(bevy_reflect::TypeUuid, TypePath)]
+    #[uuid = "44115972-f31b-46e5-be5c-2b9aece6a52f"]
+    struct MyAsset(u32, &'static str);
+
+    fn run_app<Marker>(system: impl IntoSystem<(), (), Marker>) {
+        let mut app = App::new();
+        app.add_plugins((bevy_core::FrameCountPlugin, crate::AssetPlugin::default()))
+            .add_asset::<MyAsset>()
+            .add_systems(Update, system);
+        app.run();
+    }
+
+    #[test]
+    #[should_panic]
+    fn handle_filter_with_other_query_panics() {
+        fn compatible_filter(
+            _query1: Query<&mut Handle<MyAsset>>,
+            _query2: Query<Entity, AssetChanged<MyAsset>>,
+        ) {
+            println!(
+                "this line should not run, as &mut Handle<MyAsset> conflicts with AssetChanged<MyAsset>",
+            );
+        }
+        run_app(compatible_filter);
+    }
+
+    #[test]
+    #[should_panic]
+    fn handle_query_pos_panics() {
+        fn incompatible_query(_query: Query<(AssetChanged<MyAsset>, &mut Handle<MyAsset>)>) {
+            println!(
+                "this line should not run, as &mut Handle<MyAsset> conflicts with AssetChanged<MyAsset>",
+            );
+        }
+        run_app(incompatible_query);
+    }
+
+    #[test]
+    #[should_panic]
+    fn handle_other_query_panics() {
+        fn incompatible_params(
+            _query1: Query<&mut Handle<MyAsset>>,
+            _query2: Query<AssetChanged<MyAsset>>,
+        ) {
+            println!(
+                "this line should not run, as AssetChanged<MyAsset> conflicts with &mut Handle<MyAsset>",
+            );
+        }
+        run_app(incompatible_params);
+    }
+
+    // According to a comment in QueryState::new in bevy_ecs, components on filter
+    // position shouldn't conflict with components on query position.
+    #[test]
+    fn handle_filter_pos_ok() {
+        fn compatible_filter(
+            _query: Query<&mut Handle<MyAsset>, AssetChanged<MyAsset>>,
+            mut exit: EventWriter<AppExit>,
+        ) {
+            exit.send(AppExit);
+        }
+        run_app(compatible_filter);
+    }
+
+    #[test]
+    fn asset_change() {
+        #[derive(Default, Resource)]
+        struct Counter(u32, u32);
+
+        let mut app = App::new();
+
+        fn count_update(
+            mut counter: ResMut<Counter>,
+            assets: Res<Assets<MyAsset>>,
+            query: Query<&Handle<MyAsset>, AssetChanged<MyAsset>>,
+        ) {
+            for handle in query.iter() {
+                let asset = assets.get(handle).unwrap();
+                let to_update = match asset.0 {
+                    0 => &mut counter.0,
+                    1 => &mut counter.1,
+                    _ => continue,
+                };
+                *to_update += 1;
+            }
+        }
+
+        fn update_once(mut assets: ResMut<Assets<MyAsset>>, mut run_count: Local<u32>) {
+            let mut update_index = |i| {
+                let handle = assets
+                    .iter()
+                    .find_map(|(h, a)| (a.0 == i).then_some(h))
+                    .unwrap();
+                let handle = assets.get_handle(handle);
+                let asset = assets.get_mut(&handle).unwrap();
+                asset.1 = "new_value";
+            };
+            match *run_count {
+                2 => update_index(0),
+                3.. => update_index(1),
+                _ => {}
+            };
+            *run_count += 1;
+        }
+        app.add_plugins((bevy_core::FrameCountPlugin, crate::AssetPlugin::default()))
+            .add_asset::<MyAsset>()
+            .init_resource::<Counter>()
+            .add_systems(
+                Startup,
+                |mut cmds: Commands, mut assets: ResMut<Assets<MyAsset>>| {
+                    let asset0 = assets.add(MyAsset(0, "init"));
+                    let asset1 = assets.add(MyAsset(1, "init"));
+                    cmds.spawn(asset0.clone());
+                    cmds.spawn(asset0);
+                    cmds.spawn(asset1.clone());
+                    cmds.spawn(asset1.clone());
+                    cmds.spawn(asset1);
+                },
+            )
+            .add_systems(Update, (count_update, update_once).chain());
+        let assert_counter = |app: &App, assert: Counter| {
+            let counter = app.world.get_resource::<Counter>().unwrap();
+            assert_eq!(counter.0, assert.0);
+            assert_eq!(counter.1, assert.1);
+        };
+        // First run of the app, `add_systems(Startup…)` runs.
+        app.update(); // run_count == 0
+
+        // This might be considered a bug, since it asserts a frame delay. So if you happen
+        // to break test on the following line, consider removing it.
+        assert_counter(&app, Counter(0, 0));
+
+        // Second run, the asset events are emitted in asset_event_system, and the
+        // AssetChanges res is updated as well.
+        // This means, our asset_change_counter are updated.
+        app.update(); // run_count == 1
+        assert_counter(&app, Counter(2, 3));
+
+        // Third run: `run_count` in `update_once` enables it to update
+        // the `MyAsset` with value 0. However, since last frame no assets
+        // were updated, asset_change_counter stays the same
+        app.update(); // run_count == 2
+        assert_counter(&app, Counter(2, 3));
+
+        // Now that `update_once` and asset_event_system ran, the `MyAsset`
+        // with value 1 got updated. There are two handles linking to it,
+        // so that does 4 updates total.
+        // This cycle, update_once updates the assets of index 1, therefore,
+        // next cycle, we will detect them
+        app.update(); // run_count == 3
+        assert_counter(&app, Counter(4, 3));
+
+        // Last cycle we updated the asset with 3 associated entities, so
+        // we got 3 new changes to 1
+        app.update(); // run_count == 4
+        assert_counter(&app, Counter(4, 6));
+        // ibid
+        app.update(); // run_count == 5
+        assert_counter(&app, Counter(4, 9));
+    }
+}