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    num::NonZeroU64,
  15    sync::{
  16        Arc, Weak,
  17        atomic::{AtomicU64, AtomicUsize, Ordering::SeqCst},
  18    },
  19    thread::panicking,
  20};
  21
  22use super::Context;
  23use crate::util::atomic_incr_if_not_zero;
  24#[cfg(any(test, feature = "leak-detection"))]
  25use collections::HashMap;
  26
  27slotmap::new_key_type! {
  28    /// A unique identifier for a entity across the application.
  29    pub struct EntityId;
  30}
  31
  32impl From<u64> for EntityId {
  33    fn from(value: u64) -> Self {
  34        Self(KeyData::from_ffi(value))
  35    }
  36}
  37
  38impl EntityId {
  39    /// Converts this entity id to a [NonZeroU64]
  40    pub fn as_non_zero_u64(self) -> NonZeroU64 {
  41        NonZeroU64::new(self.0.as_ffi()).unwrap()
  42    }
  43
  44    /// Converts this entity id to a [u64]
  45    pub fn as_u64(self) -> u64 {
  46        self.0.as_ffi()
  47    }
  48}
  49
  50impl Display for EntityId {
  51    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  52        write!(f, "{}", self.as_u64())
  53    }
  54}
  55
  56pub(crate) struct EntityMap {
  57    entities: SecondaryMap<EntityId, Box<dyn Any>>,
  58    pub accessed_entities: RefCell<FxHashSet<EntityId>>,
  59    ref_counts: Arc<RwLock<EntityRefCounts>>,
  60}
  61
  62struct EntityRefCounts {
  63    counts: SlotMap<EntityId, AtomicUsize>,
  64    dropped_entity_ids: Vec<EntityId>,
  65    #[cfg(any(test, feature = "leak-detection"))]
  66    leak_detector: LeakDetector,
  67}
  68
  69impl EntityMap {
  70    pub fn new() -> Self {
  71        Self {
  72            entities: SecondaryMap::new(),
  73            accessed_entities: RefCell::new(FxHashSet::default()),
  74            ref_counts: Arc::new(RwLock::new(EntityRefCounts {
  75                counts: SlotMap::with_key(),
  76                dropped_entity_ids: Vec::new(),
  77                #[cfg(any(test, feature = "leak-detection"))]
  78                leak_detector: LeakDetector {
  79                    next_handle_id: 0,
  80                    entity_handles: HashMap::default(),
  81                },
  82            })),
  83        }
  84    }
  85
  86    /// Reserve a slot for an entity, which you can subsequently use with `insert`.
  87    pub fn reserve<T: 'static>(&self) -> Slot<T> {
  88        let id = self.ref_counts.write().counts.insert(1.into());
  89        Slot(Entity::new(id, Arc::downgrade(&self.ref_counts)))
  90    }
  91
  92    /// Insert an entity into a slot obtained by calling `reserve`.
  93    pub fn insert<T>(&mut self, slot: Slot<T>, entity: T) -> Entity<T>
  94    where
  95        T: 'static,
  96    {
  97        let mut accessed_entities = self.accessed_entities.get_mut();
  98        accessed_entities.insert(slot.entity_id);
  99
 100        let handle = slot.0;
 101        self.entities.insert(handle.entity_id, Box::new(entity));
 102        handle
 103    }
 104
 105    /// Move an entity to the stack.
 106    #[track_caller]
 107    pub fn lease<T>(&mut self, pointer: &Entity<T>) -> Lease<T> {
 108        self.assert_valid_context(pointer);
 109        let mut accessed_entities = self.accessed_entities.get_mut();
 110        accessed_entities.insert(pointer.entity_id);
 111
 112        let entity = Some(
 113            self.entities
 114                .remove(pointer.entity_id)
 115                .unwrap_or_else(|| double_lease_panic::<T>("update")),
 116        );
 117        Lease {
 118            entity,
 119            id: pointer.entity_id,
 120            entity_type: PhantomData,
 121        }
 122    }
 123
 124    /// Returns an entity after moving it to the stack.
 125    pub fn end_lease<T>(&mut self, mut lease: Lease<T>) {
 126        self.entities.insert(lease.id, lease.entity.take().unwrap());
 127    }
 128
 129    pub fn read<T: 'static>(&self, entity: &Entity<T>) -> &T {
 130        self.assert_valid_context(entity);
 131        let mut accessed_entities = self.accessed_entities.borrow_mut();
 132        accessed_entities.insert(entity.entity_id);
 133
 134        self.entities
 135            .get(entity.entity_id)
 136            .and_then(|entity| entity.downcast_ref())
 137            .unwrap_or_else(|| double_lease_panic::<T>("read"))
 138    }
 139
 140    fn assert_valid_context(&self, entity: &AnyEntity) {
 141        debug_assert!(
 142            Weak::ptr_eq(&entity.entity_map, &Arc::downgrade(&self.ref_counts)),
 143            "used a entity with the wrong context"
 144        );
 145    }
 146
 147    pub fn extend_accessed(&mut self, entities: &FxHashSet<EntityId>) {
 148        self.accessed_entities
 149            .get_mut()
 150            .extend(entities.iter().copied());
 151    }
 152
 153    pub fn clear_accessed(&mut self) {
 154        self.accessed_entities.get_mut().clear();
 155    }
 156
 157    pub fn take_dropped(&mut self) -> Vec<(EntityId, Box<dyn Any>)> {
 158        let mut ref_counts = &mut *self.ref_counts.write();
 159        let dropped_entity_ids = ref_counts.dropped_entity_ids.drain(..);
 160        let mut accessed_entities = self.accessed_entities.get_mut();
 161
 162        dropped_entity_ids
 163            .filter_map(|entity_id| {
 164                let count = ref_counts.counts.remove(entity_id).unwrap();
 165                debug_assert_eq!(
 166                    count.load(SeqCst),
 167                    0,
 168                    "dropped an entity that was referenced"
 169                );
 170                accessed_entities.remove(&entity_id);
 171                // If the EntityId was allocated with `Context::reserve`,
 172                // the entity may not have been inserted.
 173                Some((entity_id, self.entities.remove(entity_id)?))
 174            })
 175            .collect()
 176    }
 177}
 178
 179#[track_caller]
 180fn double_lease_panic<T>(operation: &str) -> ! {
 181    panic!(
 182        "cannot {operation} {} while it is already being updated",
 183        std::any::type_name::<T>()
 184    )
 185}
 186
 187pub(crate) struct Lease<T> {
 188    entity: Option<Box<dyn Any>>,
 189    pub id: EntityId,
 190    entity_type: PhantomData<T>,
 191}
 192
 193impl<T: 'static> core::ops::Deref for Lease<T> {
 194    type Target = T;
 195
 196    fn deref(&self) -> &Self::Target {
 197        self.entity.as_ref().unwrap().downcast_ref().unwrap()
 198    }
 199}
 200
 201impl<T: 'static> core::ops::DerefMut for Lease<T> {
 202    fn deref_mut(&mut self) -> &mut Self::Target {
 203        self.entity.as_mut().unwrap().downcast_mut().unwrap()
 204    }
 205}
 206
 207impl<T> Drop for Lease<T> {
 208    fn drop(&mut self) {
 209        if self.entity.is_some() && !panicking() {
 210            panic!("Leases must be ended with EntityMap::end_lease")
 211        }
 212    }
 213}
 214
 215#[derive(Deref, DerefMut)]
 216pub(crate) struct Slot<T>(Entity<T>);
 217
 218/// A dynamically typed reference to a entity, which can be downcast into a `Entity<T>`.
 219pub struct AnyEntity {
 220    pub(crate) entity_id: EntityId,
 221    pub(crate) entity_type: TypeId,
 222    entity_map: Weak<RwLock<EntityRefCounts>>,
 223    #[cfg(any(test, feature = "leak-detection"))]
 224    handle_id: HandleId,
 225}
 226
 227impl AnyEntity {
 228    fn new(id: EntityId, entity_type: TypeId, entity_map: Weak<RwLock<EntityRefCounts>>) -> Self {
 229        Self {
 230            entity_id: id,
 231            entity_type,
 232            #[cfg(any(test, feature = "leak-detection"))]
 233            handle_id: entity_map
 234                .clone()
 235                .upgrade()
 236                .unwrap()
 237                .write()
 238                .leak_detector
 239                .handle_created(id),
 240            entity_map,
 241        }
 242    }
 243
 244    /// Returns the id associated with this entity.
 245    #[inline]
 246    pub fn entity_id(&self) -> EntityId {
 247        self.entity_id
 248    }
 249
 250    /// Returns the [TypeId] associated with this entity.
 251    #[inline]
 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    #[inline]
 336    fn from(entity: Entity<T>) -> Self {
 337        entity.any_entity
 338    }
 339}
 340
 341impl Hash for AnyEntity {
 342    #[inline]
 343    fn hash<H: Hasher>(&self, state: &mut H) {
 344        self.entity_id.hash(state);
 345    }
 346}
 347
 348impl PartialEq for AnyEntity {
 349    #[inline]
 350    fn eq(&self, other: &Self) -> bool {
 351        self.entity_id == other.entity_id
 352    }
 353}
 354
 355impl Eq for AnyEntity {}
 356
 357impl Ord for AnyEntity {
 358    #[inline]
 359    fn cmp(&self, other: &Self) -> Ordering {
 360        self.entity_id.cmp(&other.entity_id)
 361    }
 362}
 363
 364impl PartialOrd for AnyEntity {
 365    #[inline]
 366    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
 367        Some(self.cmp(other))
 368    }
 369}
 370
 371impl std::fmt::Debug for AnyEntity {
 372    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 373        f.debug_struct("AnyEntity")
 374            .field("entity_id", &self.entity_id.as_u64())
 375            .finish()
 376    }
 377}
 378
 379/// A strong, well-typed reference to a struct which is managed
 380/// by GPUI
 381#[derive(Deref, DerefMut)]
 382pub struct Entity<T> {
 383    #[deref]
 384    #[deref_mut]
 385    pub(crate) any_entity: AnyEntity,
 386    pub(crate) entity_type: PhantomData<fn(T) -> T>,
 387}
 388
 389impl<T> Sealed for Entity<T> {}
 390
 391impl<T: 'static> Entity<T> {
 392    #[inline]
 393    fn new(id: EntityId, entity_map: Weak<RwLock<EntityRefCounts>>) -> Self
 394    where
 395        T: 'static,
 396    {
 397        Self {
 398            any_entity: AnyEntity::new(id, TypeId::of::<T>(), entity_map),
 399            entity_type: PhantomData,
 400        }
 401    }
 402
 403    /// Get the entity ID associated with this entity
 404    #[inline]
 405    pub fn entity_id(&self) -> EntityId {
 406        self.any_entity.entity_id
 407    }
 408
 409    /// Downgrade this entity pointer to a non-retaining weak pointer
 410    #[inline]
 411    pub fn downgrade(&self) -> WeakEntity<T> {
 412        WeakEntity {
 413            any_entity: self.any_entity.downgrade(),
 414            entity_type: self.entity_type,
 415        }
 416    }
 417
 418    /// Convert this into a dynamically typed entity.
 419    #[inline]
 420    pub fn into_any(self) -> AnyEntity {
 421        self.any_entity
 422    }
 423
 424    /// Grab a reference to this entity from the context.
 425    #[inline]
 426    pub fn read<'a>(&self, cx: &'a App) -> &'a T {
 427        cx.entities.read(self)
 428    }
 429
 430    /// Read the entity referenced by this handle with the given function.
 431    #[inline]
 432    pub fn read_with<R, C: AppContext>(&self, cx: &C, f: impl FnOnce(&T, &App) -> R) -> R {
 433        cx.read_entity(self, f)
 434    }
 435
 436    /// Updates the entity referenced by this handle with the given function.
 437    #[inline]
 438    pub fn update<R, C: AppContext>(
 439        &self,
 440        cx: &mut C,
 441        update: impl FnOnce(&mut T, &mut Context<T>) -> R,
 442    ) -> R {
 443        cx.update_entity(self, update)
 444    }
 445
 446    /// Updates the entity referenced by this handle with the given function.
 447    #[inline]
 448    pub fn as_mut<'a, C: AppContext>(&self, cx: &'a mut C) -> GpuiBorrow<'a, T> {
 449        cx.as_mut(self)
 450    }
 451
 452    /// Updates the entity referenced by this handle with the given function.
 453    pub fn write<C: AppContext>(&self, cx: &mut C, value: T) {
 454        self.update(cx, |entity, cx| {
 455            *entity = value;
 456            cx.notify();
 457        })
 458    }
 459
 460    /// Updates the entity referenced by this handle with the given function if
 461    /// the referenced entity still exists, within a visual context that has a window.
 462    /// Returns an error if the window has been closed.
 463    #[inline]
 464    pub fn update_in<R, C: VisualContext>(
 465        &self,
 466        cx: &mut C,
 467        update: impl FnOnce(&mut T, &mut Window, &mut Context<T>) -> R,
 468    ) -> C::Result<R> {
 469        cx.update_window_entity(self, update)
 470    }
 471}
 472
 473impl<T> Clone for Entity<T> {
 474    #[inline]
 475    fn clone(&self) -> Self {
 476        Self {
 477            any_entity: self.any_entity.clone(),
 478            entity_type: self.entity_type,
 479        }
 480    }
 481}
 482
 483impl<T> std::fmt::Debug for Entity<T> {
 484    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 485        f.debug_struct("Entity")
 486            .field("entity_id", &self.any_entity.entity_id)
 487            .field("entity_type", &type_name::<T>())
 488            .finish()
 489    }
 490}
 491
 492impl<T> Hash for Entity<T> {
 493    #[inline]
 494    fn hash<H: Hasher>(&self, state: &mut H) {
 495        self.any_entity.hash(state);
 496    }
 497}
 498
 499impl<T> PartialEq for Entity<T> {
 500    #[inline]
 501    fn eq(&self, other: &Self) -> bool {
 502        self.any_entity == other.any_entity
 503    }
 504}
 505
 506impl<T> Eq for Entity<T> {}
 507
 508impl<T> PartialEq<WeakEntity<T>> for Entity<T> {
 509    #[inline]
 510    fn eq(&self, other: &WeakEntity<T>) -> bool {
 511        self.any_entity.entity_id() == other.entity_id()
 512    }
 513}
 514
 515impl<T: 'static> Ord for Entity<T> {
 516    #[inline]
 517    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
 518        self.entity_id().cmp(&other.entity_id())
 519    }
 520}
 521
 522impl<T: 'static> PartialOrd for Entity<T> {
 523    #[inline]
 524    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
 525        Some(self.cmp(other))
 526    }
 527}
 528
 529/// A type erased, weak reference to a entity.
 530#[derive(Clone)]
 531pub struct AnyWeakEntity {
 532    pub(crate) entity_id: EntityId,
 533    entity_type: TypeId,
 534    entity_ref_counts: Weak<RwLock<EntityRefCounts>>,
 535}
 536
 537impl AnyWeakEntity {
 538    /// Get the entity ID associated with this weak reference.
 539    #[inline]
 540    pub fn entity_id(&self) -> EntityId {
 541        self.entity_id
 542    }
 543
 544    /// Check if this weak handle can be upgraded, or if the entity has already been dropped
 545    pub fn is_upgradable(&self) -> bool {
 546        let ref_count = self
 547            .entity_ref_counts
 548            .upgrade()
 549            .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
 550            .unwrap_or(0);
 551        ref_count > 0
 552    }
 553
 554    /// Upgrade this weak entity reference to a strong reference.
 555    pub fn upgrade(&self) -> Option<AnyEntity> {
 556        let ref_counts = &self.entity_ref_counts.upgrade()?;
 557        let ref_counts = ref_counts.read();
 558        let ref_count = ref_counts.counts.get(self.entity_id)?;
 559
 560        if atomic_incr_if_not_zero(ref_count) == 0 {
 561            // entity_id is in dropped_entity_ids
 562            return None;
 563        }
 564        drop(ref_counts);
 565
 566        Some(AnyEntity {
 567            entity_id: self.entity_id,
 568            entity_type: self.entity_type,
 569            entity_map: self.entity_ref_counts.clone(),
 570            #[cfg(any(test, feature = "leak-detection"))]
 571            handle_id: self
 572                .entity_ref_counts
 573                .upgrade()
 574                .unwrap()
 575                .write()
 576                .leak_detector
 577                .handle_created(self.entity_id),
 578        })
 579    }
 580
 581    /// Asserts that the entity referenced by this weak handle has been fully released.
 582    ///
 583    /// # Example
 584    ///
 585    /// ```ignore
 586    /// let entity = cx.new(|_| MyEntity::new());
 587    /// let weak = entity.downgrade();
 588    /// drop(entity);
 589    ///
 590    /// // Verify the entity was released
 591    /// weak.assert_released();
 592    /// ```
 593    ///
 594    /// # Debugging Leaks
 595    ///
 596    /// If this method panics due to leaked handles, set the `LEAK_BACKTRACE` environment
 597    /// variable to see where the leaked handles were allocated:
 598    ///
 599    /// ```bash
 600    /// LEAK_BACKTRACE=1 cargo test my_test
 601    /// ```
 602    ///
 603    /// # Panics
 604    ///
 605    /// - Panics if any strong handles to the entity are still alive.
 606    /// - Panics if the entity was recently dropped but cleanup hasn't completed yet
 607    ///   (resources are retained until the end of the effect cycle).
 608    #[cfg(any(test, feature = "leak-detection"))]
 609    pub fn assert_released(&self) {
 610        self.entity_ref_counts
 611            .upgrade()
 612            .unwrap()
 613            .write()
 614            .leak_detector
 615            .assert_released(self.entity_id);
 616
 617        if self
 618            .entity_ref_counts
 619            .upgrade()
 620            .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
 621            .is_some()
 622        {
 623            panic!(
 624                "entity was recently dropped but resources are retained until the end of the effect cycle."
 625            )
 626        }
 627    }
 628
 629    /// Creates a weak entity that can never be upgraded.
 630    pub fn new_invalid() -> Self {
 631        /// To hold the invariant that all ids are unique, and considering that slotmap
 632        /// increases their IDs from `0`, we can decrease ours from `u64::MAX` so these
 633        /// two will never conflict (u64 is way too large).
 634        static UNIQUE_NON_CONFLICTING_ID_GENERATOR: AtomicU64 = AtomicU64::new(u64::MAX);
 635        let entity_id = UNIQUE_NON_CONFLICTING_ID_GENERATOR.fetch_sub(1, SeqCst);
 636
 637        Self {
 638            // Safety:
 639            //   Docs say this is safe but can be unspecified if slotmap changes the representation
 640            //   after `1.0.7`, that said, providing a valid entity_id here is not necessary as long
 641            //   as we guarantee that `entity_id` is never used if `entity_ref_counts` equals
 642            //   to `Weak::new()` (that is, it's unable to upgrade), that is the invariant that
 643            //   actually needs to be hold true.
 644            //
 645            //   And there is no sane reason to read an entity slot if `entity_ref_counts` can't be
 646            //   read in the first place, so we're good!
 647            entity_id: entity_id.into(),
 648            entity_type: TypeId::of::<()>(),
 649            entity_ref_counts: Weak::new(),
 650        }
 651    }
 652}
 653
 654impl std::fmt::Debug for AnyWeakEntity {
 655    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 656        f.debug_struct(type_name::<Self>())
 657            .field("entity_id", &self.entity_id)
 658            .field("entity_type", &self.entity_type)
 659            .finish()
 660    }
 661}
 662
 663impl<T> From<WeakEntity<T>> for AnyWeakEntity {
 664    #[inline]
 665    fn from(entity: WeakEntity<T>) -> Self {
 666        entity.any_entity
 667    }
 668}
 669
 670impl Hash for AnyWeakEntity {
 671    #[inline]
 672    fn hash<H: Hasher>(&self, state: &mut H) {
 673        self.entity_id.hash(state);
 674    }
 675}
 676
 677impl PartialEq for AnyWeakEntity {
 678    #[inline]
 679    fn eq(&self, other: &Self) -> bool {
 680        self.entity_id == other.entity_id
 681    }
 682}
 683
 684impl Eq for AnyWeakEntity {}
 685
 686impl Ord for AnyWeakEntity {
 687    #[inline]
 688    fn cmp(&self, other: &Self) -> Ordering {
 689        self.entity_id.cmp(&other.entity_id)
 690    }
 691}
 692
 693impl PartialOrd for AnyWeakEntity {
 694    #[inline]
 695    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
 696        Some(self.cmp(other))
 697    }
 698}
 699
 700/// A weak reference to a entity of the given type.
 701#[derive(Deref, DerefMut)]
 702pub struct WeakEntity<T> {
 703    #[deref]
 704    #[deref_mut]
 705    any_entity: AnyWeakEntity,
 706    entity_type: PhantomData<fn(T) -> T>,
 707}
 708
 709impl<T> std::fmt::Debug for WeakEntity<T> {
 710    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 711        f.debug_struct(type_name::<Self>())
 712            .field("entity_id", &self.any_entity.entity_id)
 713            .field("entity_type", &type_name::<T>())
 714            .finish()
 715    }
 716}
 717
 718impl<T> Clone for WeakEntity<T> {
 719    fn clone(&self) -> Self {
 720        Self {
 721            any_entity: self.any_entity.clone(),
 722            entity_type: self.entity_type,
 723        }
 724    }
 725}
 726
 727impl<T: 'static> WeakEntity<T> {
 728    /// Upgrade this weak entity reference into a strong entity reference
 729    pub fn upgrade(&self) -> Option<Entity<T>> {
 730        Some(Entity {
 731            any_entity: self.any_entity.upgrade()?,
 732            entity_type: self.entity_type,
 733        })
 734    }
 735
 736    /// Updates the entity referenced by this handle with the given function if
 737    /// the referenced entity still exists. Returns an error if the entity has
 738    /// been released.
 739    pub fn update<C, R>(
 740        &self,
 741        cx: &mut C,
 742        update: impl FnOnce(&mut T, &mut Context<T>) -> R,
 743    ) -> Result<R>
 744    where
 745        C: AppContext,
 746    {
 747        let entity = self.upgrade().context("entity released")?;
 748        Ok(cx.update_entity(&entity, update))
 749    }
 750
 751    /// Updates the entity referenced by this handle with the given function if
 752    /// the referenced entity still exists, within a visual context that has a window.
 753    /// Returns an error if the entity has been released.
 754    pub fn update_in<C, R>(
 755        &self,
 756        cx: &mut C,
 757        update: impl FnOnce(&mut T, &mut Window, &mut Context<T>) -> R,
 758    ) -> Result<R>
 759    where
 760        C: VisualContext,
 761    {
 762        let window = cx.window_handle();
 763        let entity = self.upgrade().context("entity released")?;
 764
 765        window.update(cx, |_, window, cx| {
 766            entity.update(cx, |entity, cx| update(entity, window, cx))
 767        })
 768    }
 769
 770    /// Reads the entity referenced by this handle with the given function if
 771    /// the referenced entity still exists. Returns an error if the entity has
 772    /// been released.
 773    pub fn read_with<C, R>(&self, cx: &C, read: impl FnOnce(&T, &App) -> R) -> Result<R>
 774    where
 775        C: AppContext,
 776    {
 777        let entity = self.upgrade().context("entity released")?;
 778        Ok(cx.read_entity(&entity, read))
 779    }
 780
 781    /// Create a new weak entity that can never be upgraded.
 782    #[inline]
 783    pub fn new_invalid() -> Self {
 784        Self {
 785            any_entity: AnyWeakEntity::new_invalid(),
 786            entity_type: PhantomData,
 787        }
 788    }
 789}
 790
 791impl<T> Hash for WeakEntity<T> {
 792    #[inline]
 793    fn hash<H: Hasher>(&self, state: &mut H) {
 794        self.any_entity.hash(state);
 795    }
 796}
 797
 798impl<T> PartialEq for WeakEntity<T> {
 799    #[inline]
 800    fn eq(&self, other: &Self) -> bool {
 801        self.any_entity == other.any_entity
 802    }
 803}
 804
 805impl<T> Eq for WeakEntity<T> {}
 806
 807impl<T> PartialEq<Entity<T>> for WeakEntity<T> {
 808    #[inline]
 809    fn eq(&self, other: &Entity<T>) -> bool {
 810        self.entity_id() == other.any_entity.entity_id()
 811    }
 812}
 813
 814impl<T: 'static> Ord for WeakEntity<T> {
 815    #[inline]
 816    fn cmp(&self, other: &Self) -> Ordering {
 817        self.entity_id().cmp(&other.entity_id())
 818    }
 819}
 820
 821impl<T: 'static> PartialOrd for WeakEntity<T> {
 822    #[inline]
 823    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
 824        Some(self.cmp(other))
 825    }
 826}
 827
 828/// Controls whether backtraces are captured when entity handles are created.
 829///
 830/// Set the `LEAK_BACKTRACE` environment variable to any non-empty value to enable
 831/// backtrace capture. This helps identify where leaked handles were allocated.
 832#[cfg(any(test, feature = "leak-detection"))]
 833static LEAK_BACKTRACE: std::sync::LazyLock<bool> =
 834    std::sync::LazyLock::new(|| std::env::var("LEAK_BACKTRACE").is_ok_and(|b| !b.is_empty()));
 835
 836/// Unique identifier for a specific entity handle instance.
 837///
 838/// This is distinct from `EntityId` - while multiple handles can point to the same
 839/// entity (same `EntityId`), each handle has its own unique `HandleId`.
 840#[cfg(any(test, feature = "leak-detection"))]
 841#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq)]
 842pub(crate) struct HandleId {
 843    id: u64,
 844}
 845
 846/// Tracks entity handle allocations to detect leaks.
 847///
 848/// The leak detector is enabled in tests and when the `leak-detection` feature is active.
 849/// It tracks every `Entity<T>` and `AnyEntity` handle that is created and released,
 850/// allowing you to verify that all handles to an entity have been properly dropped.
 851///
 852/// # How do leaks happen?
 853///
 854/// Entities are reference-counted structures that can own other entities
 855/// allowing to form cycles. If such a strong-reference counted cycle is
 856/// created, all participating strong entities in this cycle will effectively
 857/// leak as they cannot be released anymore.
 858///
 859/// # Usage
 860///
 861/// You can use `WeakEntity::assert_released` or `AnyWeakEntity::assert_released`
 862/// to verify that an entity has been fully released:
 863///
 864/// ```ignore
 865/// let entity = cx.new(|_| MyEntity::new());
 866/// let weak = entity.downgrade();
 867/// drop(entity);
 868///
 869/// // This will panic if any handles to the entity are still alive
 870/// weak.assert_released();
 871/// ```
 872///
 873/// # Debugging Leaks
 874///
 875/// When a leak is detected, the detector will panic with information about the leaked
 876/// handles. To see where the leaked handles were allocated, set the `LEAK_BACKTRACE`
 877/// environment variable:
 878///
 879/// ```bash
 880/// LEAK_BACKTRACE=1 cargo test my_test
 881/// ```
 882///
 883/// This will capture and display backtraces for each leaked handle, helping you
 884/// identify where handles were created but not released.
 885///
 886/// # How It Works
 887///
 888/// - When an entity handle is created (via `Entity::new`, `Entity::clone`, or
 889///   `WeakEntity::upgrade`), `handle_created` is called to register the handle.
 890/// - When a handle is dropped, `handle_released` removes it from tracking.
 891/// - `assert_released` verifies that no handles remain for a given entity.
 892#[cfg(any(test, feature = "leak-detection"))]
 893pub(crate) struct LeakDetector {
 894    next_handle_id: u64,
 895    entity_handles: HashMap<EntityId, HashMap<HandleId, Option<backtrace::Backtrace>>>,
 896}
 897
 898#[cfg(any(test, feature = "leak-detection"))]
 899impl LeakDetector {
 900    /// Records that a new handle has been created for the given entity.
 901    ///
 902    /// Returns a unique `HandleId` that must be passed to `handle_released` when
 903    /// the handle is dropped. If `LEAK_BACKTRACE` is set, captures a backtrace
 904    /// at the allocation site.
 905    #[track_caller]
 906    pub fn handle_created(&mut self, entity_id: EntityId) -> HandleId {
 907        let id = gpui_util::post_inc(&mut self.next_handle_id);
 908        let handle_id = HandleId { id };
 909        let handles = self.entity_handles.entry(entity_id).or_default();
 910        handles.insert(
 911            handle_id,
 912            LEAK_BACKTRACE.then(backtrace::Backtrace::new_unresolved),
 913        );
 914        handle_id
 915    }
 916
 917    /// Records that a handle has been released (dropped).
 918    ///
 919    /// This removes the handle from tracking. The `handle_id` should be the same
 920    /// one returned by `handle_created` when the handle was allocated.
 921    pub fn handle_released(&mut self, entity_id: EntityId, handle_id: HandleId) {
 922        let handles = self.entity_handles.entry(entity_id).or_default();
 923        handles.remove(&handle_id);
 924    }
 925
 926    /// Asserts that all handles to the given entity have been released.
 927    ///
 928    /// # Panics
 929    ///
 930    /// Panics if any handles to the entity are still alive. The panic message
 931    /// includes backtraces for each leaked handle if `LEAK_BACKTRACE` is set,
 932    /// otherwise it suggests setting the environment variable to get more info.
 933    pub fn assert_released(&mut self, entity_id: EntityId) {
 934        use std::fmt::Write as _;
 935        let handles = self.entity_handles.entry(entity_id).or_default();
 936        if !handles.is_empty() {
 937            let mut out = String::new();
 938            for backtrace in handles.values_mut() {
 939                if let Some(mut backtrace) = backtrace.take() {
 940                    backtrace.resolve();
 941                    writeln!(out, "Leaked handle:\n{:?}", backtrace).unwrap();
 942                } else {
 943                    writeln!(
 944                        out,
 945                        "Leaked handle: (export LEAK_BACKTRACE to find allocation site)"
 946                    )
 947                    .unwrap();
 948                }
 949            }
 950            panic!("{out}");
 951        }
 952    }
 953}
 954
 955#[cfg(test)]
 956mod test {
 957    use crate::EntityMap;
 958
 959    struct TestEntity {
 960        pub i: i32,
 961    }
 962
 963    #[test]
 964    fn test_entity_map_slot_assignment_before_cleanup() {
 965        // Tests that slots are not re-used before take_dropped.
 966        let mut entity_map = EntityMap::new();
 967
 968        let slot = entity_map.reserve::<TestEntity>();
 969        entity_map.insert(slot, TestEntity { i: 1 });
 970
 971        let slot = entity_map.reserve::<TestEntity>();
 972        entity_map.insert(slot, TestEntity { i: 2 });
 973
 974        let dropped = entity_map.take_dropped();
 975        assert_eq!(dropped.len(), 2);
 976
 977        assert_eq!(
 978            dropped
 979                .into_iter()
 980                .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
 981                .collect::<Vec<i32>>(),
 982            vec![1, 2],
 983        );
 984    }
 985
 986    #[test]
 987    fn test_entity_map_weak_upgrade_before_cleanup() {
 988        // Tests that weak handles are not upgraded before take_dropped
 989        let mut entity_map = EntityMap::new();
 990
 991        let slot = entity_map.reserve::<TestEntity>();
 992        let handle = entity_map.insert(slot, TestEntity { i: 1 });
 993        let weak = handle.downgrade();
 994        drop(handle);
 995
 996        let strong = weak.upgrade();
 997        assert_eq!(strong, None);
 998
 999        let dropped = entity_map.take_dropped();
1000        assert_eq!(dropped.len(), 1);
1001
1002        assert_eq!(
1003            dropped
1004                .into_iter()
1005                .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
1006                .collect::<Vec<i32>>(),
1007            vec![1],
1008        );
1009    }
1010}