From 9d60563adf4167059fcd1a0addc11487a4394c15 Mon Sep 17 00:00:00 2001
From: TheRawMeatball <therawmeatball@gmail.com>
Date: Mon, 8 Mar 2021 21:21:47 +0000
Subject: [PATCH] Query::get_unique (#1263)

Adds `get_unique` and `get_unique_mut` to extend the query api and cover a common use case. Also establishes a second impl block where non-core APIs that don't access the internal fields of queries can live.
---
 crates/bevy_ecs/src/system/query.rs | 42 ++++++++++++++++++++++++++++-
 examples/game/alien_cake_addict.rs  |  5 ++--
 examples/game/breakout.rs           | 11 ++++----
 3 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs
index a98f5cd36c63d..d19056b86c42f 100644
--- a/crates/bevy_ecs/src/system/query.rs
+++ b/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,38 @@ where
             Err(QueryComponentError::MissingWriteAccess)
         }
     }
+
+    pub fn single(&self) -> Result<<Q::Fetch as Fetch<'_>>::Item, QuerySingleError>
+    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(QuerySingleError::NoEntities(std::any::type_name::<Self>())),
+            (Some(_), _) => Err(QuerySingleError::MultipleEntities(std::any::type_name::<
+                Self,
+            >())),
+        }
+    }
+
+    /// See [`Query::single`]
+    pub fn single_mut(&mut self) -> Result<<Q::Fetch as Fetch<'_>>::Item, QuerySingleError> {
+        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(QuerySingleError::NoEntities(std::any::type_name::<Self>())),
+            (Some(_), _) => Err(QuerySingleError::MultipleEntities(std::any::type_name::<
+                Self,
+            >())),
+        }
+    }
 }
 
 /// An error that occurs when retrieving a specific [Entity]'s component from a [Query]
@@ -226,3 +258,11 @@ pub enum QueryComponentError {
     #[error("The requested entity does not exist.")]
     NoSuchEntity,
 }
+
+#[derive(Debug, Error)]
+pub enum QuerySingleError {
+    #[error("No entities fit the query {0}")]
+    NoEntities(&'static str),
+    #[error("Multiple entities fit the query {0}!")]
+    MultipleEntities(&'static str),
+}
diff --git a/examples/game/alien_cake_addict.rs b/examples/game/alien_cake_addict.rs
index 3c50a992a6b4a..d1b0546c6552f 100644
--- a/examples/game/alien_cake_addict.rs
+++ b/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.single_mut().unwrap();
+    text.sections[0].value = format!("Sugar Rush: {}", game.score);
 }
 
 // restart the game when pressing spacebar
diff --git a/examples/game/breakout.rs b/examples/game/breakout.rs
index 1cd312439a429..dea6e834c8ac8 100644
--- a/examples/game/breakout.rs
+++ b/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.single_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.single_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.single_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.single_mut() {
         let ball_size = sprite.size;
         let velocity = &mut ball.velocity;