flax/fetch/
maybe_mut.rs

1use alloc::vec::Vec;
2use atomic_refcell::AtomicRef;
3use core::marker::PhantomData;
4
5use crate::{
6    archetype::{Cell, RefMut, Slot},
7    component::ComponentValue,
8    system::{Access, AccessKind},
9    Component, Entity, Fetch, FetchItem,
10};
11
12use super::{FetchAccessData, PreparedFetch, RandomFetch};
13
14/// A query for conservative mutablility.
15///
16/// This is useful for not triggering change detection when a component in a query isn't always
17/// modified.
18///
19/// Implements `ReadOnlyFetch` as the mutation is explicit and the returned reference is limited
20/// to the loop body, rather than the iterator.
21pub struct MaybeMut<T>(pub(crate) Component<T>);
22
23impl<'q, T: ComponentValue> FetchItem<'q> for MaybeMut<T> {
24    type Item = MutGuard<'q, T>;
25}
26
27impl<'w, T: ComponentValue> Fetch<'w> for MaybeMut<T> {
28    const MUTABLE: bool = false;
29
30    type Prepared = PreparedMaybeMut<'w, T>;
31
32    fn prepare(&'w self, data: super::FetchPrepareData<'w>) -> Option<Self::Prepared> {
33        let cell = data.arch.cell(self.0.key())?;
34        Some(PreparedMaybeMut {
35            cell,
36            new_tick: data.new_tick,
37            entities: data.arch.entities(),
38            _marker: PhantomData,
39        })
40    }
41
42    fn filter_arch(&self, data: FetchAccessData) -> bool {
43        data.arch.has(self.0.key())
44    }
45
46    fn access(&self, data: FetchAccessData, dst: &mut Vec<Access>) {
47        if data.arch.has(self.0.key()) {
48            dst.extend_from_slice(&[Access {
49                kind: AccessKind::Archetype {
50                    id: data.arch_id,
51                    component: self.0.key(),
52                },
53                mutable: true,
54            }])
55        }
56    }
57
58    fn describe(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
59        f.write_str("mut ")?;
60        f.write_str(self.0.name())
61    }
62
63    fn searcher(&self, searcher: &mut crate::ArchetypeSearcher) {
64        searcher.add_required(self.0.key())
65    }
66
67    fn by_ref(&self) -> crate::filter::RefFetch<Self>
68    where
69        Self: Sized,
70    {
71        crate::filter::RefFetch(self)
72    }
73}
74
75pub struct PreparedMaybeMut<'w, T> {
76    cell: &'w Cell,
77    new_tick: u32,
78    entities: &'w [Entity],
79    _marker: PhantomData<T>,
80}
81
82pub struct Batch<'a> {
83    cell: &'a Cell,
84    new_tick: u32,
85    ids: &'a [Entity],
86    slot: Slot,
87}
88
89impl<'w, 'q, T: ComponentValue> PreparedFetch<'q> for PreparedMaybeMut<'w, T> {
90    type Item = MutGuard<'q, T>;
91    type Chunk = Batch<'q>;
92
93    const HAS_FILTER: bool = false;
94
95    unsafe fn create_chunk(&'q mut self, slice: crate::archetype::Slice) -> Self::Chunk {
96        Batch {
97            cell: self.cell,
98            new_tick: self.new_tick,
99            ids: self.entities,
100            slot: slice.start,
101        }
102    }
103
104    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
105        let slot = chunk.slot;
106        chunk.slot += 1;
107
108        MutGuard {
109            slot,
110            cell: chunk.cell,
111            new_tick: chunk.new_tick,
112            id: *chunk.ids.get_unchecked(slot),
113            _marker: PhantomData,
114        }
115    }
116}
117
118impl<'w, 'q, T: ComponentValue> RandomFetch<'q> for PreparedMaybeMut<'w, T> {
119    #[inline]
120    unsafe fn fetch_shared(&'q self, slot: usize) -> Self::Item {
121        MutGuard {
122            slot,
123            cell: self.cell,
124            new_tick: self.new_tick,
125            id: self.entities[slot],
126            _marker: PhantomData,
127        }
128    }
129
130    unsafe fn fetch_shared_chunk(chunk: &Self::Chunk, slot: Slot) -> Self::Item {
131        MutGuard {
132            slot,
133            cell: chunk.cell,
134            new_tick: chunk.new_tick,
135            id: chunk.ids[slot],
136            _marker: PhantomData,
137        }
138    }
139}
140
141/// Protects against accidental mutation.
142///
143/// See: [`MaybeMut`]
144pub struct MutGuard<'w, T> {
145    slot: Slot,
146    id: Entity,
147    cell: &'w Cell,
148    new_tick: u32,
149    _marker: PhantomData<T>,
150}
151
152impl<'w, T: ComponentValue> MutGuard<'w, T> {
153    /// Acquire a shared reference to the current value without triggering a change
154    pub fn read(&self) -> AtomicRef<T> {
155        // Type is guaranteed by fetch constructor
156        unsafe { self.cell.get(self.slot).unwrap() }
157    }
158
159    /// Acquire a mutable reference to the current value.
160    ///
161    /// Triggers a change
162    pub fn write(&self) -> RefMut<T> {
163        // Type is guaranteed by constructor
164        self.cell
165            .get_mut(self.id, self.slot, self.new_tick)
166            .unwrap()
167    }
168}
flax/fetch/maybe_mut.rs

flax/fetch/
maybe_mut.rs
