1use crate::Context;
2use anyhow::{anyhow, Result};
3use derive_more::{Deref, DerefMut};
4use parking_lot::{Mutex, RwLock};
5use slotmap::{SecondaryMap, SlotMap};
6use std::{
7 any::Any,
8 marker::PhantomData,
9 sync::{
10 atomic::{AtomicUsize, Ordering::SeqCst},
11 Arc, Weak,
12 },
13};
14
15slotmap::new_key_type! { pub struct EntityId; }
16
17#[derive(Deref, DerefMut)]
18pub struct Lease<T> {
19 #[deref]
20 #[deref_mut]
21 entity: Box<T>,
22 pub id: EntityId,
23}
24
25pub(crate) struct EntityMap {
26 ref_counts: Arc<RwLock<RefCounts>>,
27 entities: Arc<Mutex<SecondaryMap<EntityId, Box<dyn Any + Send + Sync>>>>,
28}
29
30impl EntityMap {
31 pub fn new() -> Self {
32 Self {
33 ref_counts: Arc::new(RwLock::new(SlotMap::with_key())),
34 entities: Arc::new(Mutex::new(SecondaryMap::new())),
35 }
36 }
37
38 pub fn reserve<T: 'static + Send + Sync>(&self) -> Slot<T> {
39 let id = self.ref_counts.write().insert(1.into());
40 Slot(Handle::new(id, Arc::downgrade(&self.ref_counts)))
41 }
42
43 pub fn redeem<T: 'static + Any + Send + Sync>(&self, slot: Slot<T>, entity: T) -> Handle<T> {
44 let handle = slot.0;
45 self.entities.lock().insert(handle.id, Box::new(entity));
46 handle
47 }
48
49 pub fn lease<T: 'static + Send + Sync>(&self, handle: &Handle<T>) -> Lease<T> {
50 let id = handle.id;
51 let entity = self
52 .entities
53 .lock()
54 .remove(id)
55 .expect("Circular entity lease. Is the entity already being updated?")
56 .downcast::<T>()
57 .unwrap();
58 Lease { id, entity }
59 }
60
61 pub fn end_lease<T: 'static + Send + Sync>(&mut self, lease: Lease<T>) {
62 self.entities.lock().insert(lease.id, lease.entity);
63 }
64
65 pub fn weak_handle<T: 'static + Send + Sync>(&self, id: EntityId) -> WeakHandle<T> {
66 WeakHandle {
67 id,
68 entity_type: PhantomData,
69 ref_counts: Arc::downgrade(&self.ref_counts),
70 }
71 }
72}
73
74#[derive(Deref, DerefMut)]
75pub struct Slot<T: Send + Sync + 'static>(Handle<T>);
76
77pub struct Handle<T: Send + Sync> {
78 pub(crate) id: EntityId,
79 entity_type: PhantomData<T>,
80 ref_counts: Weak<RwLock<RefCounts>>,
81}
82
83type RefCounts = SlotMap<EntityId, AtomicUsize>;
84
85impl<T: 'static + Send + Sync> Handle<T> {
86 pub fn new(id: EntityId, ref_counts: Weak<RwLock<RefCounts>>) -> Self {
87 Self {
88 id,
89 entity_type: PhantomData,
90 ref_counts,
91 }
92 }
93
94 pub fn downgrade(&self) -> WeakHandle<T> {
95 WeakHandle {
96 id: self.id,
97 entity_type: self.entity_type,
98 ref_counts: self.ref_counts.clone(),
99 }
100 }
101
102 /// Update the entity referenced by this handle with the given function.
103 ///
104 /// The update function receives a context appropriate for its environment.
105 /// When updating in an `AppContext`, it receives a `ModelContext`.
106 /// When updating an a `WindowContext`, it receives a `ViewContext`.
107 pub fn update<C: Context, R>(
108 &self,
109 cx: &mut C,
110 update: impl FnOnce(&mut T, &mut C::EntityContext<'_, '_, T>) -> R,
111 ) -> C::Result<R> {
112 cx.update_entity(self, update)
113 }
114}
115
116impl<T: Send + Sync> Clone for Handle<T> {
117 fn clone(&self) -> Self {
118 Self {
119 id: self.id,
120 entity_type: PhantomData,
121 ref_counts: self.ref_counts.clone(),
122 }
123 }
124}
125
126impl<T: Send + Sync> Drop for Handle<T> {
127 fn drop(&mut self) {
128 if let Some(ref_counts) = self.ref_counts.upgrade() {
129 if let Some(count) = ref_counts.read().get(self.id) {
130 let prev_count = count.fetch_sub(1, SeqCst);
131 assert_ne!(prev_count, 0, "Detected over-release of a handle.");
132 }
133 }
134 }
135}
136
137pub struct WeakHandle<T> {
138 pub(crate) id: EntityId,
139 pub(crate) entity_type: PhantomData<T>,
140 pub(crate) ref_counts: Weak<RwLock<RefCounts>>,
141}
142
143impl<T: Send + Sync + 'static> WeakHandle<T> {
144 pub fn upgrade(&self, _: &impl Context) -> Option<Handle<T>> {
145 let ref_counts = self.ref_counts.upgrade()?;
146 ref_counts.read().get(self.id).unwrap().fetch_add(1, SeqCst);
147 Some(Handle {
148 id: self.id,
149 entity_type: self.entity_type,
150 ref_counts: self.ref_counts.clone(),
151 })
152 }
153
154 /// Update the entity referenced by this handle with the given function if
155 /// the referenced entity still exists. Returns an error if the entity has
156 /// been released.
157 ///
158 /// The update function receives a context appropriate for its environment.
159 /// When updating in an `AppContext`, it receives a `ModelContext`.
160 /// When updating an a `WindowContext`, it receives a `ViewContext`.
161 pub fn update<C: Context, R>(
162 &self,
163 cx: &mut C,
164 update: impl FnOnce(&mut T, &mut C::EntityContext<'_, '_, T>) -> R,
165 ) -> Result<R>
166 where
167 Result<C::Result<R>>: crate::Flatten<R>,
168 {
169 crate::Flatten::flatten(
170 self.upgrade(cx)
171 .ok_or_else(|| anyhow!("entity release"))
172 .map(|this| cx.update_entity(&this, update)),
173 )
174 }
175}