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