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 // todo!()
130 // if let Some(count) = ref_counts.read().get(self.id) {
131 // let prev_count = count.fetch_sub(1, SeqCst);
132 // assert_ne!(prev_count, 0, "Detected over-release of a handle.");
133 // }
134 }
135 }
136}
137
138pub struct WeakHandle<T> {
139 pub(crate) id: EntityId,
140 pub(crate) entity_type: PhantomData<T>,
141 pub(crate) ref_counts: Weak<RwLock<RefCounts>>,
142}
143
144impl<T: Send + Sync + 'static> WeakHandle<T> {
145 pub fn upgrade(&self, _: &impl Context) -> Option<Handle<T>> {
146 let ref_counts = self.ref_counts.upgrade()?;
147 ref_counts.read().get(self.id).unwrap().fetch_add(1, SeqCst);
148 Some(Handle {
149 id: self.id,
150 entity_type: self.entity_type,
151 ref_counts: self.ref_counts.clone(),
152 })
153 }
154
155 /// Update the entity referenced by this handle with the given function if
156 /// the referenced entity still exists. Returns an error if the entity has
157 /// been released.
158 ///
159 /// The update function receives a context appropriate for its environment.
160 /// When updating in an `AppContext`, it receives a `ModelContext`.
161 /// When updating an a `WindowContext`, it receives a `ViewContext`.
162 pub fn update<C: Context, R>(
163 &self,
164 cx: &mut C,
165 update: impl FnOnce(&mut T, &mut C::EntityContext<'_, '_, T>) -> R,
166 ) -> Result<R>
167 where
168 Result<C::Result<R>>: crate::Flatten<R>,
169 {
170 crate::Flatten::flatten(
171 self.upgrade(cx)
172 .ok_or_else(|| anyhow!("entity release"))
173 .map(|this| cx.update_entity(&this, update)),
174 )
175 }
176}