[Merged by Bors] - Query::get_unique

crates/bevy_ecs/src/system/query/extensions.rs

@@ -0,0 +1,63 @@
+use std::fmt::Debug;
+
+use thiserror::Error;
+
+use crate::{Fetch, Query, QueryFilter, ReadOnlyFetch, WorldQuery};
+
+impl<'a, Q: WorldQuery, F: QueryFilter> Query<'a, Q, F> {
+    /// Takes exactly one result from the query. If there are no results, or more than 1 result, this will return an error instead.
+    pub fn get_unique(&self) -> Result<<Q::Fetch as Fetch<'_>>::Item, OnlyQueryError<'_, Q>>
+    where
+        Q::Fetch: ReadOnlyFetch,
+    {
+        let mut query = self.iter();
+        let first = query.next();
+        let extra_count = query.count();
+
+        match (first, extra_count) {
+            (Some(r), 0) => Ok(r),
+            (None, _) => Err(OnlyQueryError::NoEntities(std::any::type_name::<Self>())),
+            (Some(r), extra) => Err(OnlyQueryError::MultipleEntities(
+                r,
+                extra,
+                std::any::type_name::<Self>(),
+            )),
+        }
+    }
+
+    /// See [`Query::get_unique`]
+    pub fn get_unique_mut(
+        &mut self,
+    ) -> Result<<Q::Fetch as Fetch<'_>>::Item, OnlyQueryError<'_, Q>> {
+        let mut query = self.iter_mut();
+        let first = query.next();
+        let extra_count = query.count();
+
+        match (first, extra_count) {
+            (Some(r), 0) => Ok(r),
+            (None, _) => Err(OnlyQueryError::NoEntities(std::any::type_name::<Self>())),
+            (Some(r), extra) => Err(OnlyQueryError::MultipleEntities(
+                r,
+                extra,
+                std::any::type_name::<Self>(),
+            )),
+        }
+    }
+}
+
+#[derive(Error)]
+pub enum OnlyQueryError<'a, Q: WorldQuery> {
+    #[error("No entities fit the query {0}")]
+    NoEntities(&'static str),
+    #[error("Multiple entities ({1} extra) fit the query {2}!")]
+    MultipleEntities(<Q::Fetch as Fetch<'a>>::Item, usize, &'static str),
+}
+
+impl<'a, Q: WorldQuery> Debug for OnlyQueryError<'a, Q> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            OnlyQueryError::NoEntities(_) => f.debug_tuple("NoEntities").finish(),
+            OnlyQueryError::MultipleEntities(_, _, _) => f.debug_tuple("MultipleEntities").finish(),
+        }
+    }
+}

crates/bevy_ecs/src/system/query/mod.rs

@@ -1,3 +1,4 @@
+mod extensions;
 mod query_set;
 pub use query_set::*;
 

examples/game/breakout.rs

@@ -162,7 +162,7 @@ fn paddle_movement_system(
     keyboard_input: Res<Input<KeyCode>>,
     mut query: Query<(&Paddle, &mut Transform)>,
 ) {
-    for (paddle, mut transform) in query.iter_mut() {
+    if let Ok((paddle, mut transform)) = query.get_unique_mut() {
         let mut direction = 0.0;
         if keyboard_input.pressed(KeyCode::Left) {
             direction -= 1.0;
@@ -184,15 +184,13 @@ fn ball_movement_system(time: Res<Time>, mut ball_query: Query<(&Ball, &mut Tran
     // clamp the timestep to stop the ball from escaping when the game starts
     let delta_seconds = f32::min(0.2, time.delta_seconds());
 
-    for (ball, mut transform) in ball_query.iter_mut() {
+    if let Ok((ball, mut transform)) = ball_query.get_unique_mut() {
         transform.translation += ball.velocity * delta_seconds;
     }
 }
 
 fn scoreboard_system(scoreboard: Res<Scoreboard>, mut query: Query<&mut Text>) {
-    for mut text in query.iter_mut() {
-        text.value = format!("Score: {}", scoreboard.score);
-    }
+    query.get_unique_mut().unwrap().value = format!("Score: {}", scoreboard.score)
 }
 
 fn ball_collision_system(
@@ -201,7 +199,7 @@ fn ball_collision_system(
     mut ball_query: Query<(&mut Ball, &Transform, &Sprite)>,
     collider_query: Query<(Entity, &Collider, &Transform, &Sprite)>,
 ) {
-    for (mut ball, ball_transform, sprite) in ball_query.iter_mut() {
+    if let Ok((mut ball, ball_transform, sprite)) = ball_query.get_unique_mut() {
         let ball_size = sprite.size;
         let velocity = &mut ball.velocity;