From bb3f24fd4e162ac11eb10f56497b6547dc83a3fd Mon Sep 17 00:00:00 2001
From: Mike Hsu <mike.hsu@gmail.com>
Date: Sat, 20 Jan 2024 23:12:14 -0800
Subject: [PATCH] change use query to use system state

---
 examples/demo.rs |  2 +-
 src/ecs_hooks.rs | 28 ++++++++++------------------
 2 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/examples/demo.rs b/examples/demo.rs
index 41cb797..087d829 100644
--- a/examples/demo.rs
+++ b/examples/demo.rs
@@ -33,7 +33,7 @@ fn Editor(cx: Scope) -> Element {
 
 #[component]
 fn SceneTree<'a>(cx: Scope, selected_entity: &'a UseStateSendable<Option<Entity>>) -> Element {
-    let entities = use_query_filtered::<(Entity, DebugName), Without<Node>>(cx);
+    let mut entities = use_query_filtered::<(Entity, DebugName), Without<Node>>(cx);
     let entities = entities.query();
     let mut entities = entities.into_iter().collect::<Vec<_>>();
     entities.sort_by_key(|(entity, _)| *entity);
diff --git a/src/ecs_hooks.rs b/src/ecs_hooks.rs
index aa1e735..339f49a 100644
--- a/src/ecs_hooks.rs
+++ b/src/ecs_hooks.rs
@@ -3,9 +3,9 @@ use bevy::{
     ecs::{
         component::ComponentId,
         event::{Event, EventIterator, Events, ManualEventReader},
-        query::{QueryState, ReadOnlyWorldQuery},
-        system::{Query, Resource},
-        world::{unsafe_world_cell::UnsafeWorldCell, World},
+        query::ReadOnlyWorldQuery,
+        system::{Query, Resource, SystemState},
+        world::World,
     },
     utils::{HashMap, HashSet},
 };
@@ -115,8 +115,8 @@ where
     });
 
     UseQuery {
-        query_state: QueryState::new(world),
-        world_cell: world.as_unsafe_world_cell(),
+        system_state: SystemState::new(world),
+        world_ref: world,
     }
 }
 
@@ -128,9 +128,9 @@ pub fn use_event_reader<E: Event>(cx: &ScopeState) -> EventIterator<'_, E> {
     event_reader.read(events)
 }
 
-pub struct UseQuery<'a, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery> {
-    query_state: QueryState<Q, F>,
-    world_cell: UnsafeWorldCell<'a>,
+pub struct UseQuery<'a, Q: ReadOnlyWorldQuery + 'static, F: ReadOnlyWorldQuery + 'static> {
+    system_state: SystemState<Query<'static, 'static, Q, F>>,
+    world_ref: &'a World,
 }
 
 impl<'a, Q, F> UseQuery<'a, Q, F>
@@ -138,15 +138,7 @@ where
     Q: ReadOnlyWorldQuery,
     F: ReadOnlyWorldQuery,
 {
-    pub fn query(&self) -> Query<Q, F> {
-        unsafe {
-            Query::new(
-                self.world_cell,
-                &self.query_state,
-                self.world_cell.last_change_tick(),
-                self.world_cell.change_tick(),
-                true,
-            )
-        }
+    pub fn query(&mut self) -> Query<Q, F> {
+        self.system_state.get(self.world_ref)
     }
 }