hecs/
entity_ref.rs

1use core::any::TypeId;
2use core::fmt::{self, Debug, Display, Formatter};
3use core::marker::PhantomData;
4use core::ops::{Deref, DerefMut, FnOnce};
5use core::ptr::NonNull;
6
7use crate::archetype::Archetype;
8use crate::entities::EntityMeta;
9use crate::{
10    ArchetypeColumn, ArchetypeColumnMut, Component, Entity, Fetch, MissingComponent, Query,
11    QueryOne,
12};
13
14/// Handle to an entity with any component types
15#[derive(Copy, Clone)]
16pub struct EntityRef<'a> {
17    meta: &'a [EntityMeta],
18    archetype: &'a Archetype,
19    /// Position of this entity in `archetype`
20    index: u32,
21}
22
23impl<'a> EntityRef<'a> {
24    pub(crate) unsafe fn new(meta: &'a [EntityMeta], archetype: &'a Archetype, index: u32) -> Self {
25        Self {
26            meta,
27            archetype,
28            index,
29        }
30    }
31
32    /// Get the [`Entity`] handle associated with this entity
33    #[inline]
34    pub fn entity(&self) -> Entity {
35        let id = self.archetype.entity_id(self.index);
36        Entity {
37            id,
38            generation: self.meta[id as usize].generation,
39        }
40    }
41
42    /// Determine whether this entity would satisfy the query `Q` without borrowing any components
43    pub fn satisfies<Q: Query>(&self) -> bool {
44        Q::Fetch::access(self.archetype).is_some()
45    }
46
47    /// Determine whether this entity has a `T` component without borrowing it
48    ///
49    /// Equivalent to [`satisfies::<&T>`](Self::satisfies)
50    pub fn has<T: Component>(&self) -> bool {
51        self.archetype.has::<T>()
52    }
53
54    /// Borrow a single component, if it exists
55    ///
56    /// `T` must be a shared or unique reference to a component type.
57    ///
58    /// # Example
59    /// ```
60    /// # use hecs::*;
61    /// let mut world = World::new();
62    /// let a = world.spawn((42, "abc"));
63    /// let e = world.entity(a).unwrap();
64    /// *e.get::<&mut i32>().unwrap() = 17;
65    /// assert_eq!(*e.get::<&i32>().unwrap(), 17);
66    /// ```
67    ///
68    /// Panics if `T` is a unique reference and the component is already borrowed, or if the
69    /// component is already uniquely borrowed.
70    pub fn get<T: ComponentRef<'a>>(&self) -> Option<T::Ref> {
71        T::get_component(*self)
72    }
73
74    /// Run a query against this entity
75    ///
76    /// Equivalent to invoking [`World::query_one`](crate::World::query_one) on the entity. May
77    /// outlive `self`.
78    ///
79    /// # Example
80    /// ```
81    /// # use hecs::*;
82    /// let mut world = World::new();
83    /// let a = world.spawn((123, true, "abc"));
84    /// // The returned query must outlive the borrow made by `get`
85    /// let mut query = world.entity(a).unwrap().query::<(&mut i32, &bool)>();
86    /// let (number, flag) = query.get().unwrap();
87    /// if *flag { *number *= 2; }
88    /// assert_eq!(*number, 246);
89    /// ```
90    pub fn query<Q: Query>(&self) -> QueryOne<'a, Q> {
91        unsafe { QueryOne::new(self.meta, self.archetype, self.index) }
92    }
93
94    /// Enumerate the types of the entity's components
95    ///
96    /// Convenient for dispatching component-specific logic for a single entity. For example, this
97    /// can be combined with a `HashMap<TypeId, Box<dyn Handler>>` where `Handler` is some
98    /// user-defined trait with methods for serialization, or to be called after spawning or before
99    /// despawning to maintain secondary indices.
100    pub fn component_types(&self) -> impl Iterator<Item = TypeId> + 'a {
101        self.archetype.types().iter().map(|ty| ty.id())
102    }
103
104    /// Number of components in this entity
105    pub fn len(&self) -> usize {
106        self.archetype.types().len()
107    }
108
109    /// Shorthand for `self.len() == 0`
110    pub fn is_empty(&self) -> bool {
111        self.len() == 0
112    }
113}
114
115unsafe impl Send for EntityRef<'_> {}
116unsafe impl Sync for EntityRef<'_> {}
117
118/// Shared borrow of an entity's component
119pub struct Ref<'a, T: ?Sized> {
120    borrow: ComponentBorrow<'a>,
121    target: NonNull<T>,
122    _phantom: PhantomData<&'a T>,
123}
124
125impl<'a, T: Component> Ref<'a, T> {
126    pub(crate) unsafe fn new(
127        archetype: &'a Archetype,
128        index: u32,
129    ) -> Result<Self, MissingComponent> {
130        let (target, borrow) = ComponentBorrow::for_component::<T>(archetype, index)?;
131        Ok(Self {
132            borrow,
133            target,
134            _phantom: PhantomData,
135        })
136    }
137}
138
139unsafe impl<T: ?Sized + Sync> Send for Ref<'_, T> {}
140unsafe impl<T: ?Sized + Sync> Sync for Ref<'_, T> {}
141
142impl<'a, T: ?Sized> Ref<'a, T> {
143    /// Transform the `Ref<'_, T>` to point to a part of the borrowed data, e.g.
144    /// a struct field.
145    ///
146    /// The `Ref<'_, T>` is already borrowed, so this cannot fail.
147    ///
148    /// # Examples
149    ///
150    /// ```no_run
151    /// # use hecs::{EntityRef, Ref};
152    /// struct Component {
153    ///     member: i32,
154    /// }
155    ///
156    /// # fn example(entity_ref: EntityRef<'_>) {
157    /// let component_ref = entity_ref.get::<&Component>()
158    ///     .expect("Entity does not contain an instance of \"Component\"");
159    /// let member_ref = Ref::map(component_ref, |component| &component.member);
160    /// println!("member = {:?}", *member_ref);
161    /// # }
162    /// ```
163    pub fn map<U: ?Sized, F>(orig: Ref<'a, T>, f: F) -> Ref<'a, U>
164    where
165        F: FnOnce(&T) -> &U,
166    {
167        let target = NonNull::from(f(&*orig));
168        Ref {
169            borrow: orig.borrow,
170            target,
171            _phantom: PhantomData,
172        }
173    }
174}
175
176impl<T: ?Sized> Deref for Ref<'_, T> {
177    type Target = T;
178    fn deref(&self) -> &T {
179        unsafe { self.target.as_ref() }
180    }
181}
182
183impl<T: ?Sized + Debug> Debug for Ref<'_, T> {
184    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
185        Debug::fmt(self.deref(), f)
186    }
187}
188
189impl<T: ?Sized + Display> Display for Ref<'_, T> {
190    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
191        Display::fmt(self.deref(), f)
192    }
193}
194
195impl<T: ?Sized> Clone for Ref<'_, T> {
196    fn clone(&self) -> Self {
197        Self {
198            borrow: self.borrow.clone(),
199            target: self.target,
200            _phantom: self._phantom,
201        }
202    }
203}
204
205/// Unique borrow of an entity's component
206pub struct RefMut<'a, T: ?Sized> {
207    borrow: ComponentBorrowMut<'a>,
208    target: NonNull<T>,
209    _phantom: PhantomData<&'a mut T>,
210}
211
212impl<'a, T: Component> RefMut<'a, T> {
213    pub(crate) unsafe fn new(
214        archetype: &'a Archetype,
215        index: u32,
216    ) -> Result<Self, MissingComponent> {
217        let (target, borrow) = ComponentBorrowMut::for_component::<T>(archetype, index)?;
218        Ok(Self {
219            borrow,
220            target,
221            _phantom: PhantomData,
222        })
223    }
224}
225
226unsafe impl<T: ?Sized + Send> Send for RefMut<'_, T> {}
227unsafe impl<T: ?Sized + Sync> Sync for RefMut<'_, T> {}
228
229impl<'a, T: ?Sized> RefMut<'a, T> {
230    /// Transform the `RefMut<'_, T>` to point to a part of the borrowed data, e.g.
231    /// a struct field.
232    ///
233    /// The `RefMut<'_, T>` is already mutably borrowed, so this cannot fail.
234    ///
235    /// # Examples
236    ///
237    /// ```no_run
238    /// # use hecs::{EntityRef, RefMut};
239    /// struct Component {
240    ///     member: i32,
241    /// }
242    ///
243    /// # fn example(entity_ref: EntityRef<'_>) {
244    /// let component_ref = entity_ref.get::<&mut Component>()
245    ///     .expect("Entity does not contain an instance of \"Component\"");
246    /// let mut member_ref = RefMut::map(component_ref, |component| &mut component.member);
247    /// *member_ref = 21;
248    /// println!("member = {:?}", *member_ref);
249    /// # }
250    /// ```
251    pub fn map<U: ?Sized, F>(mut orig: RefMut<'a, T>, f: F) -> RefMut<'a, U>
252    where
253        F: FnOnce(&mut T) -> &mut U,
254    {
255        let target = NonNull::from(f(&mut *orig));
256        RefMut {
257            borrow: orig.borrow,
258            target,
259            _phantom: PhantomData,
260        }
261    }
262}
263
264impl<T: ?Sized> Deref for RefMut<'_, T> {
265    type Target = T;
266    fn deref(&self) -> &T {
267        unsafe { self.target.as_ref() }
268    }
269}
270
271impl<T: ?Sized> DerefMut for RefMut<'_, T> {
272    fn deref_mut(&mut self) -> &mut T {
273        unsafe { self.target.as_mut() }
274    }
275}
276
277impl<T: ?Sized + Debug> Debug for RefMut<'_, T> {
278    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
279        Debug::fmt(self.deref(), f)
280    }
281}
282
283impl<T: ?Sized + Display> Display for RefMut<'_, T> {
284    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
285        Display::fmt(self.deref(), f)
286    }
287}
288
289/// `&T` or `&mut T` where `T` is some component type
290///
291/// The interface of this trait is a private implementation detail.
292pub trait ComponentRef<'a> {
293    /// Smart pointer to a component of the referenced type
294    #[doc(hidden)]
295    type Ref;
296
297    /// Smart pointer to a column of the referenced type in an [`Archetype`](crate::Archetype)
298    #[doc(hidden)]
299    type Column;
300
301    /// Component type referenced by `Ref`
302    #[doc(hidden)]
303    type Component: Component;
304
305    /// Fetch the component from `entity`
306    #[doc(hidden)]
307    fn get_component(entity: EntityRef<'a>) -> Option<Self::Ref>;
308
309    /// Construct from a raw pointer
310    ///
311    /// # Safety
312    ///
313    /// Dereferencing `raw` for lifetime `'a` must be sound
314    #[doc(hidden)]
315    unsafe fn from_raw(raw: *mut Self::Component) -> Self;
316
317    /// Borrow a column from an archetype
318    #[doc(hidden)]
319    fn get_column(archetype: &'a Archetype) -> Option<Self::Column>;
320}
321
322impl<'a, T: Component> ComponentRef<'a> for &'a T {
323    type Ref = Ref<'a, T>;
324
325    type Column = ArchetypeColumn<'a, T>;
326
327    type Component = T;
328
329    fn get_component(entity: EntityRef<'a>) -> Option<Self::Ref> {
330        Some(unsafe { Ref::new(entity.archetype, entity.index).ok()? })
331    }
332
333    unsafe fn from_raw(raw: *mut Self::Component) -> Self {
334        &*raw
335    }
336
337    fn get_column(archetype: &'a Archetype) -> Option<Self::Column> {
338        ArchetypeColumn::new(archetype)
339    }
340}
341
342impl<'a, T: Component> ComponentRef<'a> for &'a mut T {
343    type Ref = RefMut<'a, T>;
344
345    type Column = ArchetypeColumnMut<'a, T>;
346
347    type Component = T;
348
349    fn get_component(entity: EntityRef<'a>) -> Option<Self::Ref> {
350        Some(unsafe { RefMut::new(entity.archetype, entity.index).ok()? })
351    }
352
353    unsafe fn from_raw(raw: *mut Self::Component) -> Self {
354        &mut *raw
355    }
356
357    fn get_column(archetype: &'a Archetype) -> Option<Self::Column> {
358        ArchetypeColumnMut::new(archetype)
359    }
360}
361
362/// `&T` where `T` is some component type
363///
364/// Used when consistency demands that references to component types, rather than component types
365/// themselves, be supplied as a type parameter to a function that cannot operate on unique
366/// references.
367pub trait ComponentRefShared<'a>: ComponentRef<'a> {}
368
369impl<'a, T: Component> ComponentRefShared<'a> for &'a T {}
370
371struct ComponentBorrow<'a> {
372    archetype: &'a Archetype,
373    /// State index for the borrowed component in the `archetype`.
374    state: usize,
375}
376
377impl<'a> ComponentBorrow<'a> {
378    // This method is unsafe as if the `index` is out of bounds,
379    // then this will cause undefined behavior as the returned
380    // `target` will point to undefined memory.
381    unsafe fn for_component<T: Component>(
382        archetype: &'a Archetype,
383        index: u32,
384    ) -> Result<(NonNull<T>, Self), MissingComponent> {
385        let state = archetype
386            .get_state::<T>()
387            .ok_or_else(MissingComponent::new::<T>)?;
388
389        let target =
390            NonNull::new_unchecked(archetype.get_base::<T>(state).as_ptr().add(index as usize));
391
392        archetype.borrow::<T>(state);
393
394        Ok((target, Self { archetype, state }))
395    }
396}
397
398impl Clone for ComponentBorrow<'_> {
399    fn clone(&self) -> Self {
400        unsafe {
401            self.archetype.borrow_raw(self.state);
402        }
403        Self {
404            archetype: self.archetype,
405            state: self.state,
406        }
407    }
408}
409
410impl Drop for ComponentBorrow<'_> {
411    fn drop(&mut self) {
412        unsafe {
413            self.archetype.release_raw(self.state);
414        }
415    }
416}
417
418struct ComponentBorrowMut<'a> {
419    archetype: &'a Archetype,
420    /// State index for the borrowed component in the `archetype`.
421    state: usize,
422}
423
424impl<'a> ComponentBorrowMut<'a> {
425    // This method is unsafe as if the `index` is out of bounds,
426    // then this will cause undefined behavior as the returned
427    // `target` will point to undefined memory.
428    unsafe fn for_component<T: Component>(
429        archetype: &'a Archetype,
430        index: u32,
431    ) -> Result<(NonNull<T>, Self), MissingComponent> {
432        let state = archetype
433            .get_state::<T>()
434            .ok_or_else(MissingComponent::new::<T>)?;
435
436        let target =
437            NonNull::new_unchecked(archetype.get_base::<T>(state).as_ptr().add(index as usize));
438
439        archetype.borrow_mut::<T>(state);
440
441        Ok((target, Self { archetype, state }))
442    }
443}
444
445impl Drop for ComponentBorrowMut<'_> {
446    fn drop(&mut self) {
447        unsafe {
448            self.archetype.release_raw_mut(self.state);
449        }
450    }
451}