From 523f599fc306d1fe9b04b4760557bfe282b9592f Mon Sep 17 00:00:00 2001
From: MiniaczQ <xnetroidpl@gmail.com>
Date: Sat, 28 Sep 2024 04:47:15 +0200
Subject: [PATCH] init

---
 crates/bevy_ecs/src/system/query.rs        | 29 +++++++++++
 crates/bevy_ecs/src/system/system_param.rs | 59 ++++++++++++++++++++++
 examples/ecs/fallible_params.rs            |  9 ++--
 3 files changed, 93 insertions(+), 4 deletions(-)

diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs
index 13155057ff91b8..e74b21dc4a7603 100644
--- a/crates/bevy_ecs/src/system/query.rs
+++ b/crates/bevy_ecs/src/system/query.rs
@@ -1667,3 +1667,32 @@ impl<'w, D: QueryData, F: QueryFilter> QuerySingle<'w, D, F> {
         self.item
     }
 }
+
+/// [System parameter] that works very much like [`Query`] except it always contains at least one matching entity.
+///
+/// This [`SystemParam`](crate::system::SystemParam) fails validation if no matching entities exist.
+/// This will cause systems that use this parameter to be skipped.
+///
+/// See [`Query`] for more details.
+pub struct QueryNonEmpty<'w, 's, D: QueryData, F: QueryFilter = ()>(pub(crate) Query<'w, 's, D, F>);
+
+impl<'w, 's, D: QueryData, F: QueryFilter> Deref for QueryNonEmpty<'w, 's, D, F> {
+    type Target = Query<'w, 's, D, F>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<D: QueryData, F: QueryFilter> DerefMut for QueryNonEmpty<'_, '_, D, F> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl<'w, 's, D: QueryData, F: QueryFilter> QueryNonEmpty<'w, 's, D, F> {
+    /// Returns the inner item with ownership.
+    pub fn into_inner(self) -> Query<'w, 's, D, F> {
+        self.0
+    }
+}
diff --git a/crates/bevy_ecs/src/system/system_param.rs b/crates/bevy_ecs/src/system/system_param.rs
index 8dc4edc43f08e7..43fb9e2d3b2580 100644
--- a/crates/bevy_ecs/src/system/system_param.rs
+++ b/crates/bevy_ecs/src/system/system_param.rs
@@ -26,6 +26,8 @@ use core::{
     ops::{Deref, DerefMut},
 };
 
+use super::QueryNonEmpty;
+
 /// A parameter that can be used in a [`System`](super::System).
 ///
 /// # Derive
@@ -489,6 +491,63 @@ unsafe impl<'a, D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> ReadOn
 {
 }
 
+// SAFETY: Relevant query ComponentId and ArchetypeComponentId access is applied to SystemMeta. If
+// this Query conflicts with any prior access, a panic will occur.
+unsafe impl<'world, 'state, D: QueryData + 'static, F: QueryFilter + 'static> SystemParam
+    for QueryNonEmpty<'world, 'state, D, F>
+{
+    type State = QueryState<D, F>;
+    type Item<'w, 's> = QueryNonEmpty<'w, 's, D, F>;
+
+    fn init_state(world: &mut World, system_meta: &mut SystemMeta) -> Self::State {
+        Query::<'world, 'state, D, F>::init_state(world, system_meta)
+    }
+
+    unsafe fn new_archetype(
+        state: &mut Self::State,
+        archetype: &Archetype,
+        system_meta: &mut SystemMeta,
+    ) {
+        // SAFETY: Delegate to existing `SystemParam` implementations.
+        unsafe { Query::<'world, 'state, D, F>::new_archetype(state, archetype, system_meta) };
+    }
+
+    #[inline]
+    unsafe fn get_param<'w, 's>(
+        state: &'s mut Self::State,
+        system_meta: &SystemMeta,
+        world: UnsafeWorldCell<'w>,
+        change_tick: Tick,
+    ) -> Self::Item<'w, 's> {
+        // SAFETY: Delegate to existing `SystemParam` implementations.
+        let query = unsafe {
+            Query::<'world, 'state, D, F>::get_param(state, system_meta, world, change_tick)
+        };
+        QueryNonEmpty(query)
+    }
+
+    #[inline]
+    unsafe fn validate_param(
+        state: &Self::State,
+        system_meta: &SystemMeta,
+        world: UnsafeWorldCell,
+    ) -> bool {
+        state.validate_world(world.id());
+        // SAFETY:
+        // - We have read-only access to the components accessed by query.
+        // - The world has been validated.
+        !unsafe {
+            state.is_empty_unsafe_world_cell(world, system_meta.last_run, world.change_tick())
+        }
+    }
+}
+
+// SAFETY: QueryState is constrained to read-only fetches, so it only reads World.
+unsafe impl<'w, 's, D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> ReadOnlySystemParam
+    for QueryNonEmpty<'w, 's, D, F>
+{
+}
+
 /// A collection of potentially conflicting [`SystemParam`]s allowed by disjoint access.
 ///
 /// Allows systems to safely access and interact with up to 8 mutually exclusive [`SystemParam`]s, such as
diff --git a/examples/ecs/fallible_params.rs b/examples/ecs/fallible_params.rs
index 6965fbcf1d321f..4671834b681fa7 100644
--- a/examples/ecs/fallible_params.rs
+++ b/examples/ecs/fallible_params.rs
@@ -5,8 +5,9 @@
 //! - [`Res<R>`], [`ResMut<R>`] - If resource doesn't exist.
 //! - [`QuerySingle<D, F>`] - If there is no or more than one entities matching.
 //! - [`Option<QuerySingle<D, F>>`] - If there are more than one entities matching.
+//! - [`QueryNonEmpty<D, F>`] - If there are no matching entities matching.
 
-use bevy::prelude::*;
+use bevy::{ecs::system::QueryNonEmpty, prelude::*};
 use rand::Rng;
 
 fn main() {
@@ -105,9 +106,9 @@ fn user_input(
 }
 
 // System that moves the enemies in a circle.
-// TODO: Use [`NonEmptyQuery`] when it exists.
-fn move_targets(mut enemies: Query<(&mut Transform, &mut Enemy)>, time: Res<Time>) {
-    for (mut transform, mut target) in &mut enemies {
+// Only runs if there are enemies.
+fn move_targets(mut enemies: QueryNonEmpty<(&mut Transform, &mut Enemy)>, time: Res<Time>) {
+    for (mut transform, mut target) in &mut *enemies {
         target.rotation += target.rotation_speed * time.delta_seconds();
         transform.rotation = Quat::from_rotation_z(target.rotation);
         let offset = transform.right() * target.radius;