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