bevy_trait_resource/
lib.rs

1//! # bevy-trait-resource
2//! A way to get resources that implements a specific trait.
3//! 
4//! ```rust
5//!# use bevy_trait_resource::*;
6//! use bevy::prelude::*;
7//! use bevy_trait_resource::{trait_resource, TraitResourceExt};
8//! 
9//! #[trait_resource]
10//! pub trait IncrementTrait {
11//!     fn value(&self) -> i32;
12//!     fn increment(&mut self);
13//! }
14//! 
15//! #[derive(Resource, Default)]
16//! struct NumberValueResource {
17//!     value: i32,
18//! }
19//! 
20//! impl IncrementTrait for NumberValueResource {
21//!#     fn value(&self) -> i32 {
22//!#         self.value
23//!#     }
24//!#     fn increment(&mut self) {
25//!#         self.value += 1;
26//!#     }
27//!     // Trait implementation...
28//! }
29//! 
30//! pub fn increment_value_system(world: &mut World) {
31//!     for res_opt in world.get_resources_trait_mut::<dyn IncrementTrait>() {
32//!         if let Some(res) = res_opt {
33//!             res.increment();
34//!         }
35//!     }
36//! } 
37//! 
38//! struct SomePlugin;
39//! 
40//! impl Plugin for SomePlugin {
41//!     fn build(&self, app: &mut App) {
42//!         app.init_resource_as::<dyn IncrementTrait, NumberValueResource>();
43//!         app.add_systems(Update, increment_value_system);
44//!     }
45//! }
46//!# 
47//!# fn main() {
48//!#    App::new()
49//!#        .add_plugins(SomePlugin)
50//!#        .run();
51//!# }
52//! ```
53//! You can also register a resource to multiple traits.
54//! ```rust
55//!# use bevy_trait_resource::*;
56//!# use bevy::prelude::*;
57//!# use bevy_trait_resource::{trait_resource, TraitResourceExt};
58//!# 
59//!# #[trait_resource]
60//!# pub trait IncrementTrait {
61//!#     fn value(&self) -> i32;
62//!#     fn increment(&mut self);
63//!# }
64//!# 
65//!# #[derive(Resource, Default)]
66//!# struct NumberValueResource {
67//!#     value: i32,
68//!# }
69//!# 
70//!# impl IncrementTrait for NumberValueResource {
71//!#     fn value(&self) -> i32 {
72//!#         self.value
73//!#     }
74//!#     fn increment(&mut self) {
75//!#         self.value += 1;
76//!#     }
77//!# }
78//!# 
79//!# #[trait_resource]
80//!# pub trait SomeOtherTrait  {
81//!#     fn some_other_value(&self) -> i32;
82//!# }
83//!# 
84//!# impl SomeOtherTrait for NumberValueResource {
85//!#     fn some_other_value(&self) -> i32 {
86//!#         self.value
87//!#     }
88//!# }
89//!# 
90//!# pub fn increment_value_system(world: &mut World) {
91//!#     for res_opt in world.get_resources_trait_mut::<dyn IncrementTrait>() {
92//!#         if let Some(res) = res_opt {
93//!#             res.increment();
94//!#         }
95//!#     }
96//!# } 
97//!# 
98//!# struct SomePlugin;
99//!# 
100//! impl Plugin for SomePlugin {
101//!     fn build(&self, app: &mut App) {
102//!         app.init_resource::<NumberValueResource>();
103//! 
104//!         app.register_resource_as::<dyn IncrementTrait, NumberValueResource>()
105//!            .register_resource_as::<dyn SomeOtherTrait, NumberValueResource>();
106//!     }
107//! }
108//! 
109//!# fn main() {
110//!#    App::new()
111//!#        .add_plugins(SomePlugin)
112//!#        .run();
113//!# }
114//! ```
115//! Unregistering.
116//! ```rust
117//!# use bevy_trait_resource::*;
118//!# use bevy::prelude::*;
119//!# use bevy_trait_resource::{trait_resource, TraitResourceExt};
120//!# 
121//!# #[trait_resource]
122//!# pub trait IncrementTrait {
123//!#     fn value(&self) -> i32;
124//!#     fn increment(&mut self);
125//!# }
126//!# 
127//!# #[derive(Resource, Default)]
128//!# struct NumberValueResource {
129//!#     value: i32,
130//!# }
131//!# 
132//!# impl IncrementTrait for NumberValueResource {
133//!#     fn value(&self) -> i32 {
134//!#         self.value
135//!#     }
136//!#     fn increment(&mut self) {
137//!#         self.value += 1;
138//!#     }
139//!# }
140//!# 
141//!# struct SomePlugin; 
142//!#
143//!# impl Plugin for SomePlugin {
144//!#    fn build(&self, app: &mut App) {
145//! app.unregister_resource_from_trait::<dyn IncrementTrait, NumberValueResource>();
146//!#    }
147//!# }
148//!#
149//!# fn main() {
150//!#    App::new()
151//!#        .add_plugins(SomePlugin)
152//!#        .run();
153//!# }
154//! ```
155
156use std::marker::PhantomData;
157use bevy::{
158    ecs::component::ComponentId,
159    prelude::*,
160    ptr::{Ptr, PtrMut}
161};
162
163#[cfg(test)]
164mod tests;
165
166pub use bevy_trait_resource_macro::trait_resource;
167
168#[doc(hidden)]
169pub mod imports {
170    pub use bevy::ecs::{
171            world::World,
172            system::Resource,
173    };
174}
175
176pub trait TraitResource: 'static {}
177
178#[doc(hidden)]
179pub trait TraitResourceMarker<Trait: ?Sized + TraitResource> {
180    type Covered: Resource;
181    fn cast(_: *mut u8) -> *mut Trait;
182}
183
184/// Turns an untyped pointer into a trait object pointer,
185/// for a specific erased concrete type.
186struct DynCtor<Trait: ?Sized> {
187    cast: unsafe fn(*mut u8) -> *mut Trait,
188}
189
190impl<T: ?Sized> Copy for DynCtor<T> {}
191impl<T: ?Sized> Clone for DynCtor<T> {
192    fn clone(&self) -> Self {
193        *self
194    }
195}
196
197impl<Trait: ?Sized> DynCtor<Trait> {
198    #[inline]
199    unsafe fn cast(self, ptr: Ptr) -> &Trait {
200        &*(self.cast)(ptr.as_ptr())
201    }
202    #[inline]
203    unsafe fn cast_mut(self, ptr: PtrMut) -> &mut Trait {
204        &mut *(self.cast)(ptr.as_ptr())
205    }
206}
207
208struct TraitData<Trait: ?Sized> {
209    resource_component_id: ComponentId,
210    trait_ptr: DynCtor<Trait>,
211}
212
213impl<T: ?Sized> Copy for TraitData<T> {}
214impl<T: ?Sized> Clone for TraitData<T> {
215    fn clone(&self) -> Self {
216        *self
217    }
218}
219
220#[derive(Resource)]
221struct TraitResourceRegistry<Trait: ?Sized> {
222    trait_data: Vec<TraitData<Trait>>,
223}
224
225impl<Trait: ?Sized> TraitResourceRegistry<Trait> {
226    fn empty() -> Self {
227        Self { trait_data: vec![] }
228    }
229}
230
231impl<Trait: ?Sized> Default for TraitResourceRegistry<Trait> {
232    fn default() -> Self {
233        Self::empty()
234    }
235}
236
237impl<Trait: ?Sized + TraitResource> TraitResourceRegistry<Trait> {
238    /// Registers data for resource trait.
239    /// Overrides data with same component id.
240    fn register(&mut self, trait_data: TraitData<Trait>) {
241        let exists_in_index = self
242            .trait_data
243            .iter()
244            .position(|data| data.resource_component_id == trait_data.resource_component_id);
245
246        if let Some(index) = exists_in_index {
247            self.trait_data[index] = trait_data;
248        } else {
249            self.trait_data.push(trait_data);
250        }
251    }
252
253    /// Unregister the component from the registry.
254    /// Returns the new length of the registry.
255    fn unregister(&mut self, resource_component_id: ComponentId) -> usize {
256        self.trait_data.retain(|data| data.resource_component_id != resource_component_id);
257        self.trait_data.len()
258    }
259}
260
261/// An [`Iterator`] over resources as mutables that implements a trait.
262pub struct TraitResourceIteratorMut<'w, Trait: ?Sized + TraitResource> {
263    world: &'w mut World,
264    cursor: usize,
265    _marker: PhantomData<Trait>
266}
267
268impl<'w, Trait: ?Sized + TraitResource> TraitResourceIteratorMut<'w, Trait> {
269    fn new(world: &'w mut World) -> Self {
270        Self {
271            world,
272            cursor: 0,
273            _marker: PhantomData
274        }
275    }
276}
277
278impl<'w, Trait: ?Sized + TraitResource> Iterator for TraitResourceIteratorMut<'w, Trait> {
279    type Item = Option<&'w mut Trait>;
280
281    fn next(&mut self) -> Option<Self::Item> {
282        // TODO: Avoid getting resource every next call? (This is done because we cant have two mutable borrows of world).
283        if let Some(registry) = self.world.get_resource_mut::<TraitResourceRegistry<Trait>>()  {
284            if self.cursor >= registry.trait_data.len() {
285                return None;
286            }
287
288            let data = registry.trait_data[self.cursor];
289            self.cursor += 1;
290
291            if let Some(mut ptr) = self.world.get_resource_mut_by_id(data.resource_component_id) {
292                // SAFETY: We know that the pointer is valid because we got it from the world.
293                // We are also guaranteed that there is no duplicates of resource_component_id and resources
294                // and we can therefore "guarantee" that the iterator will only return one mutable
295                // reference to a single resource. So this should be OK?
296                let raw_ptr = unsafe { data.trait_ptr.cast_mut(ptr.as_mut()) as *mut Trait };
297                Some(Some(unsafe { &mut *raw_ptr }))
298            } else {
299                Some(None)
300            }
301        } else {
302            None
303        }
304    }
305}
306
307/// An [`Iterator`] over resources that implements a trait.
308pub struct TraitResourceIterator<'w, Trait: ?Sized + TraitResource> {
309    registry: Option<&'w TraitResourceRegistry<Trait>>,
310    world: &'w World,
311    cursor: usize,
312}
313
314impl<'w, Trait: ?Sized + TraitResource> TraitResourceIterator<'w, Trait> {
315    fn new(world: &'w World) -> Self {
316        Self {
317            registry: world.get_resource::<TraitResourceRegistry<Trait>>(),
318            world,
319            cursor: 0,
320        }
321    }
322}
323
324impl<'w, Trait: ?Sized + TraitResource> Iterator for TraitResourceIterator<'w, Trait> {
325    type Item = Option<&'w Trait>;
326
327    fn next(&mut self) -> Option<Self::Item> {
328        if let Some(registry) = self.registry {
329            if self.cursor >= registry.trait_data.len() {
330                return None
331            }
332
333            let data = registry.trait_data[self.cursor];
334            self.cursor += 1;
335
336            if let Some(ptr) = self.world.get_resource_by_id(data.resource_component_id) {
337                Some(Some(unsafe { data.trait_ptr.cast(ptr) }))
338            } else {
339                Some(None)
340            }
341        } else {
342            None
343        }
344    }
345}
346
347pub trait TraitResourceExt {
348    /// Inserts a resource into the world and registers its trait.
349    /// if the resource already exists, it will be overridden.
350    fn insert_resource_as<Trait: ?Sized + TraitResource, R: Resource>(&mut self, resource: R) -> &mut Self
351    where (R,): TraitResourceMarker<Trait, Covered = R>;
352
353    /// Initializes a resource into the world and registers its trait.
354    /// if the resource already exists, it will be overridden.
355    fn init_resource_as<Trait: ?Sized + TraitResource, R: Resource + Default>(&mut self) -> &mut Self
356    where (R,): TraitResourceMarker<Trait, Covered = R>;
357
358    // ?: Should this maybe warn instead of panic?
359    /// Registers a resource as implementing a trait.
360    /// If the resource is already registered nothing will happen.
361    /// # Panics 
362    /// Panics if the resource does not exist.
363    fn register_resource_as<Trait: ?Sized + TraitResource, R: Resource>(&mut self) -> &mut Self
364    where (R,): TraitResourceMarker<Trait, Covered = R>;
365
366    /// Get [`TraitResourceIterator<Trait>`]
367    fn get_resources_trait<Trait: ?Sized + TraitResource>(&self) -> TraitResourceIterator<Trait>;
368
369    /// Get [`TraitResourceIteratorMut<Trait>`]
370    fn get_resources_trait_mut<Trait: ?Sized + TraitResource>(&mut self) -> TraitResourceIteratorMut<Trait>;
371
372    /// Unregister a resource from trait.
373    fn unregister_resource_from_trait<Trait: ?Sized + TraitResource, R: Resource>(&mut self)
374    where (R,): TraitResourceMarker<Trait, Covered = R>;
375}
376
377impl TraitResourceExt for World {
378    fn insert_resource_as<Trait: ?Sized + TraitResource, R: Resource>(
379        &mut self,
380        resource: R,
381    ) -> &mut Self
382    where
383        (R,): TraitResourceMarker<Trait, Covered = R>,
384    {
385        self.insert_resource(resource);
386        self.register_resource_as::<Trait, R>();
387        self
388    }
389
390    fn init_resource_as<Trait: ?Sized + TraitResource, R: Resource + Default>(&mut self) -> &mut Self
391    where
392        (R,): TraitResourceMarker<Trait, Covered = R>
393    {
394        self.insert_resource_as::<Trait, R>(R::default());
395        self
396    }
397
398    fn register_resource_as<Trait: ?Sized + TraitResource, R: Resource>(&mut self) -> &mut Self
399    where
400        (R,): TraitResourceMarker<Trait, Covered = R>,
401    {
402        let resource_id = self
403            .components()
404            .resource_id::<R>()
405            .expect("Trying to register a nonexistent resource");
406
407        let resource_registry = self
408            .get_resource_or_insert_with::<TraitResourceRegistry<Trait>>(default)
409            .into_inner();
410
411        let trait_data = TraitData {
412            resource_component_id: resource_id,
413            trait_ptr: DynCtor { cast: <(R,)>::cast },
414        };
415
416        resource_registry.register(trait_data);
417        self
418    }
419
420    fn get_resources_trait<Trait: ?Sized + TraitResource>(&self) -> TraitResourceIterator<Trait> {
421        TraitResourceIterator::new(self)
422    }
423
424    fn get_resources_trait_mut<Trait: ?Sized + TraitResource>(&mut self) -> TraitResourceIteratorMut<Trait> {
425        TraitResourceIteratorMut::new(self)
426    }
427
428    // ? Should this maybe warn if the resource does not exists, because we should 
429    // ? unregister before removing the resource if we also want to remove the resource
430    // Also removes the registry if it is empty.
431    fn unregister_resource_from_trait<Trait: ?Sized + TraitResource, R: Resource>(&mut self)
432    where
433        (R,): TraitResourceMarker<Trait, Covered = R>
434    {
435        let resource_id_opt = self
436            .components()
437            .resource_id::<R>();
438
439        if let Some(resource_id) = resource_id_opt {
440            match self.get_resource_mut::<TraitResourceRegistry<Trait>>() {
441                Some(mut registry) => {
442                    let new_size = registry.unregister(resource_id);
443                    if new_size == 0 {
444                        self.remove_resource::<TraitResourceRegistry<Trait>>();
445                    }
446                },
447                _ => {}
448            }
449        }
450    }
451}
452
453impl TraitResourceExt for App {
454    fn insert_resource_as<Trait: ?Sized + TraitResource, R: Resource>(
455        &mut self,
456        resource: R,
457    ) -> &mut Self
458    where
459        (R,): TraitResourceMarker<Trait, Covered = R>,
460    {
461        self.world.insert_resource_as::<Trait, R>(resource);
462        self
463    }
464
465    fn init_resource_as<Trait: ?Sized + TraitResource, R: Resource + Default>(&mut self) -> &mut Self
466    where
467        (R,): TraitResourceMarker<Trait, Covered = R>
468    {
469        self.world.init_resource_as::<Trait, R>();
470        self
471    }
472
473    fn register_resource_as<Trait: ?Sized + TraitResource, R: Resource>(&mut self) -> &mut Self
474    where
475        (R,): TraitResourceMarker<Trait, Covered = R> 
476    {
477        self.world.register_resource_as::<Trait, R>();
478        self
479    }
480
481    fn get_resources_trait<Trait: ?Sized + TraitResource>(&self) -> TraitResourceIterator<Trait> {
482        self.world.get_resources_trait::<Trait>()
483    }
484
485    fn get_resources_trait_mut<Trait: ?Sized + TraitResource>(&mut self) -> TraitResourceIteratorMut<Trait> {
486        self.world.get_resources_trait_mut::<Trait>()
487    }
488
489    fn unregister_resource_from_trait<Trait: ?Sized + TraitResource, R: Resource>(&mut self)
490    where
491        (R,): TraitResourceMarker<Trait, Covered = R>
492    {
493        self.world.unregister_resource_from_trait::<Trait, R>();
494    }
495}
bevy_trait_resource/lib.rs

bevy_trait_resource/
lib.rs
