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            entity_map: entity_map.clone(),
235            #[cfg(any(test, feature = "leak-detection"))]
236            handle_id: entity_map
237                .upgrade()
238                .unwrap()
239                .write()
240                .leak_detector
241                .handle_created(id),
242        }
243    }
244
245    /// Returns the id associated with this entity.
246    pub fn entity_id(&self) -> EntityId {
247        self.entity_id
248    }
249
250    /// Returns the [TypeId] associated with this entity.
251    pub fn entity_type(&self) -> TypeId {
252        self.entity_type
253    }
254
255    /// Converts this entity handle into a weak variant, which does not prevent it from being released.
256    pub fn downgrade(&self) -> AnyWeakEntity {
257        AnyWeakEntity {
258            entity_id: self.entity_id,
259            entity_type: self.entity_type,
260            entity_ref_counts: self.entity_map.clone(),
261        }
262    }
263
264    /// Converts this entity handle into a strongly-typed entity handle of the given type.
265    /// If this entity handle is not of the specified type, returns itself as an error variant.
266    pub fn downcast<T: 'static>(self) -> Result<Entity<T>, AnyEntity> {
267        if TypeId::of::<T>() == self.entity_type {
268            Ok(Entity {
269                any_entity: self,
270                entity_type: PhantomData,
271            })
272        } else {
273            Err(self)
274        }
275    }
276}
277
278impl Clone for AnyEntity {
279    fn clone(&self) -> Self {
280        if let Some(entity_map) = self.entity_map.upgrade() {
281            let entity_map = entity_map.read();
282            let count = entity_map
283                .counts
284                .get(self.entity_id)
285                .expect("detected over-release of a entity");
286            let prev_count = count.fetch_add(1, SeqCst);
287            assert_ne!(prev_count, 0, "Detected over-release of a entity.");
288        }
289
290        Self {
291            entity_id: self.entity_id,
292            entity_type: self.entity_type,
293            entity_map: self.entity_map.clone(),
294            #[cfg(any(test, feature = "leak-detection"))]
295            handle_id: self
296                .entity_map
297                .upgrade()
298                .unwrap()
299                .write()
300                .leak_detector
301                .handle_created(self.entity_id),
302        }
303    }
304}
305
306impl Drop for AnyEntity {
307    fn drop(&mut self) {
308        if let Some(entity_map) = self.entity_map.upgrade() {
309            let entity_map = entity_map.upgradable_read();
310            let count = entity_map
311                .counts
312                .get(self.entity_id)
313                .expect("detected over-release of a handle.");
314            let prev_count = count.fetch_sub(1, SeqCst);
315            assert_ne!(prev_count, 0, "Detected over-release of a entity.");
316            if prev_count == 1 {
317                // We were the last reference to this entity, so we can remove it.
318                let mut entity_map = RwLockUpgradableReadGuard::upgrade(entity_map);
319                entity_map.dropped_entity_ids.push(self.entity_id);
320            }
321        }
322
323        #[cfg(any(test, feature = "leak-detection"))]
324        if let Some(entity_map) = self.entity_map.upgrade() {
325            entity_map
326                .write()
327                .leak_detector
328                .handle_released(self.entity_id, self.handle_id)
329        }
330    }
331}
332
333impl<T> From<Entity<T>> for AnyEntity {
334    fn from(entity: Entity<T>) -> Self {
335        entity.any_entity
336    }
337}
338
339impl Hash for AnyEntity {
340    fn hash<H: Hasher>(&self, state: &mut H) {
341        self.entity_id.hash(state);
342    }
343}
344
345impl PartialEq for AnyEntity {
346    fn eq(&self, other: &Self) -> bool {
347        self.entity_id == other.entity_id
348    }
349}
350
351impl Eq for AnyEntity {}
352
353impl Ord for AnyEntity {
354    fn cmp(&self, other: &Self) -> Ordering {
355        self.entity_id.cmp(&other.entity_id)
356    }
357}
358
359impl PartialOrd for AnyEntity {
360    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
361        Some(self.cmp(other))
362    }
363}
364
365impl std::fmt::Debug for AnyEntity {
366    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
367        f.debug_struct("AnyEntity")
368            .field("entity_id", &self.entity_id.as_u64())
369            .finish()
370    }
371}
372
373/// A strong, well-typed reference to a struct which is managed
374/// by GPUI
375#[derive(Deref, DerefMut)]
376pub struct Entity<T> {
377    #[deref]
378    #[deref_mut]
379    pub(crate) any_entity: AnyEntity,
380    pub(crate) entity_type: PhantomData<T>,
381}
382
383unsafe impl<T> Send for Entity<T> {}
384unsafe impl<T> Sync for Entity<T> {}
385impl<T> Sealed for Entity<T> {}
386
387impl<T: 'static> Entity<T> {
388    fn new(id: EntityId, entity_map: Weak<RwLock<EntityRefCounts>>) -> Self
389    where
390        T: 'static,
391    {
392        Self {
393            any_entity: AnyEntity::new(id, TypeId::of::<T>(), entity_map),
394            entity_type: PhantomData,
395        }
396    }
397
398    /// Get the entity ID associated with this entity
399    pub fn entity_id(&self) -> EntityId {
400        self.any_entity.entity_id
401    }
402
403    /// Downgrade this entity pointer to a non-retaining weak pointer
404    pub fn downgrade(&self) -> WeakEntity<T> {
405        WeakEntity {
406            any_entity: self.any_entity.downgrade(),
407            entity_type: self.entity_type,
408        }
409    }
410
411    /// Convert this into a dynamically typed entity.
412    pub fn into_any(self) -> AnyEntity {
413        self.any_entity
414    }
415
416    /// Grab a reference to this entity from the context.
417    pub fn read<'a>(&self, cx: &'a App) -> &'a T {
418        cx.entities.read(self)
419    }
420
421    /// Read the entity referenced by this handle with the given function.
422    pub fn read_with<R, C: AppContext>(
423        &self,
424        cx: &C,
425        f: impl FnOnce(&T, &App) -> R,
426    ) -> C::Result<R> {
427        cx.read_entity(self, f)
428    }
429
430    /// Updates the entity referenced by this handle with the given function.
431    pub fn update<R, C: AppContext>(
432        &self,
433        cx: &mut C,
434        update: impl FnOnce(&mut T, &mut Context<T>) -> R,
435    ) -> C::Result<R> {
436        cx.update_entity(self, update)
437    }
438
439    /// Updates the entity referenced by this handle with the given function.
440    pub fn as_mut<'a, C: AppContext>(&self, cx: &'a mut C) -> C::Result<GpuiBorrow<'a, T>> {
441        cx.as_mut(self)
442    }
443
444    /// Updates the entity referenced by this handle with the given function.
445    pub fn write<C: AppContext>(&self, cx: &mut C, value: T) -> C::Result<()> {
446        self.update(cx, |entity, cx| {
447            *entity = value;
448            cx.notify();
449        })
450    }
451
452    /// Updates the entity referenced by this handle with the given function if
453    /// the referenced entity still exists, within a visual context that has a window.
454    /// Returns an error if the entity has been released.
455    pub fn update_in<R, C: VisualContext>(
456        &self,
457        cx: &mut C,
458        update: impl FnOnce(&mut T, &mut Window, &mut Context<T>) -> R,
459    ) -> C::Result<R> {
460        cx.update_window_entity(self, update)
461    }
462}
463
464impl<T> Clone for Entity<T> {
465    fn clone(&self) -> Self {
466        Self {
467            any_entity: self.any_entity.clone(),
468            entity_type: self.entity_type,
469        }
470    }
471}
472
473impl<T> std::fmt::Debug for Entity<T> {
474    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
475        f.debug_struct("Entity")
476            .field("entity_id", &self.any_entity.entity_id)
477            .field("entity_type", &type_name::<T>())
478            .finish()
479    }
480}
481
482impl<T> Hash for Entity<T> {
483    fn hash<H: Hasher>(&self, state: &mut H) {
484        self.any_entity.hash(state);
485    }
486}
487
488impl<T> PartialEq for Entity<T> {
489    fn eq(&self, other: &Self) -> bool {
490        self.any_entity == other.any_entity
491    }
492}
493
494impl<T> Eq for Entity<T> {}
495
496impl<T> PartialEq<WeakEntity<T>> for Entity<T> {
497    fn eq(&self, other: &WeakEntity<T>) -> bool {
498        self.any_entity.entity_id() == other.entity_id()
499    }
500}
501
502impl<T: 'static> Ord for Entity<T> {
503    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
504        self.entity_id().cmp(&other.entity_id())
505    }
506}
507
508impl<T: 'static> PartialOrd for Entity<T> {
509    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
510        Some(self.cmp(other))
511    }
512}
513
514/// A type erased, weak reference to a entity.
515#[derive(Clone)]
516pub struct AnyWeakEntity {
517    pub(crate) entity_id: EntityId,
518    entity_type: TypeId,
519    entity_ref_counts: Weak<RwLock<EntityRefCounts>>,
520}
521
522impl AnyWeakEntity {
523    /// Get the entity ID associated with this weak reference.
524    pub fn entity_id(&self) -> EntityId {
525        self.entity_id
526    }
527
528    /// Check if this weak handle can be upgraded, or if the entity has already been dropped
529    pub fn is_upgradable(&self) -> bool {
530        let ref_count = self
531            .entity_ref_counts
532            .upgrade()
533            .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
534            .unwrap_or(0);
535        ref_count > 0
536    }
537
538    /// Upgrade this weak entity reference to a strong reference.
539    pub fn upgrade(&self) -> Option<AnyEntity> {
540        let ref_counts = &self.entity_ref_counts.upgrade()?;
541        let ref_counts = ref_counts.read();
542        let ref_count = ref_counts.counts.get(self.entity_id)?;
543
544        if atomic_incr_if_not_zero(ref_count) == 0 {
545            // entity_id is in dropped_entity_ids
546            return None;
547        }
548        drop(ref_counts);
549
550        Some(AnyEntity {
551            entity_id: self.entity_id,
552            entity_type: self.entity_type,
553            entity_map: self.entity_ref_counts.clone(),
554            #[cfg(any(test, feature = "leak-detection"))]
555            handle_id: self
556                .entity_ref_counts
557                .upgrade()
558                .unwrap()
559                .write()
560                .leak_detector
561                .handle_created(self.entity_id),
562        })
563    }
564
565    /// Assert that entity referenced by this weak handle has been released.
566    #[cfg(any(test, feature = "leak-detection"))]
567    pub fn assert_released(&self) {
568        self.entity_ref_counts
569            .upgrade()
570            .unwrap()
571            .write()
572            .leak_detector
573            .assert_released(self.entity_id);
574
575        if self
576            .entity_ref_counts
577            .upgrade()
578            .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
579            .is_some()
580        {
581            panic!(
582                "entity was recently dropped but resources are retained until the end of the effect cycle."
583            )
584        }
585    }
586
587    /// Creates a weak entity that can never be upgraded.
588    pub fn new_invalid() -> Self {
589        /// To hold the invariant that all ids are unique, and considering that slotmap
590        /// increases their IDs from `0`, we can decrease ours from `u64::MAX` so these
591        /// two will never conflict (u64 is way too large).
592        static UNIQUE_NON_CONFLICTING_ID_GENERATOR: AtomicU64 = AtomicU64::new(u64::MAX);
593        let entity_id = UNIQUE_NON_CONFLICTING_ID_GENERATOR.fetch_sub(1, SeqCst);
594
595        Self {
596            // Safety:
597            //   Docs say this is safe but can be unspecified if slotmap changes the representation
598            //   after `1.0.7`, that said, providing a valid entity_id here is not necessary as long
599            //   as we guarantee that `entity_id` is never used if `entity_ref_counts` equals
600            //   to `Weak::new()` (that is, it's unable to upgrade), that is the invariant that
601            //   actually needs to be hold true.
602            //
603            //   And there is no sane reason to read an entity slot if `entity_ref_counts` can't be
604            //   read in the first place, so we're good!
605            entity_id: entity_id.into(),
606            entity_type: TypeId::of::<()>(),
607            entity_ref_counts: Weak::new(),
608        }
609    }
610}
611
612impl std::fmt::Debug for AnyWeakEntity {
613    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
614        f.debug_struct(type_name::<Self>())
615            .field("entity_id", &self.entity_id)
616            .field("entity_type", &self.entity_type)
617            .finish()
618    }
619}
620
621impl<T> From<WeakEntity<T>> for AnyWeakEntity {
622    fn from(entity: WeakEntity<T>) -> Self {
623        entity.any_entity
624    }
625}
626
627impl Hash for AnyWeakEntity {
628    fn hash<H: Hasher>(&self, state: &mut H) {
629        self.entity_id.hash(state);
630    }
631}
632
633impl PartialEq for AnyWeakEntity {
634    fn eq(&self, other: &Self) -> bool {
635        self.entity_id == other.entity_id
636    }
637}
638
639impl Eq for AnyWeakEntity {}
640
641impl Ord for AnyWeakEntity {
642    fn cmp(&self, other: &Self) -> Ordering {
643        self.entity_id.cmp(&other.entity_id)
644    }
645}
646
647impl PartialOrd for AnyWeakEntity {
648    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
649        Some(self.cmp(other))
650    }
651}
652
653/// A weak reference to a entity of the given type.
654#[derive(Deref, DerefMut)]
655pub struct WeakEntity<T> {
656    #[deref]
657    #[deref_mut]
658    any_entity: AnyWeakEntity,
659    entity_type: PhantomData<T>,
660}
661
662impl<T> std::fmt::Debug for WeakEntity<T> {
663    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
664        f.debug_struct(type_name::<Self>())
665            .field("entity_id", &self.any_entity.entity_id)
666            .field("entity_type", &type_name::<T>())
667            .finish()
668    }
669}
670
671unsafe impl<T> Send for WeakEntity<T> {}
672unsafe impl<T> Sync for WeakEntity<T> {}
673
674impl<T> Clone for WeakEntity<T> {
675    fn clone(&self) -> Self {
676        Self {
677            any_entity: self.any_entity.clone(),
678            entity_type: self.entity_type,
679        }
680    }
681}
682
683impl<T: 'static> WeakEntity<T> {
684    /// Upgrade this weak entity reference into a strong entity reference
685    pub fn upgrade(&self) -> Option<Entity<T>> {
686        Some(Entity {
687            any_entity: self.any_entity.upgrade()?,
688            entity_type: self.entity_type,
689        })
690    }
691
692    /// Updates the entity referenced by this handle with the given function if
693    /// the referenced entity still exists. Returns an error if the entity has
694    /// been released.
695    pub fn update<C, R>(
696        &self,
697        cx: &mut C,
698        update: impl FnOnce(&mut T, &mut Context<T>) -> R,
699    ) -> Result<R>
700    where
701        C: AppContext,
702        Result<C::Result<R>>: crate::Flatten<R>,
703    {
704        crate::Flatten::flatten(
705            self.upgrade()
706                .context("entity released")
707                .map(|this| cx.update_entity(&this, update)),
708        )
709    }
710
711    /// Updates the entity referenced by this handle with the given function if
712    /// the referenced entity still exists, within a visual context that has a window.
713    /// Returns an error if the entity has been released.
714    pub fn update_in<C, R>(
715        &self,
716        cx: &mut C,
717        update: impl FnOnce(&mut T, &mut Window, &mut Context<T>) -> R,
718    ) -> Result<R>
719    where
720        C: VisualContext,
721        Result<C::Result<R>>: crate::Flatten<R>,
722    {
723        let window = cx.window_handle();
724        let this = self.upgrade().context("entity released")?;
725
726        crate::Flatten::flatten(window.update(cx, |_, window, cx| {
727            this.update(cx, |entity, cx| update(entity, window, cx))
728        }))
729    }
730
731    /// Reads the entity referenced by this handle with the given function if
732    /// the referenced entity still exists. Returns an error if the entity has
733    /// been released.
734    pub fn read_with<C, R>(&self, cx: &C, read: impl FnOnce(&T, &App) -> R) -> Result<R>
735    where
736        C: AppContext,
737        Result<C::Result<R>>: crate::Flatten<R>,
738    {
739        crate::Flatten::flatten(
740            self.upgrade()
741                .context("entity released")
742                .map(|this| cx.read_entity(&this, read)),
743        )
744    }
745
746    /// Create a new weak entity that can never be upgraded.
747    pub fn new_invalid() -> Self {
748        Self {
749            any_entity: AnyWeakEntity::new_invalid(),
750            entity_type: PhantomData,
751        }
752    }
753}
754
755impl<T> Hash for WeakEntity<T> {
756    fn hash<H: Hasher>(&self, state: &mut H) {
757        self.any_entity.hash(state);
758    }
759}
760
761impl<T> PartialEq for WeakEntity<T> {
762    fn eq(&self, other: &Self) -> bool {
763        self.any_entity == other.any_entity
764    }
765}
766
767impl<T> Eq for WeakEntity<T> {}
768
769impl<T> PartialEq<Entity<T>> for WeakEntity<T> {
770    fn eq(&self, other: &Entity<T>) -> bool {
771        self.entity_id() == other.any_entity.entity_id()
772    }
773}
774
775impl<T: 'static> Ord for WeakEntity<T> {
776    fn cmp(&self, other: &Self) -> Ordering {
777        self.entity_id().cmp(&other.entity_id())
778    }
779}
780
781impl<T: 'static> PartialOrd for WeakEntity<T> {
782    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
783        Some(self.cmp(other))
784    }
785}
786
787#[cfg(any(test, feature = "leak-detection"))]
788static LEAK_BACKTRACE: std::sync::LazyLock<bool> =
789    std::sync::LazyLock::new(|| std::env::var("LEAK_BACKTRACE").is_ok_and(|b| !b.is_empty()));
790
791#[cfg(any(test, feature = "leak-detection"))]
792#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq)]
793pub(crate) struct HandleId {
794    id: u64, // id of the handle itself, not the pointed at object
795}
796
797#[cfg(any(test, feature = "leak-detection"))]
798pub(crate) struct LeakDetector {
799    next_handle_id: u64,
800    entity_handles: HashMap<EntityId, HashMap<HandleId, Option<backtrace::Backtrace>>>,
801}
802
803#[cfg(any(test, feature = "leak-detection"))]
804impl LeakDetector {
805    #[track_caller]
806    pub fn handle_created(&mut self, entity_id: EntityId) -> HandleId {
807        let id = util::post_inc(&mut self.next_handle_id);
808        let handle_id = HandleId { id };
809        let handles = self.entity_handles.entry(entity_id).or_default();
810        handles.insert(
811            handle_id,
812            LEAK_BACKTRACE.then(backtrace::Backtrace::new_unresolved),
813        );
814        handle_id
815    }
816
817    pub fn handle_released(&mut self, entity_id: EntityId, handle_id: HandleId) {
818        let handles = self.entity_handles.entry(entity_id).or_default();
819        handles.remove(&handle_id);
820    }
821
822    pub fn assert_released(&mut self, entity_id: EntityId) {
823        let handles = self.entity_handles.entry(entity_id).or_default();
824        if !handles.is_empty() {
825            for backtrace in handles.values_mut() {
826                if let Some(mut backtrace) = backtrace.take() {
827                    backtrace.resolve();
828                    eprintln!("Leaked handle: {:#?}", backtrace);
829                } else {
830                    eprintln!("Leaked handle: export LEAK_BACKTRACE to find allocation site");
831                }
832            }
833            panic!();
834        }
835    }
836}
837
838#[cfg(test)]
839mod test {
840    use crate::EntityMap;
841
842    struct TestEntity {
843        pub i: i32,
844    }
845
846    #[test]
847    fn test_entity_map_slot_assignment_before_cleanup() {
848        // Tests that slots are not re-used before take_dropped.
849        let mut entity_map = EntityMap::new();
850
851        let slot = entity_map.reserve::<TestEntity>();
852        entity_map.insert(slot, TestEntity { i: 1 });
853
854        let slot = entity_map.reserve::<TestEntity>();
855        entity_map.insert(slot, TestEntity { i: 2 });
856
857        let dropped = entity_map.take_dropped();
858        assert_eq!(dropped.len(), 2);
859
860        assert_eq!(
861            dropped
862                .into_iter()
863                .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
864                .collect::<Vec<i32>>(),
865            vec![1, 2],
866        );
867    }
868
869    #[test]
870    fn test_entity_map_weak_upgrade_before_cleanup() {
871        // Tests that weak handles are not upgraded before take_dropped
872        let mut entity_map = EntityMap::new();
873
874        let slot = entity_map.reserve::<TestEntity>();
875        let handle = entity_map.insert(slot, TestEntity { i: 1 });
876        let weak = handle.downgrade();
877        drop(handle);
878
879        let strong = weak.upgrade();
880        assert_eq!(strong, None);
881
882        let dropped = entity_map.take_dropped();
883        assert_eq!(dropped.len(), 1);
884
885        assert_eq!(
886            dropped
887                .into_iter()
888                .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
889                .collect::<Vec<i32>>(),
890            vec![1],
891        );
892    }
893}