refslots 0.1.0

Reference-counted handles to slots in a slotmap
Documentation
use crate::{Handle, Registry};
use parking_lot::RwLock;
use slotmap::Key;
use std::{
    cmp::Ordering,
    fmt,
    fmt::Debug,
    hash::{Hash, Hasher},
    sync::Weak,
};

/// Weak handle to a slot in [`Registry`].
///
/// A weak handle doesn't contribute to the reference count, so it doesn't keep a slot alive. Call
/// [`upgrade`](Self::upgrade) to get a strong [`Handle`] which succeeds if the slot is still alive.
///
/// Like [`Handle`], a weak handle compares, orders, and hashes by key only and holds a weak
/// reference to the registry that created it.
///
/// # Locking
///
/// [`upgrade`](Self::upgrade) and [`is_upgradeable`](Self::is_upgradeable) acquire a read lock on
/// the registry. Calling these while holding a guard on the registry's lock can deadlock.
pub struct WeakHandle<K: Key, M = ()> {
    pub(crate) key: K,
    pub(crate) meta: M,
    pub(crate) registry: Weak<RwLock<Registry<K>>>,
}

impl<K: Key, M: Clone> WeakHandle<K, M> {
    pub(crate) fn downgrade(handle: &Handle<K, M>) -> Self {
        Self { key: handle.key, meta: handle.meta.clone(), registry: Weak::clone(&handle.registry) }
    }

    /// Attempts to get a strong handle, incrementing the reference count.
    ///
    /// Returns `None` if every strong handle has been dropped or the registry no longer exists. The
    /// liveness check and increment are atomic so a slot being released cannot be revived.
    pub fn upgrade(&self) -> Option<Handle<K, M>> {
        if !self.registry.upgrade()?.read().upgrade(self.key) {
            return None;
        }
        Some(Handle {
            key: self.key,
            meta: self.meta.clone(),
            registry: Weak::clone(&self.registry),

            #[cfg(any(test, feature = "recorder"))]
            handle_id: Registry::record_create_handle(&self.registry, self.key, None),
        })
    }
}

impl<K: Key, M> WeakHandle<K, M> {
    /// Returns `true` if the slot pointed to by the handle is alive.
    ///
    /// The result can be stale by the time it is used, since another thread might interleave
    /// dropping the last strong handle. To use the slot, call [`upgrade`](Self::upgrade).
    pub fn is_upgradeable(&self) -> bool {
        let Some(registry) = self.registry.upgrade() else { return false };
        let count = registry.read().strong_count(self.key);
        count.unwrap_or_default() > 0
    }

    /// Returns the key of the slot pointed to by the handle.
    pub fn key(&self) -> K {
        self.key
    }

    /// Returns the metadata stored with the handle.
    pub fn meta(&self) -> &M {
        &self.meta
    }
}

impl<K: Key, M: Clone> Clone for WeakHandle<K, M> {
    fn clone(&self) -> Self {
        Self { key: self.key, meta: self.meta.clone(), registry: Weak::clone(&self.registry) }
    }
}

impl<K: Key, M> PartialEq for WeakHandle<K, M> {
    fn eq(&self, other: &Self) -> bool {
        self.key == other.key
    }
}

impl<K: Key, M> Eq for WeakHandle<K, M> {}

impl<K: Key, M> PartialOrd for WeakHandle<K, M> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<K: Key, M> Ord for WeakHandle<K, M> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.key.cmp(&other.key)
    }
}

impl<K: Key, M> Hash for WeakHandle<K, M> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.key.hash(state);
    }
}

impl<K: Key, M> Debug for WeakHandle<K, M> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_tuple("WeakHandle").field(&self.key).finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use slotmap::DefaultKey;
    use static_assertions::assert_impl_all;
    use std::{any::TypeId, sync::Arc};

    type Handle = super::Handle<DefaultKey>;
    type WeakHandle = super::WeakHandle<DefaultKey>;
    type RegistryCell = Arc<RwLock<Registry<DefaultKey>>>;

    mod upgrade {
        use super::*;

        #[test]
        fn live() {
            let registry = RegistryCell::default();
            let handle = Handle::insert(&registry, ());
            let weak = handle.downgrade();
            let strong = weak.upgrade().unwrap();
            assert_eq!(handle, strong);
            assert_eq!(registry.read().strong_count(handle.key), Some(2));
        }

        #[test]
        fn unlive() {
            let registry = RegistryCell::default();
            let weak = {
                let handle = Handle::insert(&registry, ());
                handle.downgrade()
            };
            assert!(registry.read().strong_count(weak.key).is_none());
            assert!(!weak.is_upgradeable());
            assert!(weak.upgrade().is_none());
        }

        #[test]
        fn dropped_registry() {
            let weak = {
                let registry = RegistryCell::default();
                let handle = Handle::insert(&registry, ());
                handle.downgrade()
            };
            assert!(!weak.is_upgradeable());
            assert!(weak.upgrade().is_none());
        }
    }

    #[test]
    fn clone() {
        let registry = RegistryCell::default();
        let handle = Handle::insert(&registry, ());
        let weak = handle.downgrade();
        let _clone = weak.clone();
        assert_eq!(registry.read().strong_count(handle.key), Some(1));
    }

    #[test]
    fn drop() {
        let registry = RegistryCell::default();
        let handle = Handle::insert(&registry, ());
        let key = {
            let weak = handle.downgrade();
            weak.key
        };
        assert_eq!(registry.read().strong_count(key), Some(1));
    }

    #[test]
    fn send_sync() {
        assert_impl_all!(WeakHandle: Send, Sync);
        assert_impl_all!(super::WeakHandle<DefaultKey, TypeId>: Send, Sync);
    }

    #[test]
    fn debug() {
        let registry = RegistryCell::default();
        let handle = Handle::insert(&registry, ());
        let weak = handle.downgrade();
        assert_eq!(format!("{weak:?}"), "WeakHandle(DefaultKey(1v1))");
    }
}