From 1bdb9d3b009c12433c2ef31127527afa7df33417 Mon Sep 17 00:00:00 2001
From: Boutillier <boris.boutillier@gmail.com>
Date: Mon, 5 Oct 2020 19:38:12 +0200
Subject: [PATCH] Fix Added behaviour for QueryOne get. (#543)

Query unchanged as impacts performances.
Added tests in bevy_ecs/hecs
---
 crates/bevy_ecs/hecs/src/query_one.rs | 12 ++++++++--
 crates/bevy_ecs/hecs/tests/tests.rs   | 32 +++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/crates/bevy_ecs/hecs/src/query_one.rs b/crates/bevy_ecs/hecs/src/query_one.rs
index cd355fb3bf278..ae7c330bbd2c9 100644
--- a/crates/bevy_ecs/hecs/src/query_one.rs
+++ b/crates/bevy_ecs/hecs/src/query_one.rs
@@ -37,7 +37,11 @@ impl<'a, Q: Query> QueryOne<'a, Q> {
     pub fn get(&mut self) -> Option<<Q::Fetch as Fetch<'_>>::Item> {
         unsafe {
             let mut fetch = Q::Fetch::get(self.archetype, self.index)?;
-            Some(fetch.next())
+            if fetch.should_skip() {
+                None
+            } else {
+                Some(fetch.next())
+            }
         }
     }
 
@@ -104,7 +108,11 @@ where
     {
         unsafe {
             let mut fetch = Q::Fetch::get(self.archetype, self.index)?;
-            Some(fetch.next())
+            if fetch.should_skip() {
+                None
+            } else {
+                Some(fetch.next())
+            }
         }
     }
 
diff --git a/crates/bevy_ecs/hecs/tests/tests.rs b/crates/bevy_ecs/hecs/tests/tests.rs
index f3911cb0dfb27..773f9d4c84ca4 100644
--- a/crates/bevy_ecs/hecs/tests/tests.rs
+++ b/crates/bevy_ecs/hecs/tests/tests.rs
@@ -365,3 +365,35 @@ fn remove_tracking() {
         "world clears result in 'removed component' states"
     );
 }
+
+#[test]
+fn added_tracking() {
+    let mut world = World::new();
+    let a = world.spawn((123,));
+
+    assert_eq!(world.query::<&i32>().iter().count(), 1);
+    assert_eq!(world.query::<Added<i32>>().iter().count(), 1);
+    assert_eq!(world.query_mut::<&i32>().iter().count(), 1);
+    assert_eq!(world.query_mut::<Added<i32>>().iter().count(), 1);
+    assert!(world.query_one::<&i32>(a).unwrap().get().is_some());
+    assert!(world.query_one::<Added<i32>>(a).unwrap().get().is_some());
+    assert!(world.query_one_mut::<&i32>(a).unwrap().get().is_some());
+    assert!(world
+        .query_one_mut::<Added<i32>>(a)
+        .unwrap()
+        .get()
+        .is_some());
+
+    world.clear_trackers();
+
+    assert_eq!(world.query::<&i32>().iter().count(), 1);
+    assert_eq!(world.query::<Added<i32>>().iter().count(), 0);
+    assert_eq!(world.query_mut::<&i32>().iter().count(), 1);
+    assert_eq!(world.query_mut::<Added<i32>>().iter().count(), 0);
+    assert!(world.query_one_mut::<&i32>(a).unwrap().get().is_some());
+    assert!(world
+        .query_one_mut::<Added<i32>>(a)
+        .unwrap()
+        .get()
+        .is_none());
+}