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 /// Return 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 let this = 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 this
259 }
260}
261
262impl Drop for AnyModel {
263 fn drop(&mut self) {
264 if let Some(entity_map) = self.entity_map.upgrade() {
265 let entity_map = entity_map.upgradable_read();
266 let count = entity_map
267 .counts
268 .get(self.entity_id)
269 .expect("detected over-release of a handle.");
270 let prev_count = count.fetch_sub(1, SeqCst);
271 assert_ne!(prev_count, 0, "Detected over-release of a model.");
272 if prev_count == 1 {
273 // We were the last reference to this entity, so we can remove it.
274 let mut entity_map = RwLockUpgradableReadGuard::upgrade(entity_map);
275 entity_map.dropped_entity_ids.push(self.entity_id);
276 }
277 }
278
279 #[cfg(any(test, feature = "test-support"))]
280 if let Some(entity_map) = self.entity_map.upgrade() {
281 entity_map
282 .write()
283 .leak_detector
284 .handle_released(self.entity_id, self.handle_id)
285 }
286 }
287}
288
289impl<T> From<Model<T>> for AnyModel {
290 fn from(model: Model<T>) -> Self {
291 model.any_model
292 }
293}
294
295impl Hash for AnyModel {
296 fn hash<H: Hasher>(&self, state: &mut H) {
297 self.entity_id.hash(state);
298 }
299}
300
301impl PartialEq for AnyModel {
302 fn eq(&self, other: &Self) -> bool {
303 self.entity_id == other.entity_id
304 }
305}
306
307impl Eq for AnyModel {}
308
309impl std::fmt::Debug for AnyModel {
310 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
311 f.debug_struct("AnyModel")
312 .field("entity_id", &self.entity_id.as_u64())
313 .finish()
314 }
315}
316
317/// A strong, well typed reference to a struct which is managed
318/// by GPUI
319#[derive(Deref, DerefMut)]
320pub struct Model<T> {
321 #[deref]
322 #[deref_mut]
323 pub(crate) any_model: AnyModel,
324 pub(crate) entity_type: PhantomData<T>,
325}
326
327unsafe impl<T> Send for Model<T> {}
328unsafe impl<T> Sync for Model<T> {}
329impl<T> Sealed for Model<T> {}
330
331impl<T: 'static> Entity<T> for Model<T> {
332 type Weak = WeakModel<T>;
333
334 fn entity_id(&self) -> EntityId {
335 self.any_model.entity_id
336 }
337
338 fn downgrade(&self) -> Self::Weak {
339 WeakModel {
340 any_model: self.any_model.downgrade(),
341 entity_type: self.entity_type,
342 }
343 }
344
345 fn upgrade_from(weak: &Self::Weak) -> Option<Self>
346 where
347 Self: Sized,
348 {
349 Some(Model {
350 any_model: weak.any_model.upgrade()?,
351 entity_type: weak.entity_type,
352 })
353 }
354}
355
356impl<T: 'static> Model<T> {
357 fn new(id: EntityId, entity_map: Weak<RwLock<EntityRefCounts>>) -> Self
358 where
359 T: 'static,
360 {
361 Self {
362 any_model: AnyModel::new(id, TypeId::of::<T>(), entity_map),
363 entity_type: PhantomData,
364 }
365 }
366
367 /// Downgrade the this to a weak model reference
368 pub fn downgrade(&self) -> WeakModel<T> {
369 // Delegate to the trait implementation to keep behavior in one place.
370 // This method was included to improve method resolution in the presence of
371 // the Model's deref
372 Entity::downgrade(self)
373 }
374
375 /// Convert this into a dynamically typed model.
376 pub fn into_any(self) -> AnyModel {
377 self.any_model
378 }
379
380 /// Grab a reference to this entity from the context.
381 pub fn read<'a>(&self, cx: &'a AppContext) -> &'a T {
382 cx.entities.read(self)
383 }
384
385 /// Read the entity referenced by this model with the given function.
386 pub fn read_with<R, C: Context>(
387 &self,
388 cx: &C,
389 f: impl FnOnce(&T, &AppContext) -> R,
390 ) -> C::Result<R> {
391 cx.read_model(self, f)
392 }
393
394 /// Update the entity referenced by this model with the given function.
395 ///
396 /// The update function receives a context appropriate for its environment.
397 /// When updating in an `AppContext`, it receives a `ModelContext`.
398 /// When updating an a `WindowContext`, it receives a `ViewContext`.
399 pub fn update<C, R>(
400 &self,
401 cx: &mut C,
402 update: impl FnOnce(&mut T, &mut ModelContext<'_, T>) -> R,
403 ) -> C::Result<R>
404 where
405 C: Context,
406 {
407 cx.update_model(self, update)
408 }
409}
410
411impl<T> Clone for Model<T> {
412 fn clone(&self) -> Self {
413 Self {
414 any_model: self.any_model.clone(),
415 entity_type: self.entity_type,
416 }
417 }
418}
419
420impl<T> std::fmt::Debug for Model<T> {
421 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
422 write!(
423 f,
424 "Model {{ entity_id: {:?}, entity_type: {:?} }}",
425 self.any_model.entity_id,
426 type_name::<T>()
427 )
428 }
429}
430
431impl<T> Hash for Model<T> {
432 fn hash<H: Hasher>(&self, state: &mut H) {
433 self.any_model.hash(state);
434 }
435}
436
437impl<T> PartialEq for Model<T> {
438 fn eq(&self, other: &Self) -> bool {
439 self.any_model == other.any_model
440 }
441}
442
443impl<T> Eq for Model<T> {}
444
445impl<T> PartialEq<WeakModel<T>> for Model<T> {
446 fn eq(&self, other: &WeakModel<T>) -> bool {
447 self.any_model.entity_id() == other.entity_id()
448 }
449}
450
451/// A type erased, weak reference to a model.
452#[derive(Clone)]
453pub struct AnyWeakModel {
454 pub(crate) entity_id: EntityId,
455 entity_type: TypeId,
456 entity_ref_counts: Weak<RwLock<EntityRefCounts>>,
457}
458
459impl AnyWeakModel {
460 /// Get the entity ID associated with this weak reference.
461 pub fn entity_id(&self) -> EntityId {
462 self.entity_id
463 }
464
465 /// Check if this weak handle can be upgraded, or if the model has already been dropped
466 pub fn is_upgradable(&self) -> bool {
467 let ref_count = self
468 .entity_ref_counts
469 .upgrade()
470 .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
471 .unwrap_or(0);
472 ref_count > 0
473 }
474
475 /// Upgrade this weak model reference to a strong reference.
476 pub fn upgrade(&self) -> Option<AnyModel> {
477 let ref_counts = &self.entity_ref_counts.upgrade()?;
478 let ref_counts = ref_counts.read();
479 let ref_count = ref_counts.counts.get(self.entity_id)?;
480
481 // entity_id is in dropped_entity_ids
482 if ref_count.load(SeqCst) == 0 {
483 return None;
484 }
485 ref_count.fetch_add(1, SeqCst);
486 drop(ref_counts);
487
488 Some(AnyModel {
489 entity_id: self.entity_id,
490 entity_type: self.entity_type,
491 entity_map: self.entity_ref_counts.clone(),
492 #[cfg(any(test, feature = "test-support"))]
493 handle_id: self
494 .entity_ref_counts
495 .upgrade()
496 .unwrap()
497 .write()
498 .leak_detector
499 .handle_created(self.entity_id),
500 })
501 }
502
503 /// Assert that model referenced by this weak handle has been released.
504 #[cfg(any(test, feature = "test-support"))]
505 pub fn assert_released(&self) {
506 self.entity_ref_counts
507 .upgrade()
508 .unwrap()
509 .write()
510 .leak_detector
511 .assert_released(self.entity_id);
512
513 if self
514 .entity_ref_counts
515 .upgrade()
516 .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
517 .is_some()
518 {
519 panic!(
520 "entity was recently dropped but resources are retained until the end of the effect cycle."
521 )
522 }
523 }
524}
525
526impl<T> From<WeakModel<T>> for AnyWeakModel {
527 fn from(model: WeakModel<T>) -> Self {
528 model.any_model
529 }
530}
531
532impl Hash for AnyWeakModel {
533 fn hash<H: Hasher>(&self, state: &mut H) {
534 self.entity_id.hash(state);
535 }
536}
537
538impl PartialEq for AnyWeakModel {
539 fn eq(&self, other: &Self) -> bool {
540 self.entity_id == other.entity_id
541 }
542}
543
544impl Eq for AnyWeakModel {}
545
546/// A weak reference to a model of the given type.
547#[derive(Deref, DerefMut)]
548pub struct WeakModel<T> {
549 #[deref]
550 #[deref_mut]
551 any_model: AnyWeakModel,
552 entity_type: PhantomData<T>,
553}
554
555unsafe impl<T> Send for WeakModel<T> {}
556unsafe impl<T> Sync for WeakModel<T> {}
557
558impl<T> Clone for WeakModel<T> {
559 fn clone(&self) -> Self {
560 Self {
561 any_model: self.any_model.clone(),
562 entity_type: self.entity_type,
563 }
564 }
565}
566
567impl<T: 'static> WeakModel<T> {
568 /// Upgrade this weak model reference into a strong model reference
569 pub fn upgrade(&self) -> Option<Model<T>> {
570 // Delegate to the trait implementation to keep behavior in one place.
571 Model::upgrade_from(self)
572 }
573
574 /// Update the entity referenced by this model with the given function if
575 /// the referenced entity still exists. Returns an error if the entity has
576 /// been released.
577 pub fn update<C, R>(
578 &self,
579 cx: &mut C,
580 update: impl FnOnce(&mut T, &mut ModelContext<'_, T>) -> R,
581 ) -> Result<R>
582 where
583 C: Context,
584 Result<C::Result<R>>: crate::Flatten<R>,
585 {
586 crate::Flatten::flatten(
587 self.upgrade()
588 .ok_or_else(|| anyhow!("entity release"))
589 .map(|this| cx.update_model(&this, update)),
590 )
591 }
592
593 /// Reads the entity referenced by this model with the given function if
594 /// the referenced entity still exists. Returns an error if the entity has
595 /// been released.
596 pub fn read_with<C, R>(&self, cx: &C, read: impl FnOnce(&T, &AppContext) -> R) -> Result<R>
597 where
598 C: Context,
599 Result<C::Result<R>>: crate::Flatten<R>,
600 {
601 crate::Flatten::flatten(
602 self.upgrade()
603 .ok_or_else(|| anyhow!("entity release"))
604 .map(|this| cx.read_model(&this, read)),
605 )
606 }
607}
608
609impl<T> Hash for WeakModel<T> {
610 fn hash<H: Hasher>(&self, state: &mut H) {
611 self.any_model.hash(state);
612 }
613}
614
615impl<T> PartialEq for WeakModel<T> {
616 fn eq(&self, other: &Self) -> bool {
617 self.any_model == other.any_model
618 }
619}
620
621impl<T> Eq for WeakModel<T> {}
622
623impl<T> PartialEq<Model<T>> for WeakModel<T> {
624 fn eq(&self, other: &Model<T>) -> bool {
625 self.entity_id() == other.any_model.entity_id()
626 }
627}
628
629#[cfg(any(test, feature = "test-support"))]
630lazy_static::lazy_static! {
631 static ref LEAK_BACKTRACE: bool =
632 std::env::var("LEAK_BACKTRACE").map_or(false, |b| !b.is_empty());
633}
634
635#[cfg(any(test, feature = "test-support"))]
636#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq)]
637pub(crate) struct HandleId {
638 id: u64, // id of the handle itself, not the pointed at object
639}
640
641#[cfg(any(test, feature = "test-support"))]
642pub(crate) struct LeakDetector {
643 next_handle_id: u64,
644 entity_handles: HashMap<EntityId, HashMap<HandleId, Option<backtrace::Backtrace>>>,
645}
646
647#[cfg(any(test, feature = "test-support"))]
648impl LeakDetector {
649 #[track_caller]
650 pub fn handle_created(&mut self, entity_id: EntityId) -> HandleId {
651 let id = util::post_inc(&mut self.next_handle_id);
652 let handle_id = HandleId { id };
653 let handles = self.entity_handles.entry(entity_id).or_default();
654 handles.insert(
655 handle_id,
656 LEAK_BACKTRACE.then(|| backtrace::Backtrace::new_unresolved()),
657 );
658 handle_id
659 }
660
661 pub fn handle_released(&mut self, entity_id: EntityId, handle_id: HandleId) {
662 let handles = self.entity_handles.entry(entity_id).or_default();
663 handles.remove(&handle_id);
664 }
665
666 pub fn assert_released(&mut self, entity_id: EntityId) {
667 let handles = self.entity_handles.entry(entity_id).or_default();
668 if !handles.is_empty() {
669 for (_, backtrace) in handles {
670 if let Some(mut backtrace) = backtrace.take() {
671 backtrace.resolve();
672 eprintln!("Leaked handle: {:#?}", backtrace);
673 } else {
674 eprintln!("Leaked handle: export LEAK_BACKTRACE to find allocation site");
675 }
676 }
677 panic!();
678 }
679 }
680}
681
682#[cfg(test)]
683mod test {
684 use crate::EntityMap;
685
686 struct TestEntity {
687 pub i: i32,
688 }
689
690 #[test]
691 fn test_entity_map_slot_assignment_before_cleanup() {
692 // Tests that slots are not re-used before take_dropped.
693 let mut entity_map = EntityMap::new();
694
695 let slot = entity_map.reserve::<TestEntity>();
696 entity_map.insert(slot, TestEntity { i: 1 });
697
698 let slot = entity_map.reserve::<TestEntity>();
699 entity_map.insert(slot, TestEntity { i: 2 });
700
701 let dropped = entity_map.take_dropped();
702 assert_eq!(dropped.len(), 2);
703
704 assert_eq!(
705 dropped
706 .into_iter()
707 .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
708 .collect::<Vec<i32>>(),
709 vec![1, 2],
710 );
711 }
712
713 #[test]
714 fn test_entity_map_weak_upgrade_before_cleanup() {
715 // Tests that weak handles are not upgraded before take_dropped
716 let mut entity_map = EntityMap::new();
717
718 let slot = entity_map.reserve::<TestEntity>();
719 let handle = entity_map.insert(slot, TestEntity { i: 1 });
720 let weak = handle.downgrade();
721 drop(handle);
722
723 let strong = weak.upgrade();
724 assert_eq!(strong, None);
725
726 let dropped = entity_map.take_dropped();
727 assert_eq!(dropped.len(), 1);
728
729 assert_eq!(
730 dropped
731 .into_iter()
732 .map(|(_, entity)| entity.downcast::<TestEntity>().unwrap().i)
733 .collect::<Vec<i32>>(),
734 vec![1],
735 );
736 }
737}