From 97d94b66c0cdb6a997423b8e2bbab8336e47a720 Mon Sep 17 00:00:00 2001
From: James Liu <contact@jamessliu.com>
Date: Tue, 15 Nov 2022 00:19:11 +0000
Subject: [PATCH] Fix get_unchecked_manual using archetype index instead of
 table row. (#6625)

# Objective
Fix #6623.

## Solution
Use the right table row instead of the `EntityLocation` archetype index.
---
 crates/bevy_ecs/src/lib.rs         | 23 +++++++++++++++++++++++
 crates/bevy_ecs/src/query/state.rs |  5 +++--
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs
index 04bc2ed7ffec4..a91ed6634f1b6 100644
--- a/crates/bevy_ecs/src/lib.rs
+++ b/crates/bevy_ecs/src/lib.rs
@@ -707,6 +707,29 @@ mod tests {
         assert!(i32_query.get(&world, a).is_err());
     }
 
+    #[test]
+    fn query_get_works_across_sparse_removal() {
+        // Regression test for: https://github.com/bevyengine/bevy/issues/6623
+        let mut world = World::new();
+        let a = world.spawn((TableStored("abc"), SparseStored(123))).id();
+        let b = world.spawn((TableStored("def"), SparseStored(456))).id();
+        let c = world
+            .spawn((TableStored("ghi"), SparseStored(789), B(1)))
+            .id();
+
+        let mut query = world.query::<&TableStored>();
+        assert_eq!(query.get(&world, a).unwrap(), &TableStored("abc"));
+        assert_eq!(query.get(&world, b).unwrap(), &TableStored("def"));
+        assert_eq!(query.get(&world, c).unwrap(), &TableStored("ghi"));
+
+        world.entity_mut(b).remove::<SparseStored>();
+        world.entity_mut(c).remove::<SparseStored>();
+
+        assert_eq!(query.get(&world, a).unwrap(), &TableStored("abc"));
+        assert_eq!(query.get(&world, b).unwrap(), &TableStored("def"));
+        assert_eq!(query.get(&world, c).unwrap(), &TableStored("ghi"));
+    }
+
     #[test]
     fn remove_tracking() {
         let mut world = World::new();
diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs
index 18874fd107fae..245ff7dacb7b4 100644
--- a/crates/bevy_ecs/src/query/state.rs
+++ b/crates/bevy_ecs/src/query/state.rs
@@ -418,6 +418,7 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
         let mut fetch = Q::init_fetch(world, &self.fetch_state, last_change_tick, change_tick);
         let mut filter = F::init_fetch(world, &self.filter_state, last_change_tick, change_tick);
 
+        let table_row = archetype.entity_table_row(location.index);
         let table = world
             .storages()
             .tables
@@ -426,8 +427,8 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
         Q::set_archetype(&mut fetch, &self.fetch_state, archetype, table);
         F::set_archetype(&mut filter, &self.filter_state, archetype, table);
 
-        if F::filter_fetch(&mut filter, entity, location.index) {
-            Ok(Q::fetch(&mut fetch, entity, location.index))
+        if F::filter_fetch(&mut filter, entity, table_row) {
+            Ok(Q::fetch(&mut fetch, entity, table_row))
         } else {
             Err(QueryEntityError::QueryDoesNotMatch(entity))
         }