hephae_utils/
query_ext.rs

1use std::{
2    mem::MaybeUninit,
3    ops::{Deref, DerefMut},
4};
5
6use bevy_ecs::prelude::*;
7
8/// Extension traits for `Option<&T>`, `Option<&mut T>`, and `Option<Ref<T>>`.
9pub trait ComponentOption<T: Deref<Target: Component + Sized>>: Sized {
10    /// Unwraps the option or inserts a new one provided by a closure to the entity.
11    fn get_or_insert_with(self, commands: EntityCommands, insert: impl FnOnce() -> T::Target) -> RefOrInsert<T>;
12
13    /// Unwraps the option or inserts a new one by-value to the entity.
14    #[inline]
15    fn get_or_insert(self, commands: EntityCommands, insert: T::Target) -> RefOrInsert<T> {
16        self.get_or_insert_with(commands, || insert)
17    }
18
19    /// Unwraps the option or inserts a default instance to the entity.
20    #[inline]
21    fn get_or_default(self, commands: EntityCommands) -> RefOrInsert<T>
22    where
23        T::Target: Default,
24    {
25        self.get_or_insert_with(commands, T::Target::default)
26    }
27}
28
29impl<T: Deref<Target: Component + Sized>> ComponentOption<T> for Option<T> {
30    #[inline]
31    fn get_or_insert_with(self, commands: EntityCommands, insert: impl FnOnce() -> <T as Deref>::Target) -> RefOrInsert<T> {
32        RefOrInsert(match self {
33            Some(val) => RefOrInsertInner::Ref(val),
34            None => RefOrInsertInner::Spawn(MaybeUninit::new(insert()), commands),
35        })
36    }
37}
38
39/// Insert-guard returned by [`ComponentOption::get_or_insert_with`].
40pub struct RefOrInsert<'a, T: Deref<Target: Component + Sized>>(RefOrInsertInner<'a, T>);
41impl<T: Deref<Target: Component + Sized>> Deref for RefOrInsert<'_, T> {
42    type Target = T::Target;
43
44    #[inline]
45    fn deref(&self) -> &Self::Target {
46        match self.0 {
47            RefOrInsertInner::Ref(ref val) => val,
48            RefOrInsertInner::Spawn(ref val, ..) => unsafe { val.assume_init_ref() },
49        }
50    }
51}
52
53impl<T: DerefMut<Target: Component + Sized>> DerefMut for RefOrInsert<'_, T> {
54    #[inline]
55    fn deref_mut(&mut self) -> &mut Self::Target {
56        match self.0 {
57            RefOrInsertInner::Ref(ref mut val) => val,
58            RefOrInsertInner::Spawn(ref mut val, ..) => unsafe { val.assume_init_mut() },
59        }
60    }
61}
62
63impl<T: Deref<Target: Component + Sized>> Drop for RefOrInsert<'_, T> {
64    fn drop(&mut self) {
65        if let RefOrInsertInner::Spawn(insert, commands) = &mut self.0 {
66            commands.insert(unsafe { insert.assume_init_read() });
67        }
68    }
69}
70
71enum RefOrInsertInner<'a, T: Deref<Target: Component + Sized>> {
72    Ref(T),
73    Spawn(MaybeUninit<T::Target>, EntityCommands<'a>),
74}
hephae_utils/query_ext.rs

hephae_utils/
query_ext.rs
