use UnsafeWorldCell

JoJoJet · Jun 29, 2023 · 682060f · 682060f
1 parent c8f27b5
commit 682060f
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 18 deletions.
diff --git a/src/all.rs b/src/all.rs
@@ -1282,13 +1282,15 @@ unsafe impl<'a, Trait: ?Sized + TraitQuery> WorldQuery for ChangedAll<&'a Trait>
 
     #[inline]
     unsafe fn init_fetch<'w>(
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _state: &Self::State,
         last_run: Tick,
         this_run: Tick,
     ) -> ReadAllTraitsFetch<'w, Trait> {
         ReadAllTraitsFetch {
-            registry: world.resource(),
+            registry: world
+                .get_resource()
+                .unwrap_or_else(|| trait_registry_error()),
             table: None,
             sparse_sets: &world.storages().sparse_sets,
             last_run,
@@ -1403,13 +1405,15 @@ unsafe impl<'a, Trait: ?Sized + TraitQuery> WorldQuery for AddedAll<&'a Trait> {
 
     #[inline]
     unsafe fn init_fetch<'w>(
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _state: &Self::State,
         last_run: Tick,
         this_run: Tick,
     ) -> ReadAllTraitsFetch<'w, Trait> {
         ReadAllTraitsFetch {
-            registry: world.resource(),
+            registry: world
+                .get_resource()
+                .unwrap_or_else(|| trait_registry_error()),
             table: None,
             sparse_sets: &world.storages().sparse_sets,
             last_run,
@@ -1648,13 +1652,15 @@ unsafe impl<'a, Trait: ?Sized + TraitQuery> WorldQuery for AddedAll<&'a mut Trai
 
     #[inline]
     unsafe fn init_fetch<'w>(
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _state: &Self::State,
         last_run: Tick,
         this_run: Tick,
     ) -> WriteAllTraitsFetch<'w, Trait> {
         WriteAllTraitsFetch {
-            registry: world.resource(),
+            registry: world
+                .get_resource()
+                .unwrap_or_else(|| trait_registry_error()),
             table: None,
             sparse_sets: &world.storages().sparse_sets,
             last_run,
@@ -1770,13 +1776,15 @@ unsafe impl<'a, Trait: ?Sized + TraitQuery> WorldQuery for ChangedAll<&'a mut Tr
 
     #[inline]
     unsafe fn init_fetch<'w>(
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _state: &Self::State,
         last_run: Tick,
         this_run: Tick,
     ) -> WriteAllTraitsFetch<'w, Trait> {
         WriteAllTraitsFetch {
-            registry: world.resource(),
+            registry: world
+                .get_resource()
+                .unwrap_or_else(|| trait_registry_error()),
             table: None,
             sparse_sets: &world.storages().sparse_sets,
             last_run,

diff --git a/src/one.rs b/src/one.rs
@@ -286,7 +286,7 @@ unsafe impl<'a, Trait: ?Sized + TraitQuery> WorldQuery for ChangedOne<&'a Trait>
 
     #[inline]
     unsafe fn init_fetch<'w>(
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _state: &Self::State,
         last_run: Tick,
         this_run: Tick,
@@ -465,7 +465,7 @@ unsafe impl<'a, Trait: ?Sized + TraitQuery> WorldQuery for AddedOne<&'a Trait> {
 
     #[inline]
     unsafe fn init_fetch<'w>(
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _state: &Self::State,
         last_run: Tick,
         this_run: Tick,
@@ -854,7 +854,7 @@ unsafe impl<'a, Trait: ?Sized + TraitQuery> WorldQuery for ChangedOne<&'a mut Tr
 
     #[inline]
     unsafe fn init_fetch<'w>(
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _state: &Self::State,
         last_run: Tick,
         this_run: Tick,
@@ -1069,7 +1069,7 @@ unsafe impl<'a, Trait: ?Sized + TraitQuery> WorldQuery for AddedOne<&'a mut Trai
 
     #[inline]
     unsafe fn init_fetch<'w>(
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _state: &Self::State,
         last_run: Tick,
         this_run: Tick,

diff --git a/src/query_filter.rs b/src/query_filter.rs
@@ -1,10 +1,12 @@
-use crate::{debug_unreachable, TraitImplRegistry, TraitQuery, TraitQueryState};
+use crate::{
+    debug_unreachable, trait_registry_error, TraitImplRegistry, TraitQuery, TraitQueryState,
+};
 use bevy::ecs::archetype::{Archetype, ArchetypeComponentId};
 use bevy::ecs::component::{ComponentId, Tick};
 use bevy::ecs::entity::Entity;
 use bevy::ecs::query::{Access, FilteredAccess, ReadOnlyWorldQuery, WorldQuery};
 use bevy::ecs::storage::{ComponentSparseSet, SparseSets, Table, TableRow};
-use bevy::ecs::world::World;
+use bevy::ecs::world::{unsafe_world_cell::UnsafeWorldCell, World};
 use bevy::ptr::{ThinSlicePtr, UnsafeCellDeref};
 use std::cell::UnsafeCell;
 use std::marker::PhantomData;
@@ -49,13 +51,15 @@ unsafe impl<'a, Trait: ?Sized + TraitQuery> WorldQuery for OneAddedFilter<'a, Tr
     }
 
     unsafe fn init_fetch<'w>(
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _state: &Self::State,
         last_run: Tick,
         this_run: Tick,
     ) -> Self::Fetch<'w> {
         Self::Fetch::<'w> {
-            registry: world.resource(),
+            registry: world
+                .get_resource()
+                .unwrap_or_else(|| trait_registry_error()),
             storage: ChangeDetectionStorage::Uninit,
             sparse_sets: &world.storages().sparse_sets,
             last_run,
@@ -203,13 +207,15 @@ unsafe impl<'a, Trait: ?Sized + TraitQuery> WorldQuery for OneChangedFilter<'a,
     }
 
     unsafe fn init_fetch<'w>(
-        world: &'w World,
+        world: UnsafeWorldCell<'w>,
         _state: &Self::State,
         last_run: Tick,
         this_run: Tick,
     ) -> Self::Fetch<'w> {
         Self::Fetch::<'w> {
-            registry: world.resource(),
+            registry: world
+                .get_resource()
+                .unwrap_or_else(|| trait_registry_error()),
             storage: ChangeDetectionStorage::Uninit,
             sparse_sets: &world.storages().sparse_sets,
             last_run,