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
17pub(crate) struct EntityMap {
18 ref_counts: Arc<RwLock<RefCounts>>,
19 entities: Arc<Mutex<SecondaryMap<EntityId, Box<dyn Any + Send + Sync>>>>,
20}
21
22impl EntityMap {
23 pub fn new() -> Self {
24 Self {
25 ref_counts: Arc::new(RwLock::new(SlotMap::with_key())),
26 entities: Arc::new(Mutex::new(SecondaryMap::new())),
27 }
28 }
29
30 /// Reserve a slot for an entity, which you can subsequently use with `insert`.
31 pub fn reserve<T: 'static + Send + Sync>(&self) -> Slot<T> {
32 let id = self.ref_counts.write().insert(1.into());
33 Slot(Handle::new(id, Arc::downgrade(&self.ref_counts)))
34 }
35
36 /// Insert an entity into a slot obtained by calling `reserve`.
37 pub fn insert<T: 'static + Any + Send + Sync>(&self, slot: Slot<T>, entity: T) -> Handle<T> {
38 let handle = slot.0;
39 self.entities.lock().insert(handle.id, Box::new(entity));
40 handle
41 }
42
43 /// Move an entity to the stack.
44 pub fn lease<T: 'static + Send + Sync>(&self, handle: &Handle<T>) -> Lease<T> {
45 let id = handle.id;
46 let entity = Some(
47 self.entities
48 .lock()
49 .remove(id)
50 .expect("Circular entity lease. Is the entity already being updated?")
51 .downcast::<T>()
52 .unwrap(),
53 );
54 Lease { id, entity }
55 }
56
57 /// Return an entity after moving it to the stack.
58 pub fn end_lease<T: 'static + Send + Sync>(&mut self, mut lease: Lease<T>) {
59 self.entities
60 .lock()
61 .insert(lease.id, lease.entity.take().unwrap());
62 }
63
64 pub fn weak_handle<T: 'static + Send + Sync>(&self, id: EntityId) -> WeakHandle<T> {
65 WeakHandle {
66 id,
67 entity_type: PhantomData,
68 ref_counts: Arc::downgrade(&self.ref_counts),
69 }
70 }
71}
72
73pub struct Lease<T> {
74 entity: Option<Box<T>>,
75 pub id: EntityId,
76}
77
78impl<T> core::ops::Deref for Lease<T> {
79 type Target = T;
80
81 fn deref(&self) -> &Self::Target {
82 self.entity.as_ref().unwrap()
83 }
84}
85
86impl<T> core::ops::DerefMut for Lease<T> {
87 fn deref_mut(&mut self) -> &mut Self::Target {
88 self.entity.as_mut().unwrap()
89 }
90}
91
92impl<T> Drop for Lease<T> {
93 fn drop(&mut self) {
94 assert!(
95 self.entity.is_none(),
96 "Leases must be ended with EntityMap::end_lease"
97 );
98 }
99}
100
101#[derive(Deref, DerefMut)]
102pub struct Slot<T: Send + Sync + 'static>(Handle<T>);
103
104pub struct Handle<T: Send + Sync> {
105 pub(crate) id: EntityId,
106 entity_type: PhantomData<T>,
107 ref_counts: Weak<RwLock<RefCounts>>,
108}
109
110type RefCounts = SlotMap<EntityId, AtomicUsize>;
111
112impl<T: 'static + Send + Sync> Handle<T> {
113 pub fn new(id: EntityId, ref_counts: Weak<RwLock<RefCounts>>) -> Self {
114 Self {
115 id,
116 entity_type: PhantomData,
117 ref_counts,
118 }
119 }
120
121 pub fn downgrade(&self) -> WeakHandle<T> {
122 WeakHandle {
123 id: self.id,
124 entity_type: self.entity_type,
125 ref_counts: self.ref_counts.clone(),
126 }
127 }
128
129 /// Update the entity referenced by this handle with the given function.
130 ///
131 /// The update function receives a context appropriate for its environment.
132 /// When updating in an `AppContext`, it receives a `ModelContext`.
133 /// When updating an a `WindowContext`, it receives a `ViewContext`.
134 pub fn update<C: Context, R>(
135 &self,
136 cx: &mut C,
137 update: impl FnOnce(&mut T, &mut C::EntityContext<'_, '_, T>) -> R,
138 ) -> C::Result<R> {
139 cx.update_entity(self, update)
140 }
141}
142
143impl<T: Send + Sync> Clone for Handle<T> {
144 fn clone(&self) -> Self {
145 Self {
146 id: self.id,
147 entity_type: PhantomData,
148 ref_counts: self.ref_counts.clone(),
149 }
150 }
151}
152
153impl<T: Send + Sync> Drop for Handle<T> {
154 fn drop(&mut self) {
155 if let Some(_ref_counts) = self.ref_counts.upgrade() {
156 // todo!()
157 // if let Some(count) = ref_counts.read().get(self.id) {
158 // let prev_count = count.fetch_sub(1, SeqCst);
159 // assert_ne!(prev_count, 0, "Detected over-release of a handle.");
160 // }
161 }
162 }
163}
164
165pub struct WeakHandle<T> {
166 pub(crate) id: EntityId,
167 pub(crate) entity_type: PhantomData<T>,
168 pub(crate) ref_counts: Weak<RwLock<RefCounts>>,
169}
170
171impl<T: Send + Sync + 'static> WeakHandle<T> {
172 pub fn upgrade(&self, _: &impl Context) -> Option<Handle<T>> {
173 let ref_counts = self.ref_counts.upgrade()?;
174 ref_counts.read().get(self.id).unwrap().fetch_add(1, SeqCst);
175 Some(Handle {
176 id: self.id,
177 entity_type: self.entity_type,
178 ref_counts: self.ref_counts.clone(),
179 })
180 }
181
182 /// Update the entity referenced by this handle with the given function if
183 /// the referenced entity still exists. Returns an error if the entity has
184 /// been released.
185 ///
186 /// The update function receives a context appropriate for its environment.
187 /// When updating in an `AppContext`, it receives a `ModelContext`.
188 /// When updating an a `WindowContext`, it receives a `ViewContext`.
189 pub fn update<C: Context, R>(
190 &self,
191 cx: &mut C,
192 update: impl FnOnce(&mut T, &mut C::EntityContext<'_, '_, T>) -> R,
193 ) -> Result<R>
194 where
195 Result<C::Result<R>>: crate::Flatten<R>,
196 {
197 crate::Flatten::flatten(
198 self.upgrade(cx)
199 .ok_or_else(|| anyhow!("entity release"))
200 .map(|this| cx.update_entity(&this, update)),
201 )
202 }
203}