entity_map.rs

  1use crate::{App, AppContext, GpuiBorrow, VisualContext, Window, seal::Sealed};
  2use anyhow::{Context as _, Result};
  3use collections::FxHashSet;
  4use derive_more::{Deref, DerefMut};
  5use parking_lot::{RwLock, RwLockUpgradableReadGuard};
  6use slotmap::{KeyData, SecondaryMap, SlotMap};
  7use std::{
  8    any::{Any, TypeId, type_name},
  9    cell::RefCell,
 10    cmp::Ordering,
 11    fmt::{self, Display},
 12    hash::{Hash, Hasher},
 13    marker::PhantomData,
 14    mem,
 15    num::NonZeroU64,
 16    sync::{
 17        Arc, Weak,
 18        atomic::{AtomicU64, AtomicUsize, Ordering::SeqCst},
 19    },
 20    thread::panicking,
 21};
 22
 23use super::Context;
 24use crate::util::atomic_incr_if_not_zero;
 25#[cfg(any(test, feature = "leak-detection"))]
 26use collections::HashMap;
 27
 28slotmap::new_key_type! {
 29    /// A unique identifier for a entity across the application.
 30    pub struct EntityId;
 31}
 32
 33impl From<u64> for EntityId {
 34    fn from(value: u64) -> Self {
 35        Self(KeyData::from_ffi(value))
 36    }
 37}
 38
 39impl EntityId {
 40    /// Converts this entity id to a [NonZeroU64]
 41    pub fn as_non_zero_u64(self) -> NonZeroU64 {
 42        NonZeroU64::new(self.0.as_ffi()).unwrap()
 43    }
 44
 45    /// Converts this entity id to a [u64]
 46    pub fn as_u64(self) -> u64 {
 47        self.0.as_ffi()
 48    }
 49}
 50
 51impl Display for EntityId {
 52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 53        write!(f, "{}", self.as_u64())
 54    }
 55}
 56
 57pub(crate) struct EntityMap {
 58    entities: SecondaryMap<EntityId, Box<dyn Any>>,
 59    pub accessed_entities: RefCell<FxHashSet<EntityId>>,
 60    ref_counts: Arc<RwLock<EntityRefCounts>>,
 61}
 62
 63struct EntityRefCounts {
 64    counts: SlotMap<EntityId, AtomicUsize>,
 65    dropped_entity_ids: Vec<EntityId>,
 66    #[cfg(any(test, feature = "leak-detection"))]
 67    leak_detector: LeakDetector,
 68}
 69
 70impl EntityMap {
 71    pub fn new() -> Self {
 72        Self {
 73            entities: SecondaryMap::new(),
 74            accessed_entities: RefCell::new(FxHashSet::default()),
 75            ref_counts: Arc::new(RwLock::new(EntityRefCounts {
 76                counts: SlotMap::with_key(),
 77                dropped_entity_ids: Vec::new(),
 78                #[cfg(any(test, feature = "leak-detection"))]
 79                leak_detector: LeakDetector {
 80                    next_handle_id: 0,
 81                    entity_handles: HashMap::default(),
 82                },
 83            })),
 84        }
 85    }
 86
 87    /// Reserve a slot for an entity, which you can subsequently use with `insert`.
 88    pub fn reserve<T: 'static>(&self) -> Slot<T> {
 89        let id = self.ref_counts.write().counts.insert(1.into());
 90        Slot(Entity::new(id, Arc::downgrade(&self.ref_counts)))
 91    }
 92
 93    /// Insert an entity into a slot obtained by calling `reserve`.
 94    pub fn insert<T>(&mut self, slot: Slot<T>, entity: T) -> Entity<T>
 95    where
 96        T: 'static,
 97    {
 98        let mut accessed_entities = self.accessed_entities.borrow_mut();
 99        accessed_entities.insert(slot.entity_id);
100
101        let handle = slot.0;
102        self.entities.insert(handle.entity_id, Box::new(entity));
103        handle
104    }
105
106    /// Move an entity to the stack.
107    #[track_caller]
108    pub fn lease<T>(&mut self, pointer: &Entity<T>) -> Lease<T> {
109        self.assert_valid_context(pointer);
110        let mut accessed_entities = self.accessed_entities.borrow_mut();
111        accessed_entities.insert(pointer.entity_id);
112
113        let entity = Some(
114            self.entities
115                .remove(pointer.entity_id)
116                .unwrap_or_else(|| double_lease_panic::<T>("update")),
117        );
118        Lease {
119            entity,
120            id: pointer.entity_id,
121            entity_type: PhantomData,
122        }
123    }
124
125    /// Returns an entity after moving it to the stack.
126    pub fn end_lease<T>(&mut self, mut lease: Lease<T>) {
127        self.entities.insert(lease.id, lease.entity.take().unwrap());
128    }
129
130    pub fn read<T: 'static>(&self, entity: &Entity<T>) -> &T {
131        self.assert_valid_context(entity);
132        let mut accessed_entities = self.accessed_entities.borrow_mut();
133        accessed_entities.insert(entity.entity_id);
134
135        self.entities
136            .get(entity.entity_id)
137            .and_then(|entity| entity.downcast_ref())
138            .unwrap_or_else(|| double_lease_panic::<T>("read"))
139    }
140
141    fn assert_valid_context(&self, entity: &AnyEntity) {
142        debug_assert!(
143            Weak::ptr_eq(&entity.entity_map, &Arc::downgrade(&self.ref_counts)),
144            "used a entity with the wrong context"
145        );
146    }
147
148    pub fn extend_accessed(&mut self, entities: &FxHashSet<EntityId>) {
149        self.accessed_entities
150            .borrow_mut()
151            .extend(entities.iter().copied());
152    }
153
154    pub fn clear_accessed(&mut self) {
155        self.accessed_entities.borrow_mut().clear();
156    }
157
158    pub fn take_dropped(&mut self) -> Vec<(EntityId, Box<dyn Any>)> {
159        let mut ref_counts = self.ref_counts.write();
160        let dropped_entity_ids = mem::take(&mut ref_counts.dropped_entity_ids);
161        let mut accessed_entities = self.accessed_entities.borrow_mut();
162
163        dropped_entity_ids
164            .into_iter()
165            .filter_map(|entity_id| {
166                let count = ref_counts.counts.remove(entity_id).unwrap();
167                debug_assert_eq!(
168                    count.load(SeqCst),
169                    0,
170                    "dropped an entity that was referenced"
171                );
172                accessed_entities.remove(&entity_id);
173                // If the EntityId was allocated with `Context::reserve`,
174                // the entity may not have been inserted.
175                Some((entity_id, self.entities.remove(entity_id)?))
176            })
177            .collect()
178    }
179}
180
181#[track_caller]
182fn double_lease_panic<T>(operation: &str) -> ! {
183    panic!(
184        "cannot {operation} {} while it is already being updated",
185        std::any::type_name::<T>()
186    )
187}
188
189pub(crate) struct Lease<T> {
190    entity: Option<Box<dyn Any>>,
191    pub id: EntityId,
192    entity_type: PhantomData<T>,
193}
194
195impl<T: 'static> core::ops::Deref for Lease<T> {
196    type Target = T;
197
198    fn deref(&self) -> &Self::Target {
199        self.entity.as_ref().unwrap().downcast_ref().unwrap()
200    }
201}
202
203impl<T: 'static> core::ops::DerefMut for Lease<T> {
204    fn deref_mut(&mut self) -> &mut Self::Target {
205        self.entity.as_mut().unwrap().downcast_mut().unwrap()
206    }
207}
208
209impl<T> Drop for Lease<T> {
210    fn drop(&mut self) {
211        if self.entity.is_some() && !panicking() {
212            panic!("Leases must be ended with EntityMap::end_lease")
213        }
214    }
215}
216
217#[derive(Deref, DerefMut)]
218pub(crate) struct Slot<T>(Entity<T>);
219
220/// A dynamically typed reference to a entity, which can be downcast into a `Entity<T>`.
221pub struct AnyEntity {
222    pub(crate) entity_id: EntityId,
223    pub(crate) entity_type: TypeId,
224    entity_map: Weak<RwLock<EntityRefCounts>>,
225    #[cfg(any(test, feature = "leak-detection"))]
226    handle_id: HandleId,
227}
228
229impl AnyEntity {
230    fn new(id: EntityId, entity_type: TypeId, entity_map: Weak<RwLock<EntityRefCounts>>) -> Self {
231        Self {
232            entity_id: id,
233            entity_type,
234            #[cfg(any(test, feature = "leak-detection"))]
235            handle_id: entity_map
236                .clone()
237                .upgrade()
238                .unwrap()
239                .write()
240                .leak_detector
241                .handle_created(id),
242            entity_map,
243        }
244    }
245
246    /// Returns the id associated with this entity.
247    pub fn entity_id(&self) -> EntityId {
248        self.entity_id
249    }
250
251    /// Returns the [TypeId] associated with this entity.
252    pub fn entity_type(&self) -> TypeId {
253        self.entity_type
254    }
255
256    /// Converts this entity handle into a weak variant, which does not prevent it from being released.
257    pub fn downgrade(&self) -> AnyWeakEntity {
258        AnyWeakEntity {
259            entity_id: self.entity_id,
260            entity_type: self.entity_type,
261            entity_ref_counts: self.entity_map.clone(),
262        }
263    }
264
265    /// Converts this entity handle into a strongly-typed entity handle of the given type.
266    /// If this entity handle is not of the specified type, returns itself as an error variant.
267    pub fn downcast<T: 'static>(self) -> Result<Entity<T>, AnyEntity> {
268        if TypeId::of::<T>() == self.entity_type {
269            Ok(Entity {
270                any_entity: self,
271                entity_type: PhantomData,
272            })
273        } else {
274            Err(self)
275        }
276    }
277}
278
279impl Clone for AnyEntity {
280    fn clone(&self) -> Self {
281        if let Some(entity_map) = self.entity_map.upgrade() {
282            let entity_map = entity_map.read();
283            let count = entity_map
284                .counts
285                .get(self.entity_id)
286                .expect("detected over-release of a entity");
287            let prev_count = count.fetch_add(1, SeqCst);
288            assert_ne!(prev_count, 0, "Detected over-release of a entity.");
289        }
290
291        Self {
292            entity_id: self.entity_id,
293            entity_type: self.entity_type,
294            entity_map: self.entity_map.clone(),
295            #[cfg(any(test, feature = "leak-detection"))]
296            handle_id: self
297                .entity_map
298                .upgrade()
299                .unwrap()
300                .write()
301                .leak_detector
302                .handle_created(self.entity_id),
303        }
304    }
305}
306
307impl Drop for AnyEntity {
308    fn drop(&mut self) {
309        if let Some(entity_map) = self.entity_map.upgrade() {
310            let entity_map = entity_map.upgradable_read();
311            let count = entity_map
312                .counts
313                .get(self.entity_id)
314                .expect("detected over-release of a handle.");
315            let prev_count = count.fetch_sub(1, SeqCst);
316            assert_ne!(prev_count, 0, "Detected over-release of a entity.");
317            if prev_count == 1 {
318                // We were the last reference to this entity, so we can remove it.
319                let mut entity_map = RwLockUpgradableReadGuard::upgrade(entity_map);
320                entity_map.dropped_entity_ids.push(self.entity_id);
321            }
322        }
323
324        #[cfg(any(test, feature = "leak-detection"))]
325        if let Some(entity_map) = self.entity_map.upgrade() {
326            entity_map
327                .write()
328                .leak_detector
329                .handle_released(self.entity_id, self.handle_id)
330        }
331    }
332}
333
334impl<T> From<Entity<T>> for AnyEntity {
335    fn from(entity: Entity<T>) -> Self {
336        entity.any_entity
337    }
338}
339
340impl Hash for AnyEntity {
341    fn hash<H: Hasher>(&self, state: &mut H) {
342        self.entity_id.hash(state);
343    }
344}
345
346impl PartialEq for AnyEntity {
347    fn eq(&self, other: &Self) -> bool {
348        self.entity_id == other.entity_id
349    }
350}
351
352impl Eq for AnyEntity {}
353
354impl Ord for AnyEntity {
355    fn cmp(&self, other: &Self) -> Ordering {
356        self.entity_id.cmp(&other.entity_id)
357    }
358}
359
360impl PartialOrd for AnyEntity {
361    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
362        Some(self.cmp(other))
363    }
364}
365
366impl std::fmt::Debug for AnyEntity {
367    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
368        f.debug_struct("AnyEntity")
369            .field("entity_id", &self.entity_id.as_u64())
370            .finish()
371    }
372}
373
374/// A strong, well-typed reference to a struct which is managed
375/// by GPUI
376#[derive(Deref, DerefMut)]
377pub struct Entity<T> {
378    #[deref]
379    #[deref_mut]
380    pub(crate) any_entity: AnyEntity,
381    pub(crate) entity_type: PhantomData<T>,
382}
383
384unsafe impl<T> Send for Entity<T> {}
385unsafe impl<T> Sync for Entity<T> {}
386impl<T> Sealed for Entity<T> {}
387
388impl<T: 'static> Entity<T> {
389    fn new(id: EntityId, entity_map: Weak<RwLock<EntityRefCounts>>) -> Self
390    where
391        T: 'static,
392    {
393        Self {
394            any_entity: AnyEntity::new(id, TypeId::of::<T>(), entity_map),
395            entity_type: PhantomData,
396        }
397    }
398
399    /// Get the entity ID associated with this entity
400    pub fn entity_id(&self) -> EntityId {
401        self.any_entity.entity_id
402    }
403
404    /// Downgrade this entity pointer to a non-retaining weak pointer
405    pub fn downgrade(&self) -> WeakEntity<T> {
406        WeakEntity {
407            any_entity: self.any_entity.downgrade(),
408            entity_type: self.entity_type,
409        }
410    }
411
412    /// Convert this into a dynamically typed entity.
413    pub fn into_any(self) -> AnyEntity {
414        self.any_entity
415    }
416
417    /// Grab a reference to this entity from the context.
418    pub fn read<'a>(&self, cx: &'a App) -> &'a T {
419        cx.entities.read(self)
420    }
421
422    /// Read the entity referenced by this handle with the given function.
423    pub fn read_with<R, C: AppContext>(
424        &self,
425        cx: &C,
426        f: impl FnOnce(&T, &App) -> R,
427    ) -> C::Result<R> {
428        cx.read_entity(self, f)
429    }
430
431    /// Updates the entity referenced by this handle with the given function.
432    pub fn update<R, C: AppContext>(
433        &self,
434        cx: &mut C,
435        update: impl FnOnce(&mut T, &mut Context<T>) -> R,
436    ) -> C::Result<R> {
437        cx.update_entity(self, update)
438    }
439
440    /// Updates the entity referenced by this handle with the given function.
441    pub fn as_mut<'a, C: AppContext>(&self, cx: &'a mut C) -> C::Result<GpuiBorrow<'a, T>> {
442        cx.as_mut(self)
443    }
444
445    /// Updates the entity referenced by this handle with the given function.
446    pub fn write<C: AppContext>(&self, cx: &mut C, value: T) -> C::Result<()> {
447        self.update(cx, |entity, cx| {
448            *entity = value;
449            cx.notify();
450        })
451    }
452
453    /// Updates the entity referenced by this handle with the given function if
454    /// the referenced entity still exists, within a visual context that has a window.
455    /// Returns an error if the entity has been released.
456    pub fn update_in<R, C: VisualContext>(
457        &self,
458        cx: &mut C,
459        update: impl FnOnce(&mut T, &mut Window, &mut Context<T>) -> R,
460    ) -> C::Result<R> {
461        cx.update_window_entity(self, update)
462    }
463}
464
465impl<T> Clone for Entity<T> {
466    fn clone(&self) -> Self {
467        Self {
468            any_entity: self.any_entity.clone(),
469            entity_type: self.entity_type,
470        }
471    }
472}
473
474impl<T> std::fmt::Debug for Entity<T> {
475    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
476        f.debug_struct("Entity")
477            .field("entity_id", &self.any_entity.entity_id)
478            .field("entity_type", &type_name::<T>())
479            .finish()
480    }
481}
482
483impl<T> Hash for Entity<T> {
484    fn hash<H: Hasher>(&self, state: &mut H) {
485        self.any_entity.hash(state);
486    }
487}
488
489impl<T> PartialEq for Entity<T> {
490    fn eq(&self, other: &Self) -> bool {
491        self.any_entity == other.any_entity
492    }
493}
494
495impl<T> Eq for Entity<T> {}
496
497impl<T> PartialEq<WeakEntity<T>> for Entity<T> {
498    fn eq(&self, other: &WeakEntity<T>) -> bool {
499        self.any_entity.entity_id() == other.entity_id()
500    }
501}
502
503impl<T: 'static> Ord for Entity<T> {
504    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
505        self.entity_id().cmp(&other.entity_id())
506    }
507}
508
509impl<T: 'static> PartialOrd for Entity<T> {
510    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
511        Some(self.cmp(other))
512    }
513}
514
515/// A type erased, weak reference to a entity.
516#[derive(Clone)]
517pub struct AnyWeakEntity {
518    pub(crate) entity_id: EntityId,
519    entity_type: TypeId,
520    entity_ref_counts: Weak<RwLock<EntityRefCounts>>,
521}
522
523impl AnyWeakEntity {
524    /// Get the entity ID associated with this weak reference.
525    pub fn entity_id(&self) -> EntityId {
526        self.entity_id
527    }
528
529    /// Check if this weak handle can be upgraded, or if the entity has already been dropped
530    pub fn is_upgradable(&self) -> bool {
531        let ref_count = self
532            .entity_ref_counts
533            .upgrade()
534            .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
535            .unwrap_or(0);
536        ref_count > 0
537    }
538
539    /// Upgrade this weak entity reference to a strong reference.
540    pub fn upgrade(&self) -> Option<AnyEntity> {
541        let ref_counts = &self.entity_ref_counts.upgrade()?;
542        let ref_counts = ref_counts.read();
543        let ref_count = ref_counts.counts.get(self.entity_id)?;
544
545        if atomic_incr_if_not_zero(ref_count) == 0 {
546            // entity_id is in dropped_entity_ids
547            return None;
548        }
549        drop(ref_counts);
550
551        Some(AnyEntity {
552            entity_id: self.entity_id,
553            entity_type: self.entity_type,
554            entity_map: self.entity_ref_counts.clone(),
555            #[cfg(any(test, feature = "leak-detection"))]
556            handle_id: self
557                .entity_ref_counts
558                .upgrade()
559                .unwrap()
560                .write()
561                .leak_detector
562                .handle_created(self.entity_id),
563        })
564    }
565
566    /// Assert that entity referenced by this weak handle has been released.
567    #[cfg(any(test, feature = "leak-detection"))]
568    pub fn assert_released(&self) {
569        self.entity_ref_counts
570            .upgrade()
571            .unwrap()
572            .write()
573            .leak_detector
574            .assert_released(self.entity_id);
575
576        if self
577            .entity_ref_counts
578            .upgrade()
579            .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
580            .is_some()
581        {
582            panic!(
583                "entity was recently dropped but resources are retained until the end of the effect cycle."
584            )
585        }
586    }
587
588    /// Creates a weak entity that can never be upgraded.
589    pub fn new_invalid() -> Self {
590        /// To hold the invariant that all ids are unique, and considering that slotmap
591        /// increases their IDs from `0`, we can decrease ours from `u64::MAX` so these
592        /// two will never conflict (u64 is way too large).
593        static UNIQUE_NON_CONFLICTING_ID_GENERATOR: AtomicU64 = AtomicU64::new(u64::MAX);
594        let entity_id = UNIQUE_NON_CONFLICTING_ID_GENERATOR.fetch_sub(1, SeqCst);
595
596        Self {
597            // Safety:
598            //   Docs say this is safe but can be unspecified if slotmap changes the representation
599            //   after `1.0.7`, that said, providing a valid entity_id here is not necessary as long
600            //   as we guarantee that `entity_id` is never used if `entity_ref_counts` equals
601            //   to `Weak::new()` (that is, it's unable to upgrade), that is the invariant that
602            //   actually needs to be hold true.
603            //
604            //   And there is no sane reason to read an entity slot if `entity_ref_counts` can't be
605            //   read in the first place, so we're good!
606            entity_id: entity_id.into(),
607            entity_type: TypeId::of::<()>(),
608            entity_ref_counts: Weak::new(),
609        }
610    }
611}
612
613impl std::fmt::Debug for AnyWeakEntity {
614    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
615        f.debug_struct(type_name::<Self>())
616            .field("entity_id", &self.entity_id)
617            .field("entity_type", &self.entity_type)
618            .finish()
619    }
620}
621
622impl<T> From<WeakEntity<T>> for AnyWeakEntity {
623    fn from(entity: WeakEntity<T>) -> Self {
624        entity.any_entity
625    }
626}
627
628impl Hash for AnyWeakEntity {
629    fn hash<H: Hasher>(&self, state: &mut H) {
630        self.entity_id.hash(state);
631    }
632}
633
634impl PartialEq for AnyWeakEntity {
635    fn eq(&self, other: &Self) -> bool {
636        self.entity_id == other.entity_id
637    }
638}
639
640impl Eq for AnyWeakEntity {}
641
642impl Ord for AnyWeakEntity {
643    fn cmp(&self, other: &Self) -> Ordering {
644        self.entity_id.cmp(&other.entity_id)
645    }
646}
647
648impl PartialOrd for AnyWeakEntity {
649    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
650        Some(self.cmp(other))
651    }
652}
653
654/// A weak reference to a entity of the given type.
655#[derive(Deref, DerefMut)]
656pub struct WeakEntity<T> {
657    #[deref]
658    #[deref_mut]
659    any_entity: AnyWeakEntity,
660    entity_type: PhantomData<T>,
661}
662
663impl<T> std::fmt::Debug for WeakEntity<T> {
664    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
665        f.debug_struct(type_name::<Self>())
666            .field("entity_id", &self.any_entity.entity_id)
667            .field("entity_type", &type_name::<T>())
668            .finish()
669    }
670}
671
672unsafe impl<T> Send for WeakEntity<T> {}
673unsafe impl<T> Sync for WeakEntity<T> {}
674
675impl<T> Clone for WeakEntity<T> {
676    fn clone(&self) -> Self {
677        Self {
678            any_entity: self.any_entity.clone(),
679            entity_type: self.entity_type,
680        }
681    }
682}
683
684impl<T: 'static> WeakEntity<T> {
685    /// Upgrade this weak entity reference into a strong entity reference
686    pub fn upgrade(&self) -> Option<Entity<T>> {
687        Some(Entity {
688            any_entity: self.any_entity.upgrade()?,
689            entity_type: self.entity_type,
690        })
691    }
692
693    /// Updates the entity referenced by this handle with the given function if
694    /// the referenced entity still exists. Returns an error if the entity has
695    /// been released.
696    pub fn update<C, R>(
697        &self,
698        cx: &mut C,
699        update: impl FnOnce(&mut T, &mut Context<T>) -> R,
700    ) -> Result<R>
701    where
702        C: AppContext,
703        Result<C::Result<R>>: crate::Flatten<R>,
704    {
705        crate::Flatten::flatten(
706            self.upgrade()
707                .context("entity released")
708                .map(|this| cx.update_entity(&this, update)),
709        )
710    }
711
712    /// Updates the entity referenced by this handle with the given function if
713    /// the referenced entity still exists, within a visual context that has a window.
714    /// Returns an error if the entity has been released.
715    pub fn update_in<C, R>(
716        &self,
717        cx: &mut C,
718        update: impl FnOnce(&mut T, &mut Window, &mut Context<T>) -> R,
719    ) -> Result<R>
720    where
721        C: VisualContext,
722        Result<C::Result<R>>: crate::Flatten<R>,
723    {
724        let window = cx.window_handle();
725        let this = self.upgrade().context("entity released")?;
726
727        crate::Flatten::flatten(window.update(cx, |_, window, cx| {
728            this.update(cx, |entity, cx| update(entity, window, cx))
729        }))
730    }
731
732    /// Reads the entity referenced by this handle with the given function if
733    /// the referenced entity still exists. Returns an error if the entity has
734    /// been released.
735    pub fn read_with<C, R>(&self, cx: &C, read: impl FnOnce(&T, &App) -> R) -> Result<R>
736    where
737        C: AppContext,
738        Result<C::Result<R>>: crate::Flatten<R>,
739    {
740        crate::Flatten::flatten(
741            self.upgrade()
742                .context("entity released")
743                .map(|this| cx.read_entity(&this, read)),
744        )
745    }
746
747    /// Create a new weak entity that can never be upgraded.
748    pub fn new_invalid() -> Self {
749        Self {
750            any_entity: AnyWeakEntity::new_invalid(),
751            entity_type: PhantomData,
752        }
753    }
754}
755
756impl<T> Hash for WeakEntity<T> {
757    fn hash<H: Hasher>(&self, state: &mut H) {
758        self.any_entity.hash(state);
759    }
760}
761
762impl<T> PartialEq for WeakEntity<T> {
763    fn eq(&self, other: &Self) -> bool {
764        self.any_entity == other.any_entity
765    }
766}
767
768impl<T> Eq for WeakEntity<T> {}
769
770impl<T> PartialEq<Entity<T>> for WeakEntity<T> {
771    fn eq(&self, other: &Entity<T>) -> bool {
772        self.entity_id() == other.any_entity.entity_id()
773    }
774}
775
776impl<T: 'static> Ord for WeakEntity<T> {
777    fn cmp(&self, other: &Self) -> Ordering {
778        self.entity_id().cmp(&other.entity_id())
779    }
780}
781
782impl<T: 'static> PartialOrd for WeakEntity<T> {
783    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
784        Some(self.cmp(other))
785    }
786}
787
788#[cfg(any(test, feature = "leak-detection"))]
789static LEAK_BACKTRACE: std::sync::LazyLock<bool> =
790    std::sync::LazyLock::new(|| std::env::var("LEAK_BACKTRACE").is_ok_and(|b| !b.is_empty()));
791
792#[cfg(any(test, feature = "leak-detection"))]
793#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq)]
794pub(crate) struct HandleId {
795    id: u64, // id of the handle itself, not the pointed at object
796}
797
798#[cfg(any(test, feature = "leak-detection"))]
799pub(crate) struct LeakDetector {
800    next_handle_id: u64,
801    entity_handles: HashMap<EntityId, HashMap<HandleId, Option<backtrace::Backtrace>>>,
802}
803
804#[cfg(any(test, feature = "leak-detection"))]
805impl LeakDetector {
806    #[track_caller]
807    pub fn handle_created(&mut self, entity_id: EntityId) -> HandleId {
808        let id = util::post_inc(&mut self.next_handle_id);
809        let handle_id = HandleId { id };
810        let handles = self.entity_handles.entry(entity_id).or_default();
811        handles.insert(
812            handle_id,
813            LEAK_BACKTRACE.then(backtrace::Backtrace::new_unresolved),
814        );
815        handle_id
816    }
817
818    pub fn handle_released(&mut self, entity_id: EntityId, handle_id: HandleId) {
819        let handles = self.entity_handles.entry(entity_id).or_default();
820        handles.remove(&handle_id);
821    }
822
823    pub fn assert_released(&mut self, entity_id: EntityId) {
824        let handles = self.entity_handles.entry(entity_id).or_default();
825        if !handles.is_empty() {
826            for backtrace in handles.values_mut() {
827                if let Some(mut backtrace) = backtrace.take() {
828                    backtrace.resolve();
829                    eprintln!("Leaked handle: {:#?}", backtrace);
830                } else {
831                    eprintln!("Leaked handle: export LEAK_BACKTRACE to find allocation site");
832                }
833            }
834            panic!();
835        }
836    }
837}
838
839#[cfg(test)]
840mod test {
841    use crate::EntityMap;
842
843    struct TestEntity {
844        pub i: i32,
845    }
846
847    #[test]
848    fn test_entity_map_slot_assignment_before_cleanup() {
849        // Tests that slots are not re-used before take_dropped.
850        let mut entity_map = EntityMap::new();
851
852        let slot = entity_map.reserve::<TestEntity>();
853        entity_map.insert(slot, TestEntity { i: 1 });
854
855        let slot = entity_map.reserve::<TestEntity>();
856        entity_map.insert(slot, TestEntity { i: 2 });
857
858        let dropped = entity_map.take_dropped();
859        assert_eq!(dropped.len(), 2);
860
861        assert_eq!(
862            dropped
863                .into_iter()
864                .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
865                .collect::<Vec<i32>>(),
866            vec![1, 2],
867        );
868    }
869
870    #[test]
871    fn test_entity_map_weak_upgrade_before_cleanup() {
872        // Tests that weak handles are not upgraded before take_dropped
873        let mut entity_map = EntityMap::new();
874
875        let slot = entity_map.reserve::<TestEntity>();
876        let handle = entity_map.insert(slot, TestEntity { i: 1 });
877        let weak = handle.downgrade();
878        drop(handle);
879
880        let strong = weak.upgrade();
881        assert_eq!(strong, None);
882
883        let dropped = entity_map.take_dropped();
884        assert_eq!(dropped.len(), 1);
885
886        assert_eq!(
887            dropped
888                .into_iter()
889                .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
890                .collect::<Vec<i32>>(),
891            vec![1],
892        );
893    }
894}