entity_map.rs

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