[Merged by Bors] - Query::get_unique

crates/bevy_ecs/src/system/query.rs

@@ -7,7 +7,7 @@ use crate::{
     world::{Mut, World},
 };
 use bevy_tasks::TaskPool;
-use std::any::TypeId;
+use std::{any::TypeId, fmt::Debug};
 use thiserror::Error;
 
 /// Provides scoped access to a World according to a given [WorldQuery] and query filter
@@ -212,6 +212,42 @@ where
             Err(QueryComponentError::MissingWriteAccess)
         }
     }
+
+    pub fn get_unique(&self) -> Result<<Q::Fetch as Fetch<'_>>::Item, UniqueQueryError<'_, Q>>
+    where
+        Q::Fetch: ReadOnlyFetch,
+    {
+        let mut query = self.iter();
+        let first = query.next();
+        let extra = query.next().is_some();
+
+        match (first, extra) {
+            (Some(r), false) => Ok(r),
+            (None, _) => Err(UniqueQueryError::NoEntities(std::any::type_name::<Self>())),
+            (Some(r), _) => Err(UniqueQueryError::MultipleEntities {
+                result: r,
+                query_name: std::any::type_name::<Self>(),
+            }),
+        }
+    }
+
+    /// See [`Query::get_unique`]
+    pub fn get_unique_mut(
+        &mut self,
+    ) -> Result<<Q::Fetch as Fetch<'_>>::Item, UniqueQueryError<'_, Q>> {
+        let mut query = self.iter_mut();
+        let first = query.next();
+        let extra = query.next().is_some();
+
+        match (first, extra) {
+            (Some(r), false) => Ok(r),
+            (None, _) => Err(UniqueQueryError::NoEntities(std::any::type_name::<Self>())),
+            (Some(r), _) => Err(UniqueQueryError::MultipleEntities {
+                result: r,
+                query_name: std::any::type_name::<Self>(),
+            }),
+        }
+    }
 }
 
 /// An error that occurs when retrieving a specific [Entity]'s component from a [Query]
@@ -226,3 +262,23 @@ pub enum QueryComponentError {
     #[error("The requested entity does not exist.")]
     NoSuchEntity,
 }
+
+#[derive(Error)]
+pub enum UniqueQueryError<'a, Q: WorldQuery> {
+    #[error("No entities fit the query {0}")]
+    NoEntities(&'static str),
+    #[error("Multiple entities fit the query {query_name}!")]
+    MultipleEntities {
+        result: <Q::Fetch as Fetch<'a>>::Item,
+        query_name: &'static str,
+    },
+}
+
+impl<'a, Q: WorldQuery> Debug for UniqueQueryError<'a, Q> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::NoEntities(_) => f.debug_tuple("NoEntities").finish(),
+            Self::MultipleEntities { .. } => f.debug_tuple("MultipleEntities").finish(),
+        }
+    }
+}

examples/game/alien_cake_addict.rs

@@ -357,9 +357,8 @@ fn rotate_bonus(game: Res<Game>, time: Res<Time>, mut transforms: Query<&mut Tra
 
 // update the score displayed during the game
 fn scoreboard_system(game: Res<Game>, mut query: Query<&mut Text>) {
-    for mut text in query.iter_mut() {
-        text.sections[0].value = format!("Sugar Rush: {}", game.score);
-    }
+    let mut text = query.get_unique_mut().unwrap();
+    text.sections[0].value = format!("Sugar Rush: {}", game.score);
 }
 
 // restart the game when pressing spacebar

examples/game/breakout.rs

@@ -174,7 +174,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;
@@ -196,15 +196,14 @@ 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.sections[1].value = scoreboard.score.to_string();
-    }
+    let mut text = query.get_unique_mut().unwrap();
+    text.sections[0].value = format!("Score: {}", scoreboard.score);
 }
 
 fn ball_collision_system(
@@ -213,7 +212,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;