flax/
entry.rs

1use core::mem;
2
3use crate::{
4    archetype::RefMut,
5    component::ComponentValue,
6    writer::{Replace, SingleComponentWriter},
7    Component, Entity, World,
8};
9
10/// Entry like api for an entity's component
11pub enum Entry<'a, T: ComponentValue> {
12    /// A vacant entry
13    Vacant(VacantEntry<'a, T>),
14    /// An occupied entry
15    Occupied(OccupiedEntry<'a, T>),
16}
17
18/// A view into a vacant component entry
19pub struct VacantEntry<'a, T: ComponentValue> {
20    pub(crate) world: &'a mut World,
21    pub(crate) id: Entity,
22    pub(crate) component: Component<T>,
23}
24
25impl<'a, T: ComponentValue> VacantEntry<'a, T> {
26    /// Insert a value into the entry, returning a mutable reference to it
27    pub fn insert(self, value: T) -> RefMut<'a, T> {
28        let (loc, _) = self
29            .world
30            .set_with_writer(
31                self.id,
32                SingleComponentWriter::new(self.component.desc(), Replace { value }),
33            )
34            .expect("Entry is valid");
35
36        self.world.get_mut_at(loc, self.component).unwrap()
37    }
38}
39
40/// A view into an occupied component entry
41pub struct OccupiedEntry<'a, T: ComponentValue> {
42    pub(crate) borrow: RefMut<'a, T>,
43}
44
45impl<'a, T: ComponentValue> OccupiedEntry<'a, T> {
46    /// Convert the entry into a mutable reference
47    pub fn into_mut(self) -> RefMut<'a, T> {
48        self.borrow
49    }
50}
51
52impl<'a, T> Entry<'a, T>
53where
54    T: ComponentValue,
55{
56    /// Mutate the value in place
57    pub fn and_modify(mut self, mut func: impl FnMut(&mut T)) -> Self {
58        if let Self::Occupied(v) = &mut self {
59            (func)(&mut *v.borrow)
60        }
61
62        self
63    }
64
65    /// Returns the contained component or inserts a default.
66    pub fn or_insert(self, value: T) -> RefMut<'a, T> {
67        match self {
68            Entry::Vacant(slot) => slot.insert(value),
69            Entry::Occupied(slot) => slot.into_mut(),
70        }
71    }
72
73    /// Return the component in the entry or insert the default value.
74    pub fn or_default(self) -> RefMut<'a, T>
75    where
76        T: Default,
77    {
78        match self {
79            Entry::Vacant(slot) => slot.insert(Default::default()),
80            Entry::Occupied(slot) => slot.into_mut(),
81        }
82    }
83
84    /// Returns the contained component or inserts a default provided by the
85    /// function.
86    pub fn or_insert_with(self, func: impl FnOnce() -> T) -> RefMut<'a, T> {
87        match self {
88            Entry::Vacant(slot) => slot.insert((func)()),
89            Entry::Occupied(slot) => slot.into_mut(),
90        }
91    }
92
93    /// Replaces the current value and returns it
94    pub fn set(self, value: T) -> Option<T> {
95        match self {
96            Entry::Vacant(slot) => {
97                slot.insert(value);
98                None
99            }
100            Entry::Occupied(mut slot) => Some(mem::replace(&mut slot.borrow, value)),
101        }
102    }
103}