flax/fetch/
component_mut.rs

1use alloc::vec::Vec;
2
3use core::fmt::{self, Formatter};
4
5use crate::{
6    archetype::{Archetype, CellMutGuard, Slice},
7    component::ComponentValue,
8    system::{Access, AccessKind},
9    util::PtrMut,
10    Component, Fetch, FetchItem,
11};
12
13use super::{FetchAccessData, FetchPrepareData, PreparedFetch};
14
15#[derive(Debug, Clone)]
16/// Mutable component fetch
17/// See [crate::Component::as_mut]
18pub struct Mutable<T>(pub(crate) Component<T>);
19
20impl<'w, T> Fetch<'w> for Mutable<T>
21where
22    T: ComponentValue,
23{
24    const MUTABLE: bool = true;
25
26    type Prepared = WriteComponent<'w, T>;
27
28    #[inline]
29    fn prepare(&self, data: FetchPrepareData<'w>) -> Option<Self::Prepared> {
30        let guard = data.arch.borrow_mut(self.0.key())?;
31
32        Some(WriteComponent {
33            guard,
34            arch: data.arch,
35            tick: data.new_tick,
36        })
37    }
38
39    #[inline]
40    fn filter_arch(&self, data: FetchAccessData) -> bool {
41        data.arch.has(self.0.key())
42    }
43
44    #[inline]
45    fn access(&self, data: FetchAccessData, dst: &mut Vec<Access>) {
46        if data.arch.has(self.0.key()) {
47            dst.extend_from_slice(&[Access {
48                kind: AccessKind::Archetype {
49                    id: data.arch_id,
50                    component: self.0.key(),
51                },
52                mutable: true,
53            }])
54        }
55    }
56
57    fn describe(&self, f: &mut Formatter) -> fmt::Result {
58        f.write_str("mut ")?;
59        f.write_str(self.0.name())
60    }
61
62    fn searcher(&self, searcher: &mut crate::ArchetypeSearcher) {
63        searcher.add_required(self.0.key())
64    }
65}
66
67impl<'q, T: ComponentValue> FetchItem<'q> for Mutable<T> {
68    type Item = &'q mut T;
69}
70
71#[doc(hidden)]
72pub struct WriteComponent<'a, T> {
73    guard: CellMutGuard<'a, [T]>,
74    arch: &'a Archetype,
75    tick: u32,
76}
77
78impl<'w, 'q, T: 'q + ComponentValue> PreparedFetch<'q> for WriteComponent<'w, T> {
79    type Item = &'q mut T;
80    type Chunk = PtrMut<'q, T>;
81
82    const HAS_FILTER: bool = false;
83
84    unsafe fn create_chunk(&'q mut self, slots: Slice) -> Self::Chunk {
85        self.guard
86            .set_modified(&self.arch.entities[slots.as_range()], slots, self.tick);
87
88        // Convert directly into a non-overlapping subslice without reading the whole slice
89        PtrMut::new((self.guard.storage().as_ptr() as *mut T).add(slots.start))
90    }
91
92    #[inline]
93    // See: <https://godbolt.org/z/8fWa136b9>
94    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
95        let old = chunk.as_ptr();
96        chunk.advance(1);
97        &mut *old
98    }
99}