gpui_core/app/
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, gpui_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, gpui_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, gpui_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, gpui_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, gpui_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, gpui_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, gpui_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, gpui_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 the entity referenced by this weak handle has been released.
567    ///
568    /// With leak detection enabled (build with `GPUI_LEAK_DETECTION=1`), a failure
569    /// lists each leaked handle. Also set the `LEAK_BACKTRACE` environment variable
570    /// to see where they were allocated.
571    ///
572    /// # Panics
573    ///
574    /// Panics if any strong handles to the entity are still alive, or if the entity
575    /// was recently dropped but cleanup has not completed yet.
576    #[cfg(any(test, feature = "test-support", gpui_leak_detection))]
577    pub fn assert_released(&self) {
578        #[cfg(any(test, gpui_leak_detection))]
579        self.entity_ref_counts
580            .upgrade()
581            .unwrap()
582            .write()
583            .leak_detector
584            .assert_released(self.entity_id);
585
586        if self
587            .entity_ref_counts
588            .upgrade()
589            .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
590            .is_some()
591        {
592            panic!(
593                "entity is still alive, or was recently dropped and its resources are retained \
594                 until the end of the effect cycle. Build with GPUI_LEAK_DETECTION=1 to list \
595                 its live handles."
596            )
597        }
598    }
599
600    /// Creates a weak entity that can never be upgraded.
601    pub fn new_invalid() -> Self {
602        /// To hold the invariant that all ids are unique, and considering that slotmap
603        /// increases their IDs from `0`, we can decrease ours from `u64::MAX` so these
604        /// two will never conflict (u64 is way too large).
605        static UNIQUE_NON_CONFLICTING_ID_GENERATOR: AtomicU64 = AtomicU64::new(u64::MAX);
606        let entity_id = UNIQUE_NON_CONFLICTING_ID_GENERATOR.fetch_sub(1, SeqCst);
607
608        Self {
609            // Safety:
610            //   Docs say this is safe but can be unspecified if slotmap changes the representation
611            //   after `1.0.7`, that said, providing a valid entity_id here is not necessary as long
612            //   as we guarantee that `entity_id` is never used if `entity_ref_counts` equals
613            //   to `Weak::new()` (that is, it's unable to upgrade), that is the invariant that
614            //   actually needs to be hold true.
615            //
616            //   And there is no sane reason to read an entity slot if `entity_ref_counts` can't be
617            //   read in the first place, so we're good!
618            entity_id: entity_id.into(),
619            entity_type: TypeId::of::<()>(),
620            entity_ref_counts: Weak::new(),
621        }
622    }
623}
624
625impl std::fmt::Debug for AnyWeakEntity {
626    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
627        f.debug_struct(type_name::<Self>())
628            .field("entity_id", &self.entity_id)
629            .field("entity_type", &self.entity_type)
630            .finish()
631    }
632}
633
634impl<T> From<WeakEntity<T>> for AnyWeakEntity {
635    fn from(entity: WeakEntity<T>) -> Self {
636        entity.any_entity
637    }
638}
639
640impl Hash for AnyWeakEntity {
641    fn hash<H: Hasher>(&self, state: &mut H) {
642        self.entity_id.hash(state);
643    }
644}
645
646impl PartialEq for AnyWeakEntity {
647    fn eq(&self, other: &Self) -> bool {
648        self.entity_id == other.entity_id
649    }
650}
651
652impl Eq for AnyWeakEntity {}
653
654impl Ord for AnyWeakEntity {
655    fn cmp(&self, other: &Self) -> Ordering {
656        self.entity_id.cmp(&other.entity_id)
657    }
658}
659
660impl PartialOrd for AnyWeakEntity {
661    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
662        Some(self.cmp(other))
663    }
664}
665
666/// A weak reference to a entity of the given type.
667#[derive(Deref, DerefMut)]
668pub struct WeakEntity<T> {
669    #[deref]
670    #[deref_mut]
671    any_entity: AnyWeakEntity,
672    entity_type: PhantomData<T>,
673}
674
675impl<T> std::fmt::Debug for WeakEntity<T> {
676    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
677        f.debug_struct(type_name::<Self>())
678            .field("entity_id", &self.any_entity.entity_id)
679            .field("entity_type", &type_name::<T>())
680            .finish()
681    }
682}
683
684unsafe impl<T> Send for WeakEntity<T> {}
685unsafe impl<T> Sync for WeakEntity<T> {}
686
687impl<T> Clone for WeakEntity<T> {
688    fn clone(&self) -> Self {
689        Self {
690            any_entity: self.any_entity.clone(),
691            entity_type: self.entity_type,
692        }
693    }
694}
695
696impl<T: 'static> WeakEntity<T> {
697    /// Upgrade this weak entity reference into a strong entity reference
698    pub fn upgrade(&self) -> Option<Entity<T>> {
699        Some(Entity {
700            any_entity: self.any_entity.upgrade()?,
701            entity_type: self.entity_type,
702        })
703    }
704
705    /// Updates the entity referenced by this handle with the given function if
706    /// the referenced entity still exists. Returns an error if the entity has
707    /// been released.
708    pub fn update<C, R>(
709        &self,
710        cx: &mut C,
711        update: impl FnOnce(&mut T, &mut Context<T>) -> R,
712    ) -> Result<R>
713    where
714        C: AppContext,
715        Result<C::Result<R>>: crate::Flatten<R>,
716    {
717        crate::Flatten::flatten(
718            self.upgrade()
719                .context("entity released")
720                .map(|this| cx.update_entity(&this, update)),
721        )
722    }
723
724    /// Updates the entity referenced by this handle with the given function if
725    /// the referenced entity still exists, within a visual context that has a window.
726    /// Returns an error if the entity has been released.
727    pub fn update_in<C, R>(
728        &self,
729        cx: &mut C,
730        update: impl FnOnce(&mut T, &mut Window, &mut Context<T>) -> R,
731    ) -> Result<R>
732    where
733        C: VisualContext,
734        Result<C::Result<R>>: crate::Flatten<R>,
735    {
736        let window = cx.window_handle();
737        let this = self.upgrade().context("entity released")?;
738
739        crate::Flatten::flatten(window.update(cx, |_, window, cx| {
740            this.update(cx, |entity, cx| update(entity, window, cx))
741        }))
742    }
743
744    /// Reads the entity referenced by this handle with the given function if
745    /// the referenced entity still exists. Returns an error if the entity has
746    /// been released.
747    pub fn read_with<C, R>(&self, cx: &C, read: impl FnOnce(&T, &App) -> R) -> Result<R>
748    where
749        C: AppContext,
750        Result<C::Result<R>>: crate::Flatten<R>,
751    {
752        crate::Flatten::flatten(
753            self.upgrade()
754                .context("entity released")
755                .map(|this| cx.read_entity(&this, read)),
756        )
757    }
758
759    /// Create a new weak entity that can never be upgraded.
760    pub fn new_invalid() -> Self {
761        Self {
762            any_entity: AnyWeakEntity::new_invalid(),
763            entity_type: PhantomData,
764        }
765    }
766}
767
768impl<T> Hash for WeakEntity<T> {
769    fn hash<H: Hasher>(&self, state: &mut H) {
770        self.any_entity.hash(state);
771    }
772}
773
774impl<T> PartialEq for WeakEntity<T> {
775    fn eq(&self, other: &Self) -> bool {
776        self.any_entity == other.any_entity
777    }
778}
779
780impl<T> Eq for WeakEntity<T> {}
781
782impl<T> PartialEq<Entity<T>> for WeakEntity<T> {
783    fn eq(&self, other: &Entity<T>) -> bool {
784        self.entity_id() == other.any_entity.entity_id()
785    }
786}
787
788impl<T: 'static> Ord for WeakEntity<T> {
789    fn cmp(&self, other: &Self) -> Ordering {
790        self.entity_id().cmp(&other.entity_id())
791    }
792}
793
794impl<T: 'static> PartialOrd for WeakEntity<T> {
795    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
796        Some(self.cmp(other))
797    }
798}
799
800#[cfg(any(test, gpui_leak_detection))]
801static LEAK_BACKTRACE: std::sync::LazyLock<bool> =
802    std::sync::LazyLock::new(|| std::env::var("LEAK_BACKTRACE").is_ok_and(|b| !b.is_empty()));
803
804#[cfg(any(test, gpui_leak_detection))]
805#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq)]
806pub(crate) struct HandleId {
807    id: u64, // id of the handle itself, not the pointed at object
808}
809
810#[cfg(any(test, gpui_leak_detection))]
811pub(crate) struct LeakDetector {
812    next_handle_id: u64,
813    entity_handles: HashMap<EntityId, HashMap<HandleId, Option<std::backtrace::Backtrace>>>,
814}
815
816#[cfg(any(test, gpui_leak_detection))]
817impl LeakDetector {
818    #[track_caller]
819    pub fn handle_created(&mut self, entity_id: EntityId) -> HandleId {
820        let id = util::post_inc(&mut self.next_handle_id);
821        let handle_id = HandleId { id };
822        let handles = self.entity_handles.entry(entity_id).or_default();
823        handles.insert(
824            handle_id,
825            LEAK_BACKTRACE.then(std::backtrace::Backtrace::force_capture),
826        );
827        handle_id
828    }
829
830    pub fn handle_released(&mut self, entity_id: EntityId, handle_id: HandleId) {
831        let handles = self.entity_handles.entry(entity_id).or_default();
832        handles.remove(&handle_id);
833    }
834
835    pub fn assert_released(&mut self, entity_id: EntityId) {
836        let handles = self.entity_handles.entry(entity_id).or_default();
837        if !handles.is_empty() {
838            for backtrace in handles.values_mut() {
839                if let Some(backtrace) = backtrace.take() {
840                    eprintln!("Leaked handle: {backtrace:#?}");
841                } else {
842                    eprintln!("Leaked handle: export LEAK_BACKTRACE to find allocation site");
843                }
844            }
845            panic!();
846        }
847    }
848}
849
850#[cfg(test)]
851mod test {
852    use crate::EntityMap;
853
854    struct TestEntity {
855        pub i: i32,
856    }
857
858    #[test]
859    fn test_entity_map_slot_assignment_before_cleanup() {
860        // Tests that slots are not re-used before take_dropped.
861        let mut entity_map = EntityMap::new();
862
863        let slot = entity_map.reserve::<TestEntity>();
864        entity_map.insert(slot, TestEntity { i: 1 });
865
866        let slot = entity_map.reserve::<TestEntity>();
867        entity_map.insert(slot, TestEntity { i: 2 });
868
869        let dropped = entity_map.take_dropped();
870        assert_eq!(dropped.len(), 2);
871
872        assert_eq!(
873            dropped
874                .into_iter()
875                .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
876                .collect::<Vec<i32>>(),
877            vec![1, 2],
878        );
879    }
880
881    #[test]
882    fn test_entity_map_weak_upgrade_before_cleanup() {
883        // Tests that weak handles are not upgraded before take_dropped
884        let mut entity_map = EntityMap::new();
885
886        let slot = entity_map.reserve::<TestEntity>();
887        let handle = entity_map.insert(slot, TestEntity { i: 1 });
888        let weak = handle.downgrade();
889        drop(handle);
890
891        let strong = weak.upgrade();
892        assert_eq!(strong, None);
893
894        let dropped = entity_map.take_dropped();
895        assert_eq!(dropped.len(), 1);
896
897        assert_eq!(
898            dropped
899                .into_iter()
900                .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
901                .collect::<Vec<i32>>(),
902            vec![1],
903        );
904    }
905}
gpui_core/app/entity_map.rs

gpui_core/app/
entity_map.rs
