entity_map.rs

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