refslots 0.1.0

use crate::{WeakHandle, registry::Registry};
use parking_lot::{RwLock, RwLockUpgradableReadGuard};
use slotmap::Key;
use std::{
    cmp::Ordering,
    fmt::{self, Debug},
    hash::{Hash, Hasher},
    sync::{Arc, Weak},
};

#[cfg(any(test, feature = "recorder"))]
use crate::recorder::HandleId;

/// Strong, reference-counted handle to a slot in a [`Registry`].
///
/// A slot is live while at least one `Handle` that points to it exists. Cloning a handle increments
/// the count and dropping a handle decrements it. The last drop releases the slot and adds its key
/// to the registry's release queue. Handles can be cloned and dropped from any thread provided `M`
/// is `Send` and `Sync`.
///
/// `M` is metadata stored with the handle, such as a `TypeId` for dynamically typed entities. It is
/// cloned along with the handle but doesn't take part in equality, ordering, or hashing which use
/// only the key.
///
/// A handle holds a weak reference to its registry. If the registry is dropped then cloning and
/// dropping the handle do nothing.
///
/// # Locking
///
/// Cloning acquires a read lock on the registry. Dropping acquires an upgradable read lock that
/// becomes a write lock when the last reference is dropped. [`insert`](Self::insert) acquires a
/// write lock. Cloning, dropping, or inserting while holding a guard on the registry's lock can
/// deadlock.
pub struct Handle<K: Key, M = ()> {
    pub(crate) key: K,
    pub(crate) meta: M,
    pub(crate) registry: Weak<RwLock<Registry<K>>>,

    #[cfg(any(test, feature = "recorder"))]
    pub(crate) handle_id: HandleId,
}

impl<K: Key, M> Handle<K, M> {
    /// Inserts a new slot in `registry` and returns the handle that owns its initial reference.
    pub fn insert(registry: &Arc<RwLock<Registry<K>>>, meta: M) -> Self {
        Self::do_insert(registry, meta, None)
    }

    /// Like [`insert`](Self::insert), but names the slot in leak reports.
    ///
    /// The label has no effect unless the `recorder` feature is enabled.
    pub fn insert_labeled(
        registry: &Arc<RwLock<Registry<K>>>,
        meta: M,
        label: &'static str,
    ) -> Self {
        Self::do_insert(registry, meta, Some(label))
    }

    #[cfg_attr(not(any(test, feature = "recorder")), allow(unused_variables))]
    fn do_insert(
        registry: &Arc<RwLock<Registry<K>>>,
        meta: M,
        label: Option<&'static str>,
    ) -> Self {
        let key = registry.write().insert();
        let registry = Arc::downgrade(registry);

        #[cfg(any(test, feature = "recorder"))]
        let handle_id = Registry::record_create_handle(&registry, key, label);

        Self {
            key,
            meta,
            registry,

            #[cfg(any(test, feature = "recorder"))]
            handle_id,
        }
    }

    /// Creates a [`WeakHandle`] to the slot pointed to by the handle.
    pub fn downgrade(&self) -> WeakHandle<K, M>
    where
        M: Clone,
    {
        WeakHandle::downgrade(self)
    }

    /// 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
    }

    /// Returns `true` if the handle was created by `registry`.
    ///
    /// Compares by pointer identity which detects a handle being used with a registry other than
    /// the one that created it.
    pub fn check_context(&self, registry: &Weak<RwLock<Registry<K>>>) -> bool {
        Weak::ptr_eq(&self.registry, registry)
    }
}

impl<K: Key, M: Clone> Clone for Handle<K, M> {
    fn clone(&self) -> Self {
        if let Some(registry) = self.registry.upgrade() {
            registry.read().increment(self.key);
        }
        Self {
            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> Drop for Handle<K, M> {
    fn drop(&mut self) {
        if let Some(registry) = self.registry.upgrade() {
            let counts = registry.upgradable_read();
            if counts.decrement(self.key) == 1 {
                // This handle is the last reference to the slot. Release it.
                RwLockUpgradableReadGuard::upgrade(counts).release(self.key);
            }
        }

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

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

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

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

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

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

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

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

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

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

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

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

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

    #[test]
    fn drop_cloned() {
        let registry = RegistryCell::default();
        let handle = {
            let handle = Handle::insert(&registry, ());
            handle.clone()
        };
        assert_eq!(registry.read().strong_count(handle.key), Some(1));
        assert_eq!(registry.write().take_released().count(), 0);
    }

    mod recorder {
        use super::*;

        #[test]
        fn tracks_insert() {
            let registry = RegistryCell::default();
            let handle = Handle::insert(&registry, ());

            let snapshot = registry.read().snapshot();
            assert_eq!(snapshot.slots[&handle.key].handle_count, 1);
        }

        #[test]
        fn tracks_clone() {
            let registry = RegistryCell::default();
            let handle = Handle::insert(&registry, ());
            let _clone = handle.clone();

            let snapshot = registry.read().snapshot();
            assert_eq!(snapshot.slots[&handle.key].handle_count, 2);
        }

        #[test]
        fn tracks_upgrade() {
            let registry = RegistryCell::default();
            let handle = Handle::insert(&registry, ());
            let weak = handle.downgrade();
            let _upgrade = weak.upgrade().unwrap();

            let snapshot = registry.read().snapshot();
            assert_eq!(snapshot.slots[&handle.key].handle_count, 2);
        }

        #[test]
        #[should_panic(expected = "1 live handles created since last snapshot")]
        fn detects_leaked_clone() {
            let registry = RegistryCell::default();
            let handle = Handle::insert(&registry, ());
            let snapshot = registry.read().snapshot();

            let _clone = handle.clone();
            registry.read().assert_handles_released_since(&snapshot);
        }

        #[test]
        #[should_panic(expected = "1 slots with undropped handles")]
        fn drop_registry_unreleased() {
            let _handle = {
                let registry = RegistryCell::default();
                Handle::insert(&registry, ())
            };
        }
    }

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

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