1use core::mem;
2
3use crate::{
4 archetype::RefMut,
5 component::ComponentValue,
6 writer::{Replace, SingleComponentWriter},
7 Component, Entity, World,
8};
9
10pub enum Entry<'a, T: ComponentValue> {
12 Vacant(VacantEntry<'a, T>),
14 Occupied(OccupiedEntry<'a, T>),
16}
17
18pub 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 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
40pub struct OccupiedEntry<'a, T: ComponentValue> {
42 pub(crate) borrow: RefMut<'a, T>,
43}
44
45impl<'a, T: ComponentValue> OccupiedEntry<'a, T> {
46 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 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 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 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 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 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}