flax/
relation.rs

1use core::{
2    fmt::{self, Display, Formatter},
3    marker::PhantomData,
4    sync::atomic::AtomicU32,
5};
6
7use alloc::collections::btree_map::Range;
8use atomic_refcell::AtomicRef;
9
10use crate::{
11    archetype::{Archetype, RefMut, Slot},
12    component::{dummy, ComponentKey, ComponentValue},
13    entity::EntityKind,
14    fetch::{nth_relation, NthRelation},
15    filter::{WithRelation, WithoutRelation},
16    vtable::{ComponentVTable, UntypedVTable},
17    Component, Entity,
18};
19
20/// Relation helper trait
21pub trait RelationExt<T>
22where
23    T: ComponentValue,
24{
25    /// Returns the relation id
26    fn id(&self) -> Entity;
27    /// Returns the vtable of the relation
28    fn vtable(&self) -> &'static UntypedVTable;
29    /// Instantiate the relation
30    fn of(&self, target: Entity) -> Component<T>;
31    /// Construct a new filter yielding entities with this kind of relation
32    fn with_relation(self) -> WithRelation;
33    /// Construct a new filter yielding entities without this kind of relation
34    fn without_relation(self) -> WithoutRelation;
35
36    /// Convert this into a concrete relation representation
37    fn as_relation(&self) -> Relation<T> {
38        Relation {
39            id: self.id(),
40            vtable: self.vtable(),
41            marker: PhantomData,
42        }
43    }
44
45    /// Query the nth relation of the specified kind.
46    ///
47    /// This is useful for [`Exclusive`](crate::metadata::Exclusive) relations where there is only one parent
48    ///
49    /// **Note**: Fails to match if there is no nth relation, prefer using [`opt`](crate::FetchExt::opt) for
50    /// optional relations.
51    fn nth_relation(self, n: usize) -> NthRelation<T>
52    where
53        Self: Sized,
54    {
55        nth_relation(self, n)
56    }
57
58    /// Query the first relation of the specified kind.
59    fn first_relation(self) -> NthRelation<T>
60    where
61        Self: Sized,
62    {
63        nth_relation(self, 0)
64    }
65}
66
67impl<T, F> RelationExt<T> for F
68where
69    F: Fn(Entity) -> Component<T>,
70    T: ComponentValue,
71{
72    fn id(&self) -> Entity {
73        (self)(dummy()).id()
74    }
75
76    fn vtable(&self) -> &'static UntypedVTable {
77        (self)(dummy()).vtable()
78    }
79
80    fn of(&self, target: Entity) -> Component<T> {
81        (self)(target)
82    }
83
84    fn with_relation(self) -> WithRelation {
85        let c = self(dummy());
86        WithRelation {
87            relation: c.id(),
88            name: c.name(),
89        }
90    }
91
92    fn without_relation(self) -> WithoutRelation {
93        let c = self(dummy());
94        WithoutRelation {
95            relation: c.id(),
96            name: c.name(),
97        }
98    }
99}
100
101/// Represents a relation which can connect to entities
102pub struct Relation<T> {
103    pub(crate) id: Entity,
104    vtable: &'static UntypedVTable,
105    marker: PhantomData<T>,
106}
107
108impl<T> Eq for Relation<T> {}
109
110impl<T> PartialEq for Relation<T> {
111    fn eq(&self, other: &Self) -> bool {
112        self.id == other.id
113    }
114}
115
116impl<T> Copy for Relation<T> {}
117
118impl<T> Clone for Relation<T> {
119    fn clone(&self) -> Self {
120        *self
121    }
122}
123
124impl<T> fmt::Debug for Relation<T> {
125    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
126        f.debug_struct("Relation").field("id", &self.id).finish()
127    }
128}
129
130impl<T> Display for Relation<T> {
131    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
132        write!(f, "{}({})", self.vtable.name, self.id)
133    }
134}
135
136impl<T> Relation<T>
137where
138    T: ComponentValue,
139{
140    pub(crate) fn new(id: Entity, vtable: &'static ComponentVTable<T>) -> Self {
141        Self {
142            id,
143            vtable,
144            marker: PhantomData,
145        }
146    }
147
148    #[doc(hidden)]
149    pub fn static_init(
150        id: &AtomicU32,
151        kind: EntityKind,
152        vtable: &'static ComponentVTable<T>,
153    ) -> Self {
154        let id = Entity::static_init(id, kind);
155
156        Self {
157            id,
158            vtable,
159            marker: PhantomData,
160        }
161    }
162
163    /// Returns the relation name
164    pub fn name(&self) -> &'static str {
165        self.vtable.name
166    }
167}
168
169impl<T: ComponentValue> RelationExt<T> for Relation<T> {
170    #[inline]
171    fn id(&self) -> Entity {
172        self.id
173    }
174
175    #[inline]
176    fn vtable(&self) -> &'static UntypedVTable {
177        self.vtable
178    }
179
180    fn of(&self, target: Entity) -> Component<T> {
181        Component::from_raw_parts(ComponentKey::new(self.id, Some(target)), self.vtable)
182    }
183
184    #[inline]
185    fn with_relation(self) -> WithRelation {
186        WithRelation {
187            relation: self.id(),
188            name: self.name(),
189        }
190    }
191
192    #[inline]
193    fn without_relation(self) -> WithoutRelation {
194        WithoutRelation {
195            relation: self.id(),
196            name: self.name(),
197        }
198    }
199}
200
201/// Allows to iterate all relations of a specific type for an entity
202pub struct RelationIter<'a, T> {
203    cells: Range<'a, ComponentKey, usize>,
204    arch: &'a Archetype,
205    slot: Slot,
206    marker: PhantomData<T>,
207}
208
209impl<'a, T: ComponentValue> RelationIter<'a, T> {
210    pub(crate) fn new(relation: impl RelationExt<T>, arch: &'a Archetype, slot: Slot) -> Self {
211        let relation = relation.id();
212        Self {
213            cells: arch.components().range(
214                ComponentKey::new(relation, Some(Entity::MIN))
215                    ..=ComponentKey::new(relation, Some(Entity::MAX)),
216            ),
217            slot,
218            marker: PhantomData,
219            arch,
220        }
221    }
222}
223
224impl<'a, T> Iterator for RelationIter<'a, T>
225where
226    T: ComponentValue,
227{
228    type Item = (Entity, AtomicRef<'a, T>);
229
230    fn next(&mut self) -> Option<Self::Item> {
231        let (&key, &cell_index) = self.cells.next()?;
232        // Safety: the type matches the relation ext
233        Some((key.target.unwrap(), unsafe {
234            self.arch.cells()[cell_index].get::<T>(self.slot).unwrap()
235        }))
236    }
237}
238
239/// See: [RelationIter]
240pub struct RelationIterMut<'a, T> {
241    entities: &'a [Entity],
242    cells: Range<'a, ComponentKey, usize>,
243    arch: &'a Archetype,
244    slot: Slot,
245    change_tick: u32,
246    marker: PhantomData<T>,
247}
248
249impl<'a, T: ComponentValue> RelationIterMut<'a, T> {
250    pub(crate) fn new(
251        relation: impl RelationExt<T>,
252        arch: &'a Archetype,
253        slot: Slot,
254        change_tick: u32,
255    ) -> Self {
256        let relation = relation.id();
257        Self {
258            cells: arch.components().range(
259                ComponentKey::new(relation, Some(Entity::MIN))
260                    ..=ComponentKey::new(relation, Some(Entity::MAX)),
261            ),
262            slot,
263            marker: PhantomData,
264            change_tick,
265            entities: arch.entities(),
266            arch,
267        }
268    }
269}
270
271impl<'a, T> Iterator for RelationIterMut<'a, T>
272where
273    T: ComponentValue,
274{
275    type Item = (Entity, RefMut<'a, T>);
276
277    fn next(&mut self) -> Option<Self::Item> {
278        let (&key, &cell_index) = self.cells.next()?;
279        Some((
280            key.target.unwrap(),
281            self.arch.cells()[cell_index]
282                .get_mut::<T>(self.entities[self.slot], self.slot, self.change_tick)
283                .unwrap(),
284        ))
285    }
286}
flax/relation.rs

flax/
relation.rs
