1use crate::{seal::Sealed, AppContext, Context, Entity, ModelContext};
2use anyhow::{anyhow, Result};
3use derive_more::{Deref, DerefMut};
4use parking_lot::{RwLock, RwLockUpgradableReadGuard};
5use slotmap::{KeyData, SecondaryMap, SlotMap};
6use std::{
7 any::{type_name, Any, TypeId},
8 fmt::{self, Display},
9 hash::{Hash, Hasher},
10 marker::PhantomData,
11 mem,
12 sync::{
13 atomic::{AtomicUsize, Ordering::SeqCst},
14 Arc, Weak,
15 },
16 thread::panicking,
17};
18
19#[cfg(any(test, feature = "test-support"))]
20use collections::HashMap;
21
22slotmap::new_key_type! {
23 /// A unique identifier for a model or view across the application.
24 pub struct EntityId;
25}
26
27impl From<u64> for EntityId {
28 fn from(value: u64) -> Self {
29 Self(KeyData::from_ffi(value))
30 }
31}
32
33impl EntityId {
34 /// Converts this entity id to a [u64]
35 pub fn as_u64(self) -> u64 {
36 self.0.as_ffi()
37 }
38}
39
40impl Display for EntityId {
41 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42 write!(f, "{}", self.as_u64())
43 }
44}
45
46pub(crate) struct EntityMap {
47 entities: SecondaryMap<EntityId, Box<dyn Any>>,
48 ref_counts: Arc<RwLock<EntityRefCounts>>,
49}
50
51struct EntityRefCounts {
52 counts: SlotMap<EntityId, AtomicUsize>,
53 dropped_entity_ids: Vec<EntityId>,
54 #[cfg(any(test, feature = "test-support"))]
55 leak_detector: LeakDetector,
56}
57
58impl EntityMap {
59 pub fn new() -> Self {
60 Self {
61 entities: SecondaryMap::new(),
62 ref_counts: Arc::new(RwLock::new(EntityRefCounts {
63 counts: SlotMap::with_key(),
64 dropped_entity_ids: Vec::new(),
65 #[cfg(any(test, feature = "test-support"))]
66 leak_detector: LeakDetector {
67 next_handle_id: 0,
68 entity_handles: HashMap::default(),
69 },
70 })),
71 }
72 }
73
74 /// Reserve a slot for an entity, which you can subsequently use with `insert`.
75 pub fn reserve<T: 'static>(&self) -> Slot<T> {
76 let id = self.ref_counts.write().counts.insert(1.into());
77 Slot(Model::new(id, Arc::downgrade(&self.ref_counts)))
78 }
79
80 /// Insert an entity into a slot obtained by calling `reserve`.
81 pub fn insert<T>(&mut self, slot: Slot<T>, entity: T) -> Model<T>
82 where
83 T: 'static,
84 {
85 let model = slot.0;
86 self.entities.insert(model.entity_id, Box::new(entity));
87 model
88 }
89
90 /// Move an entity to the stack.
91 #[track_caller]
92 pub fn lease<'a, T>(&mut self, model: &'a Model<T>) -> Lease<'a, T> {
93 self.assert_valid_context(model);
94 let entity = Some(self.entities.remove(model.entity_id).unwrap_or_else(|| {
95 panic!(
96 "Circular entity lease of {}. Is it already being updated?",
97 std::any::type_name::<T>()
98 )
99 }));
100 Lease {
101 model,
102 entity,
103 entity_type: PhantomData,
104 }
105 }
106
107 /// Returns an entity after moving it to the stack.
108 pub fn end_lease<T>(&mut self, mut lease: Lease<T>) {
109 self.entities
110 .insert(lease.model.entity_id, lease.entity.take().unwrap());
111 }
112
113 pub fn read<T: 'static>(&self, model: &Model<T>) -> &T {
114 self.assert_valid_context(model);
115 self.entities[model.entity_id].downcast_ref().unwrap()
116 }
117
118 fn assert_valid_context(&self, model: &AnyModel) {
119 debug_assert!(
120 Weak::ptr_eq(&model.entity_map, &Arc::downgrade(&self.ref_counts)),
121 "used a model with the wrong context"
122 );
123 }
124
125 pub fn take_dropped(&mut self) -> Vec<(EntityId, Box<dyn Any>)> {
126 let mut ref_counts = self.ref_counts.write();
127 let dropped_entity_ids = mem::take(&mut ref_counts.dropped_entity_ids);
128
129 dropped_entity_ids
130 .into_iter()
131 .map(|entity_id| {
132 let count = ref_counts.counts.remove(entity_id).unwrap();
133 debug_assert_eq!(
134 count.load(SeqCst),
135 0,
136 "dropped an entity that was referenced"
137 );
138 (entity_id, self.entities.remove(entity_id).unwrap())
139 })
140 .collect()
141 }
142}
143
144pub(crate) struct Lease<'a, T> {
145 entity: Option<Box<dyn Any>>,
146 pub model: &'a Model<T>,
147 entity_type: PhantomData<T>,
148}
149
150impl<'a, T: 'static> core::ops::Deref for Lease<'a, T> {
151 type Target = T;
152
153 fn deref(&self) -> &Self::Target {
154 self.entity.as_ref().unwrap().downcast_ref().unwrap()
155 }
156}
157
158impl<'a, T: 'static> core::ops::DerefMut for Lease<'a, T> {
159 fn deref_mut(&mut self) -> &mut Self::Target {
160 self.entity.as_mut().unwrap().downcast_mut().unwrap()
161 }
162}
163
164impl<'a, T> Drop for Lease<'a, T> {
165 fn drop(&mut self) {
166 if self.entity.is_some() && !panicking() {
167 panic!("Leases must be ended with EntityMap::end_lease")
168 }
169 }
170}
171
172#[derive(Deref, DerefMut)]
173pub(crate) struct Slot<T>(Model<T>);
174
175/// A dynamically typed reference to a model, which can be downcast into a `Model<T>`.
176pub struct AnyModel {
177 pub(crate) entity_id: EntityId,
178 pub(crate) entity_type: TypeId,
179 entity_map: Weak<RwLock<EntityRefCounts>>,
180 #[cfg(any(test, feature = "test-support"))]
181 handle_id: HandleId,
182}
183
184impl AnyModel {
185 fn new(id: EntityId, entity_type: TypeId, entity_map: Weak<RwLock<EntityRefCounts>>) -> Self {
186 Self {
187 entity_id: id,
188 entity_type,
189 entity_map: entity_map.clone(),
190 #[cfg(any(test, feature = "test-support"))]
191 handle_id: entity_map
192 .upgrade()
193 .unwrap()
194 .write()
195 .leak_detector
196 .handle_created(id),
197 }
198 }
199
200 /// Returns the id associated with this model.
201 pub fn entity_id(&self) -> EntityId {
202 self.entity_id
203 }
204
205 /// Returns the [TypeId] associated with this model.
206 pub fn entity_type(&self) -> TypeId {
207 self.entity_type
208 }
209
210 /// Converts this model handle into a weak variant, which does not prevent it from being released.
211 pub fn downgrade(&self) -> AnyWeakModel {
212 AnyWeakModel {
213 entity_id: self.entity_id,
214 entity_type: self.entity_type,
215 entity_ref_counts: self.entity_map.clone(),
216 }
217 }
218
219 /// Converts this model handle into a strongly-typed model handle of the given type.
220 /// If this model handle is not of the specified type, returns itself as an error variant.
221 pub fn downcast<T: 'static>(self) -> Result<Model<T>, AnyModel> {
222 if TypeId::of::<T>() == self.entity_type {
223 Ok(Model {
224 any_model: self,
225 entity_type: PhantomData,
226 })
227 } else {
228 Err(self)
229 }
230 }
231}
232
233impl Clone for AnyModel {
234 fn clone(&self) -> Self {
235 if let Some(entity_map) = self.entity_map.upgrade() {
236 let entity_map = entity_map.read();
237 let count = entity_map
238 .counts
239 .get(self.entity_id)
240 .expect("detected over-release of a model");
241 let prev_count = count.fetch_add(1, SeqCst);
242 assert_ne!(prev_count, 0, "Detected over-release of a model.");
243 }
244
245 Self {
246 entity_id: self.entity_id,
247 entity_type: self.entity_type,
248 entity_map: self.entity_map.clone(),
249 #[cfg(any(test, feature = "test-support"))]
250 handle_id: self
251 .entity_map
252 .upgrade()
253 .unwrap()
254 .write()
255 .leak_detector
256 .handle_created(self.entity_id),
257 }
258 }
259}
260
261impl Drop for AnyModel {
262 fn drop(&mut self) {
263 if let Some(entity_map) = self.entity_map.upgrade() {
264 let entity_map = entity_map.upgradable_read();
265 let count = entity_map
266 .counts
267 .get(self.entity_id)
268 .expect("detected over-release of a handle.");
269 let prev_count = count.fetch_sub(1, SeqCst);
270 assert_ne!(prev_count, 0, "Detected over-release of a model.");
271 if prev_count == 1 {
272 // We were the last reference to this entity, so we can remove it.
273 let mut entity_map = RwLockUpgradableReadGuard::upgrade(entity_map);
274 entity_map.dropped_entity_ids.push(self.entity_id);
275 }
276 }
277
278 #[cfg(any(test, feature = "test-support"))]
279 if let Some(entity_map) = self.entity_map.upgrade() {
280 entity_map
281 .write()
282 .leak_detector
283 .handle_released(self.entity_id, self.handle_id)
284 }
285 }
286}
287
288impl<T> From<Model<T>> for AnyModel {
289 fn from(model: Model<T>) -> Self {
290 model.any_model
291 }
292}
293
294impl Hash for AnyModel {
295 fn hash<H: Hasher>(&self, state: &mut H) {
296 self.entity_id.hash(state);
297 }
298}
299
300impl PartialEq for AnyModel {
301 fn eq(&self, other: &Self) -> bool {
302 self.entity_id == other.entity_id
303 }
304}
305
306impl Eq for AnyModel {}
307
308impl std::fmt::Debug for AnyModel {
309 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
310 f.debug_struct("AnyModel")
311 .field("entity_id", &self.entity_id.as_u64())
312 .finish()
313 }
314}
315
316/// A strong, well typed reference to a struct which is managed
317/// by GPUI
318#[derive(Deref, DerefMut)]
319pub struct Model<T> {
320 #[deref]
321 #[deref_mut]
322 pub(crate) any_model: AnyModel,
323 pub(crate) entity_type: PhantomData<T>,
324}
325
326unsafe impl<T> Send for Model<T> {}
327unsafe impl<T> Sync for Model<T> {}
328impl<T> Sealed for Model<T> {}
329
330impl<T: 'static> Entity<T> for Model<T> {
331 type Weak = WeakModel<T>;
332
333 fn entity_id(&self) -> EntityId {
334 self.any_model.entity_id
335 }
336
337 fn downgrade(&self) -> Self::Weak {
338 WeakModel {
339 any_model: self.any_model.downgrade(),
340 entity_type: self.entity_type,
341 }
342 }
343
344 fn upgrade_from(weak: &Self::Weak) -> Option<Self>
345 where
346 Self: Sized,
347 {
348 Some(Model {
349 any_model: weak.any_model.upgrade()?,
350 entity_type: weak.entity_type,
351 })
352 }
353}
354
355impl<T: 'static> Model<T> {
356 fn new(id: EntityId, entity_map: Weak<RwLock<EntityRefCounts>>) -> Self
357 where
358 T: 'static,
359 {
360 Self {
361 any_model: AnyModel::new(id, TypeId::of::<T>(), entity_map),
362 entity_type: PhantomData,
363 }
364 }
365
366 /// Downgrade the this to a weak model reference
367 pub fn downgrade(&self) -> WeakModel<T> {
368 // Delegate to the trait implementation to keep behavior in one place.
369 // This method was included to improve method resolution in the presence of
370 // the Model's deref
371 Entity::downgrade(self)
372 }
373
374 /// Convert this into a dynamically typed model.
375 pub fn into_any(self) -> AnyModel {
376 self.any_model
377 }
378
379 /// Grab a reference to this entity from the context.
380 pub fn read<'a>(&self, cx: &'a AppContext) -> &'a T {
381 cx.entities.read(self)
382 }
383
384 /// Read the entity referenced by this model with the given function.
385 pub fn read_with<R, C: Context>(
386 &self,
387 cx: &C,
388 f: impl FnOnce(&T, &AppContext) -> R,
389 ) -> C::Result<R> {
390 cx.read_model(self, f)
391 }
392
393 /// Updates the entity referenced by this model with the given function.
394 ///
395 /// The update function receives a context appropriate for its environment.
396 /// When updating in an `AppContext`, it receives a `ModelContext`.
397 /// When updating an a `WindowContext`, it receives a `ViewContext`.
398 pub fn update<C, R>(
399 &self,
400 cx: &mut C,
401 update: impl FnOnce(&mut T, &mut ModelContext<'_, T>) -> R,
402 ) -> C::Result<R>
403 where
404 C: Context,
405 {
406 cx.update_model(self, update)
407 }
408}
409
410impl<T> Clone for Model<T> {
411 fn clone(&self) -> Self {
412 Self {
413 any_model: self.any_model.clone(),
414 entity_type: self.entity_type,
415 }
416 }
417}
418
419impl<T> std::fmt::Debug for Model<T> {
420 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
421 write!(
422 f,
423 "Model {{ entity_id: {:?}, entity_type: {:?} }}",
424 self.any_model.entity_id,
425 type_name::<T>()
426 )
427 }
428}
429
430impl<T> Hash for Model<T> {
431 fn hash<H: Hasher>(&self, state: &mut H) {
432 self.any_model.hash(state);
433 }
434}
435
436impl<T> PartialEq for Model<T> {
437 fn eq(&self, other: &Self) -> bool {
438 self.any_model == other.any_model
439 }
440}
441
442impl<T> Eq for Model<T> {}
443
444impl<T> PartialEq<WeakModel<T>> for Model<T> {
445 fn eq(&self, other: &WeakModel<T>) -> bool {
446 self.any_model.entity_id() == other.entity_id()
447 }
448}
449
450/// A type erased, weak reference to a model.
451#[derive(Clone)]
452pub struct AnyWeakModel {
453 pub(crate) entity_id: EntityId,
454 entity_type: TypeId,
455 entity_ref_counts: Weak<RwLock<EntityRefCounts>>,
456}
457
458impl AnyWeakModel {
459 /// Get the entity ID associated with this weak reference.
460 pub fn entity_id(&self) -> EntityId {
461 self.entity_id
462 }
463
464 /// Check if this weak handle can be upgraded, or if the model has already been dropped
465 pub fn is_upgradable(&self) -> bool {
466 let ref_count = self
467 .entity_ref_counts
468 .upgrade()
469 .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
470 .unwrap_or(0);
471 ref_count > 0
472 }
473
474 /// Upgrade this weak model reference to a strong reference.
475 pub fn upgrade(&self) -> Option<AnyModel> {
476 let ref_counts = &self.entity_ref_counts.upgrade()?;
477 let ref_counts = ref_counts.read();
478 let ref_count = ref_counts.counts.get(self.entity_id)?;
479
480 // entity_id is in dropped_entity_ids
481 if ref_count.load(SeqCst) == 0 {
482 return None;
483 }
484 ref_count.fetch_add(1, SeqCst);
485 drop(ref_counts);
486
487 Some(AnyModel {
488 entity_id: self.entity_id,
489 entity_type: self.entity_type,
490 entity_map: self.entity_ref_counts.clone(),
491 #[cfg(any(test, feature = "test-support"))]
492 handle_id: self
493 .entity_ref_counts
494 .upgrade()
495 .unwrap()
496 .write()
497 .leak_detector
498 .handle_created(self.entity_id),
499 })
500 }
501
502 /// Assert that model referenced by this weak handle has been released.
503 #[cfg(any(test, feature = "test-support"))]
504 pub fn assert_released(&self) {
505 self.entity_ref_counts
506 .upgrade()
507 .unwrap()
508 .write()
509 .leak_detector
510 .assert_released(self.entity_id);
511
512 if self
513 .entity_ref_counts
514 .upgrade()
515 .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
516 .is_some()
517 {
518 panic!(
519 "entity was recently dropped but resources are retained until the end of the effect cycle."
520 )
521 }
522 }
523}
524
525impl<T> From<WeakModel<T>> for AnyWeakModel {
526 fn from(model: WeakModel<T>) -> Self {
527 model.any_model
528 }
529}
530
531impl Hash for AnyWeakModel {
532 fn hash<H: Hasher>(&self, state: &mut H) {
533 self.entity_id.hash(state);
534 }
535}
536
537impl PartialEq for AnyWeakModel {
538 fn eq(&self, other: &Self) -> bool {
539 self.entity_id == other.entity_id
540 }
541}
542
543impl Eq for AnyWeakModel {}
544
545/// A weak reference to a model of the given type.
546#[derive(Deref, DerefMut)]
547pub struct WeakModel<T> {
548 #[deref]
549 #[deref_mut]
550 any_model: AnyWeakModel,
551 entity_type: PhantomData<T>,
552}
553
554unsafe impl<T> Send for WeakModel<T> {}
555unsafe impl<T> Sync for WeakModel<T> {}
556
557impl<T> Clone for WeakModel<T> {
558 fn clone(&self) -> Self {
559 Self {
560 any_model: self.any_model.clone(),
561 entity_type: self.entity_type,
562 }
563 }
564}
565
566impl<T: 'static> WeakModel<T> {
567 /// Upgrade this weak model reference into a strong model reference
568 pub fn upgrade(&self) -> Option<Model<T>> {
569 // Delegate to the trait implementation to keep behavior in one place.
570 Model::upgrade_from(self)
571 }
572
573 /// Updates the entity referenced by this model with the given function if
574 /// the referenced entity still exists. Returns an error if the entity has
575 /// been released.
576 pub fn update<C, R>(
577 &self,
578 cx: &mut C,
579 update: impl FnOnce(&mut T, &mut ModelContext<'_, T>) -> R,
580 ) -> Result<R>
581 where
582 C: Context,
583 Result<C::Result<R>>: crate::Flatten<R>,
584 {
585 crate::Flatten::flatten(
586 self.upgrade()
587 .ok_or_else(|| anyhow!("entity release"))
588 .map(|this| cx.update_model(&this, update)),
589 )
590 }
591
592 /// Reads the entity referenced by this model with the given function if
593 /// the referenced entity still exists. Returns an error if the entity has
594 /// been released.
595 pub fn read_with<C, R>(&self, cx: &C, read: impl FnOnce(&T, &AppContext) -> R) -> Result<R>
596 where
597 C: Context,
598 Result<C::Result<R>>: crate::Flatten<R>,
599 {
600 crate::Flatten::flatten(
601 self.upgrade()
602 .ok_or_else(|| anyhow!("entity release"))
603 .map(|this| cx.read_model(&this, read)),
604 )
605 }
606}
607
608impl<T> Hash for WeakModel<T> {
609 fn hash<H: Hasher>(&self, state: &mut H) {
610 self.any_model.hash(state);
611 }
612}
613
614impl<T> PartialEq for WeakModel<T> {
615 fn eq(&self, other: &Self) -> bool {
616 self.any_model == other.any_model
617 }
618}
619
620impl<T> Eq for WeakModel<T> {}
621
622impl<T> PartialEq<Model<T>> for WeakModel<T> {
623 fn eq(&self, other: &Model<T>) -> bool {
624 self.entity_id() == other.any_model.entity_id()
625 }
626}
627
628#[cfg(any(test, feature = "test-support"))]
629lazy_static::lazy_static! {
630 static ref LEAK_BACKTRACE: bool =
631 std::env::var("LEAK_BACKTRACE").map_or(false, |b| !b.is_empty());
632}
633
634#[cfg(any(test, feature = "test-support"))]
635#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq)]
636pub(crate) struct HandleId {
637 id: u64, // id of the handle itself, not the pointed at object
638}
639
640#[cfg(any(test, feature = "test-support"))]
641pub(crate) struct LeakDetector {
642 next_handle_id: u64,
643 entity_handles: HashMap<EntityId, HashMap<HandleId, Option<backtrace::Backtrace>>>,
644}
645
646#[cfg(any(test, feature = "test-support"))]
647impl LeakDetector {
648 #[track_caller]
649 pub fn handle_created(&mut self, entity_id: EntityId) -> HandleId {
650 let id = util::post_inc(&mut self.next_handle_id);
651 let handle_id = HandleId { id };
652 let handles = self.entity_handles.entry(entity_id).or_default();
653 handles.insert(
654 handle_id,
655 LEAK_BACKTRACE.then(|| backtrace::Backtrace::new_unresolved()),
656 );
657 handle_id
658 }
659
660 pub fn handle_released(&mut self, entity_id: EntityId, handle_id: HandleId) {
661 let handles = self.entity_handles.entry(entity_id).or_default();
662 handles.remove(&handle_id);
663 }
664
665 pub fn assert_released(&mut self, entity_id: EntityId) {
666 let handles = self.entity_handles.entry(entity_id).or_default();
667 if !handles.is_empty() {
668 for (_, backtrace) in handles {
669 if let Some(mut backtrace) = backtrace.take() {
670 backtrace.resolve();
671 eprintln!("Leaked handle: {:#?}", backtrace);
672 } else {
673 eprintln!("Leaked handle: export LEAK_BACKTRACE to find allocation site");
674 }
675 }
676 panic!();
677 }
678 }
679}
680
681#[cfg(test)]
682mod test {
683 use crate::EntityMap;
684
685 struct TestEntity {
686 pub i: i32,
687 }
688
689 #[test]
690 fn test_entity_map_slot_assignment_before_cleanup() {
691 // Tests that slots are not re-used before take_dropped.
692 let mut entity_map = EntityMap::new();
693
694 let slot = entity_map.reserve::<TestEntity>();
695 entity_map.insert(slot, TestEntity { i: 1 });
696
697 let slot = entity_map.reserve::<TestEntity>();
698 entity_map.insert(slot, TestEntity { i: 2 });
699
700 let dropped = entity_map.take_dropped();
701 assert_eq!(dropped.len(), 2);
702
703 assert_eq!(
704 dropped
705 .into_iter()
706 .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
707 .collect::<Vec<i32>>(),
708 vec![1, 2],
709 );
710 }
711
712 #[test]
713 fn test_entity_map_weak_upgrade_before_cleanup() {
714 // Tests that weak handles are not upgraded before take_dropped
715 let mut entity_map = EntityMap::new();
716
717 let slot = entity_map.reserve::<TestEntity>();
718 let handle = entity_map.insert(slot, TestEntity { i: 1 });
719 let weak = handle.downgrade();
720 drop(handle);
721
722 let strong = weak.upgrade();
723 assert_eq!(strong, None);
724
725 let dropped = entity_map.take_dropped();
726 assert_eq!(dropped.len(), 1);
727
728 assert_eq!(
729 dropped
730 .into_iter()
731 .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
732 .collect::<Vec<i32>>(),
733 vec![1],
734 );
735 }
736}