entity_map.rs

  1use crate::{private::Sealed, AnyBox, AppContext, Context, Entity};
  2use anyhow::{anyhow, Result};
  3use derive_more::{Deref, DerefMut};
  4use parking_lot::{RwLock, RwLockUpgradableReadGuard};
  5use slotmap::{SecondaryMap, SlotMap};
  6use std::{
  7    any::{type_name, 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};
 17
 18slotmap::new_key_type! { pub struct EntityId; }
 19
 20impl EntityId {
 21    pub fn as_u64(self) -> u64 {
 22        self.0.as_ffi()
 23    }
 24}
 25
 26impl Display for EntityId {
 27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 28        write!(f, "{}", self)
 29    }
 30}
 31
 32pub(crate) struct EntityMap {
 33    entities: SecondaryMap<EntityId, AnyBox>,
 34    ref_counts: Arc<RwLock<EntityRefCounts>>,
 35}
 36
 37struct EntityRefCounts {
 38    counts: SlotMap<EntityId, AtomicUsize>,
 39    dropped_entity_ids: Vec<EntityId>,
 40}
 41
 42impl EntityMap {
 43    pub fn new() -> Self {
 44        Self {
 45            entities: SecondaryMap::new(),
 46            ref_counts: Arc::new(RwLock::new(EntityRefCounts {
 47                counts: SlotMap::with_key(),
 48                dropped_entity_ids: Vec::new(),
 49            })),
 50        }
 51    }
 52
 53    /// Reserve a slot for an entity, which you can subsequently use with `insert`.
 54    pub fn reserve<T: 'static>(&self) -> Slot<T> {
 55        let id = self.ref_counts.write().counts.insert(1.into());
 56        Slot(Model::new(id, Arc::downgrade(&self.ref_counts)))
 57    }
 58
 59    /// Insert an entity into a slot obtained by calling `reserve`.
 60    pub fn insert<T>(&mut self, slot: Slot<T>, entity: T) -> Model<T>
 61    where
 62        T: 'static + Send,
 63    {
 64        let model = slot.0;
 65        self.entities.insert(model.entity_id, Box::new(entity));
 66        model
 67    }
 68
 69    /// Move an entity to the stack.
 70    pub fn lease<'a, T>(&mut self, model: &'a Model<T>) -> Lease<'a, T> {
 71        self.assert_valid_context(model);
 72        let entity = Some(
 73            self.entities
 74                .remove(model.entity_id)
 75                .expect("Circular entity lease. Is the entity already being updated?"),
 76        );
 77        Lease {
 78            model,
 79            entity,
 80            entity_type: PhantomData,
 81        }
 82    }
 83
 84    /// Return an entity after moving it to the stack.
 85    pub fn end_lease<T>(&mut self, mut lease: Lease<T>) {
 86        self.entities
 87            .insert(lease.model.entity_id, lease.entity.take().unwrap());
 88    }
 89
 90    pub fn read<T: 'static>(&self, model: &Model<T>) -> &T {
 91        self.assert_valid_context(model);
 92        self.entities[model.entity_id].downcast_ref().unwrap()
 93    }
 94
 95    fn assert_valid_context(&self, model: &AnyModel) {
 96        debug_assert!(
 97            Weak::ptr_eq(&model.entity_map, &Arc::downgrade(&self.ref_counts)),
 98            "used a model with the wrong context"
 99        );
100    }
101
102    pub fn take_dropped(&mut self) -> Vec<(EntityId, AnyBox)> {
103        let mut ref_counts = self.ref_counts.write();
104        let dropped_entity_ids = mem::take(&mut ref_counts.dropped_entity_ids);
105
106        dropped_entity_ids
107            .into_iter()
108            .map(|entity_id| {
109                ref_counts.counts.remove(entity_id);
110                (entity_id, self.entities.remove(entity_id).unwrap())
111            })
112            .collect()
113    }
114}
115
116pub struct Lease<'a, T> {
117    entity: Option<AnyBox>,
118    pub model: &'a Model<T>,
119    entity_type: PhantomData<T>,
120}
121
122impl<'a, T: 'static> core::ops::Deref for Lease<'a, T> {
123    type Target = T;
124
125    fn deref(&self) -> &Self::Target {
126        self.entity.as_ref().unwrap().downcast_ref().unwrap()
127    }
128}
129
130impl<'a, T: 'static> core::ops::DerefMut for Lease<'a, T> {
131    fn deref_mut(&mut self) -> &mut Self::Target {
132        self.entity.as_mut().unwrap().downcast_mut().unwrap()
133    }
134}
135
136impl<'a, T> Drop for Lease<'a, T> {
137    fn drop(&mut self) {
138        if self.entity.is_some() {
139            // We don't panic here, because other panics can cause us to drop the lease without ending it cleanly.
140            log::error!("Leases must be ended with EntityMap::end_lease")
141        }
142    }
143}
144
145#[derive(Deref, DerefMut)]
146pub struct Slot<T>(Model<T>);
147
148pub struct AnyModel {
149    pub(crate) entity_id: EntityId,
150    pub(crate) entity_type: TypeId,
151    entity_map: Weak<RwLock<EntityRefCounts>>,
152}
153
154impl AnyModel {
155    fn new(id: EntityId, entity_type: TypeId, entity_map: Weak<RwLock<EntityRefCounts>>) -> Self {
156        Self {
157            entity_id: id,
158            entity_type,
159            entity_map,
160        }
161    }
162
163    pub fn entity_id(&self) -> EntityId {
164        self.entity_id
165    }
166
167    // todo!() added for populating `ProjectItemBuilders` in `load_path` method
168    pub fn type_id(&self) -> TypeId {
169        self.entity_type
170    }
171
172    pub fn downgrade(&self) -> AnyWeakModel {
173        AnyWeakModel {
174            entity_id: self.entity_id,
175            entity_type: self.entity_type,
176            entity_ref_counts: self.entity_map.clone(),
177        }
178    }
179
180    pub fn downcast<T: 'static>(self) -> Result<Model<T>, AnyModel> {
181        if TypeId::of::<T>() == self.entity_type {
182            Ok(Model {
183                any_model: self,
184                entity_type: PhantomData,
185            })
186        } else {
187            Err(self)
188        }
189    }
190}
191
192impl Clone for AnyModel {
193    fn clone(&self) -> Self {
194        if let Some(entity_map) = self.entity_map.upgrade() {
195            let entity_map = entity_map.read();
196            let count = entity_map
197                .counts
198                .get(self.entity_id)
199                .expect("detected over-release of a model");
200            let prev_count = count.fetch_add(1, SeqCst);
201            assert_ne!(prev_count, 0, "Detected over-release of a model.");
202        }
203
204        Self {
205            entity_id: self.entity_id,
206            entity_type: self.entity_type,
207            entity_map: self.entity_map.clone(),
208        }
209    }
210}
211
212impl Drop for AnyModel {
213    fn drop(&mut self) {
214        if let Some(entity_map) = self.entity_map.upgrade() {
215            let entity_map = entity_map.upgradable_read();
216            let count = entity_map
217                .counts
218                .get(self.entity_id)
219                .expect("Detected over-release of a model.");
220            let prev_count = count.fetch_sub(1, SeqCst);
221            assert_ne!(prev_count, 0, "Detected over-release of a model.");
222            if prev_count == 1 {
223                // We were the last reference to this entity, so we can remove it.
224                let mut entity_map = RwLockUpgradableReadGuard::upgrade(entity_map);
225                entity_map.dropped_entity_ids.push(self.entity_id);
226            }
227        }
228    }
229}
230
231impl<T> From<Model<T>> for AnyModel {
232    fn from(model: Model<T>) -> Self {
233        model.any_model
234    }
235}
236
237impl Hash for AnyModel {
238    fn hash<H: Hasher>(&self, state: &mut H) {
239        self.entity_id.hash(state);
240    }
241}
242
243impl PartialEq for AnyModel {
244    fn eq(&self, other: &Self) -> bool {
245        self.entity_id == other.entity_id
246    }
247}
248
249impl Eq for AnyModel {}
250
251impl std::fmt::Debug for AnyModel {
252    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
253        f.debug_struct("AnyModel")
254            .field("entity_id", &self.entity_id.as_u64())
255            .finish()
256    }
257}
258
259#[derive(Deref, DerefMut)]
260pub struct Model<T> {
261    #[deref]
262    #[deref_mut]
263    pub(crate) any_model: AnyModel,
264    pub(crate) entity_type: PhantomData<T>,
265}
266
267unsafe impl<T> Send for Model<T> {}
268unsafe impl<T> Sync for Model<T> {}
269impl<T> Sealed for Model<T> {}
270
271impl<T: 'static> Entity<T> for Model<T> {
272    type Weak = WeakModel<T>;
273
274    fn entity_id(&self) -> EntityId {
275        self.any_model.entity_id
276    }
277
278    fn downgrade(&self) -> Self::Weak {
279        WeakModel {
280            any_model: self.any_model.downgrade(),
281            entity_type: self.entity_type,
282        }
283    }
284
285    fn upgrade_from(weak: &Self::Weak) -> Option<Self>
286    where
287        Self: Sized,
288    {
289        Some(Model {
290            any_model: weak.any_model.upgrade()?,
291            entity_type: weak.entity_type,
292        })
293    }
294}
295
296impl<T: 'static> Model<T> {
297    fn new(id: EntityId, entity_map: Weak<RwLock<EntityRefCounts>>) -> Self
298    where
299        T: 'static,
300    {
301        Self {
302            any_model: AnyModel::new(id, TypeId::of::<T>(), entity_map),
303            entity_type: PhantomData,
304        }
305    }
306
307    /// Downgrade the this to a weak model reference
308    pub fn downgrade(&self) -> WeakModel<T> {
309        // Delegate to the trait implementation to keep behavior in one place.
310        // This method was included to improve method resolution in the presence of
311        // the Model's deref
312        Entity::downgrade(self)
313    }
314
315    /// Convert this into a dynamically typed model.
316    pub fn into_any(self) -> AnyModel {
317        self.any_model
318    }
319
320    pub fn read<'a>(&self, cx: &'a AppContext) -> &'a T {
321        cx.entities.read(self)
322    }
323
324    /// Update the entity referenced by this model with the given function.
325    ///
326    /// The update function receives a context appropriate for its environment.
327    /// When updating in an `AppContext`, it receives a `ModelContext`.
328    /// When updating an a `WindowContext`, it receives a `ViewContext`.
329    pub fn update<C, R>(
330        &self,
331        cx: &mut C,
332        update: impl FnOnce(&mut T, &mut C::ModelContext<'_, T>) -> R,
333    ) -> C::Result<R>
334    where
335        C: Context,
336    {
337        cx.update_model(self, update)
338    }
339}
340
341impl<T> Clone for Model<T> {
342    fn clone(&self) -> Self {
343        Self {
344            any_model: self.any_model.clone(),
345            entity_type: self.entity_type,
346        }
347    }
348}
349
350impl<T> std::fmt::Debug for Model<T> {
351    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
352        write!(
353            f,
354            "Model {{ entity_id: {:?}, entity_type: {:?} }}",
355            self.any_model.entity_id,
356            type_name::<T>()
357        )
358    }
359}
360
361impl<T> Hash for Model<T> {
362    fn hash<H: Hasher>(&self, state: &mut H) {
363        self.any_model.hash(state);
364    }
365}
366
367impl<T> PartialEq for Model<T> {
368    fn eq(&self, other: &Self) -> bool {
369        self.any_model == other.any_model
370    }
371}
372
373impl<T> Eq for Model<T> {}
374
375impl<T> PartialEq<WeakModel<T>> for Model<T> {
376    fn eq(&self, other: &WeakModel<T>) -> bool {
377        self.any_model.entity_id() == other.entity_id()
378    }
379}
380
381#[derive(Clone)]
382pub struct AnyWeakModel {
383    pub(crate) entity_id: EntityId,
384    entity_type: TypeId,
385    entity_ref_counts: Weak<RwLock<EntityRefCounts>>,
386}
387
388impl AnyWeakModel {
389    pub fn entity_id(&self) -> EntityId {
390        self.entity_id
391    }
392
393    pub fn is_upgradable(&self) -> bool {
394        let ref_count = self
395            .entity_ref_counts
396            .upgrade()
397            .and_then(|ref_counts| Some(ref_counts.read().counts.get(self.entity_id)?.load(SeqCst)))
398            .unwrap_or(0);
399        ref_count > 0
400    }
401
402    pub fn upgrade(&self) -> Option<AnyModel> {
403        let entity_map = self.entity_ref_counts.upgrade()?;
404        entity_map
405            .read()
406            .counts
407            .get(self.entity_id)?
408            .fetch_add(1, SeqCst);
409        Some(AnyModel {
410            entity_id: self.entity_id,
411            entity_type: self.entity_type,
412            entity_map: self.entity_ref_counts.clone(),
413        })
414    }
415}
416
417impl<T> From<WeakModel<T>> for AnyWeakModel {
418    fn from(model: WeakModel<T>) -> Self {
419        model.any_model
420    }
421}
422
423impl Hash for AnyWeakModel {
424    fn hash<H: Hasher>(&self, state: &mut H) {
425        self.entity_id.hash(state);
426    }
427}
428
429impl PartialEq for AnyWeakModel {
430    fn eq(&self, other: &Self) -> bool {
431        self.entity_id == other.entity_id
432    }
433}
434
435impl Eq for AnyWeakModel {}
436
437#[derive(Deref, DerefMut)]
438pub struct WeakModel<T> {
439    #[deref]
440    #[deref_mut]
441    any_model: AnyWeakModel,
442    entity_type: PhantomData<T>,
443}
444
445unsafe impl<T> Send for WeakModel<T> {}
446unsafe impl<T> Sync for WeakModel<T> {}
447
448impl<T> Clone for WeakModel<T> {
449    fn clone(&self) -> Self {
450        Self {
451            any_model: self.any_model.clone(),
452            entity_type: self.entity_type,
453        }
454    }
455}
456
457impl<T: 'static> WeakModel<T> {
458    /// Upgrade this weak model reference into a strong model reference
459    pub fn upgrade(&self) -> Option<Model<T>> {
460        // Delegate to the trait implementation to keep behavior in one place.
461        Model::upgrade_from(self)
462    }
463
464    /// Update the entity referenced by this model with the given function if
465    /// the referenced entity still exists. Returns an error if the entity has
466    /// been released.
467    ///
468    /// The update function receives a context appropriate for its environment.
469    /// When updating in an `AppContext`, it receives a `ModelContext`.
470    /// When updating an a `WindowContext`, it receives a `ViewContext`.
471    pub fn update<C, R>(
472        &self,
473        cx: &mut C,
474        update: impl FnOnce(&mut T, &mut C::ModelContext<'_, T>) -> R,
475    ) -> Result<R>
476    where
477        C: Context,
478        Result<C::Result<R>>: crate::Flatten<R>,
479    {
480        crate::Flatten::flatten(
481            self.upgrade()
482                .ok_or_else(|| anyhow!("entity release"))
483                .map(|this| cx.update_model(&this, update)),
484        )
485    }
486}
487
488impl<T> Hash for WeakModel<T> {
489    fn hash<H: Hasher>(&self, state: &mut H) {
490        self.any_model.hash(state);
491    }
492}
493
494impl<T> PartialEq for WeakModel<T> {
495    fn eq(&self, other: &Self) -> bool {
496        self.any_model == other.any_model
497    }
498}
499
500impl<T> Eq for WeakModel<T> {}
501
502impl<T> PartialEq<Model<T>> for WeakModel<T> {
503    fn eq(&self, other: &Model<T>) -> bool {
504        self.entity_id() == other.any_model.entity_id()
505    }
506}