refslots 0.1.0

use slotmap::{Key, SlotMap};
use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};

#[cfg(any(test, feature = "recorder"))]
use crate::recorder::{self, Backtrace, HandleId, Recorder, RecorderSnapshot};

/// Tracks strong reference counts for slots identified by a [`Key`].
///
/// The registry only counts references for slots. It doesn't store values. To store values with a
/// slot (by key) another data structure must be used. A [`SecondaryMap`](slotmap::SecondaryMap)
/// works well for most cases.
///
/// Inserting into the registry creates a new slot with a reference count of 1. The key must be
/// given to exactly one [`Handle`](crate::Handle) which will increment the count as it is cloned,
/// and decrement the count when it is dropped. Prefer [`Handle::insert`](crate::Handle::insert)
/// which creates the slot and its owning handle in one operation.
///
/// Dropping the last handle removes the slot and queues its key on the release queue (accessed
/// with [`take_released`](Self::take_released)). This provides a mechanism for the registry owner
/// to clean up any data it associated with the released keys.
///
/// Counters are atomic, so incrementing and decrementing only need `&self`. Inserting, releasing,
/// and taking the release queue require `&mut self`.
///
/// # Panics
///
/// With the `recorder` feature enabled and always in tests, dropping a registry that has live
/// handles will panic.
#[derive(Debug)]
pub struct Registry<K: Key> {
    counters: SlotMap<K, AtomicUsize>,
    release_queue: Vec<K>,

    #[cfg(any(test, feature = "recorder"))]
    recorder: Recorder<K>,
}

impl<K: Key> Registry<K> {
    /// Inserts a new slot with a reference count of 1 and returns its key.
    ///
    /// The returned key must be given to exactly one [`Handle`](crate::Handle).
    pub fn insert(&mut self) -> K {
        self.counters.insert(AtomicUsize::new(1))
    }

    /// Increments the strong reference count of a live slot and returns the previous count.
    ///
    /// # Panics
    ///
    /// Panics if the count is 0. Use [`upgrade`](Self::upgrade) when the slot might be dead.
    pub fn increment(&self, key: K) -> usize {
        let counter = self.counter(key);
        let mut count = counter.load(SeqCst);
        loop {
            assert_ne!(count, 0, "handle: {key:?} over-release detected");
            match counter.compare_exchange_weak(count, count + 1, SeqCst, SeqCst) {
                Ok(prev) => return prev,
                Err(actual) => count = actual,
            }
        }
    }

    /// Increments the strong reference count only if it is nonzero.
    ///
    /// Returns `true` if the count was incremented. Otherwise, returns `false`.
    ///
    /// The check and increment are a single atomic step, so a slot that is being released
    /// cannot be revived. This is the primitive behind
    /// [`WeakHandle::upgrade`](crate::WeakHandle::upgrade).
    pub fn upgrade(&self, key: K) -> bool {
        let Some(counter) = self.counters.get(key) else { return false };
        let mut count = counter.load(SeqCst);
        while count > 0 {
            match counter.compare_exchange_weak(count, count + 1, SeqCst, SeqCst) {
                Ok(_) => return true,
                Err(actual) => count = actual,
            }
        }
        false
    }

    /// Decrements the strong reference count and returns the previous count.
    ///
    /// A return value of 1 means the last reference has been dropped and the caller must follow
    /// with a [`release`](Self::release) of the slot.
    ///
    /// # Panics
    ///
    /// Panics if the count is already zero, indicating the slot has been over-released.
    pub fn decrement(&self, key: K) -> usize {
        let counter = self.counter(key);
        let mut count = counter.load(SeqCst);
        loop {
            assert_ne!(count, 0, "handle: {key:?} over-release detected");
            match counter.compare_exchange_weak(count, count - 1, SeqCst, SeqCst) {
                Ok(prev) => return prev,
                Err(actual) => count = actual,
            }
        }
    }

    /// Removes a slot with a strong reference count of 0 and queues its key for cleanup.
    ///
    /// # Panics
    ///
    /// Panics if the slot is already released or its count is nonzero.
    pub fn release(&mut self, key: K) {
        let count = self.counter(key).load(SeqCst);
        assert_eq!(count, 0, "handle: {key:?} released with count: {count} ≠ 0");
        self.counters.remove(key);
        self.release_queue.push(key);
    }

    /// Drains the keys released since the last call.
    ///
    /// This to clean up anything associated with the released keys. The iterator returned borrows
    /// the registry mutably, so it is recommended to collect it before running callbacks that might
    /// create or drop handles.
    pub fn take_released(&mut self) -> impl Iterator<Item = K> + '_ {
        self.release_queue.drain(..)
    }

    /// Returns the strong reference count of the slot.
    pub fn strong_count(&self, key: K) -> Option<usize> {
        Some(self.counters.get(key)?.load(SeqCst))
    }

    fn counter(&self, key: K) -> &AtomicUsize {
        self.counters.get(key).unwrap_or_else(|| {
            panic!("handle: {key:?} already released");
        })
    }
}

#[cfg(any(test, feature = "recorder"))]
impl<K: Key> Registry<K> {
    /// Records a new handle has been created. See [`Recorder::record_create_handle`].
    pub(crate) fn record_create_handle(
        registry: &std::sync::Weak<parking_lot::RwLock<Self>>,
        key: K,
        label: Option<&'static str>,
    ) -> HandleId {
        let Some(registry) = registry.upgrade() else { return HandleId::DETACHED };
        // Capture backtrace before acquiring the write lock.
        let backtrace = Backtrace::new_unresolved();
        registry.write().recorder.record_create_handle(key, label, backtrace)
    }

    /// Records a handle has been dropped. See [`Recorder::record_drop_handle`].
    pub(crate) fn record_drop_handle(
        registry: &std::sync::Weak<parking_lot::RwLock<Self>>,
        key: K,
        handle_id: HandleId,
    ) {
        let Some(registry) = registry.upgrade() else { return };
        registry.write().recorder.record_drop_handle(key, handle_id);
    }

    /// Captures the current set of live slots in the registry.
    pub fn snapshot(&self) -> RecorderSnapshot<K> {
        self.recorder.snapshot()
    }

    /// Asserts that all handles for the slot `key` have been released.
    ///
    /// # Panics
    ///
    /// Panics if the assertion is false.
    pub fn assert_handles_released(&self, key: K) {
        recorder::assert::handles_released(&self.recorder, key);
    }

    /// Asserts that all handles created after `snapshot` have been released.
    ///
    /// Handles that existed when the snapshot was taken are ignored, even if they are still alive.
    ///
    /// # Panics
    ///
    /// Panics if the assertion is false.
    pub fn assert_handles_released_since(&self, snapshot: &RecorderSnapshot<K>) {
        recorder::assert::handles_released_since(&self.recorder, snapshot);
    }
}

impl<K: Key> Default for Registry<K> {
    fn default() -> Self {
        Self {
            counters: SlotMap::with_key(),
            release_queue: vec![],

            #[cfg(any(test, feature = "recorder"))]
            recorder: Recorder::default(),
        }
    }
}

#[cfg(test)]
mod tests {
    use slotmap::DefaultKey;

    type Registry = super::Registry<DefaultKey>;

    #[test]
    fn insert() {
        let mut registry = Registry::default();
        let key = registry.insert();
        assert_eq!(registry.strong_count(key), Some(1));
        assert!(registry.release_queue.is_empty());
    }

    mod increment {
        use super::*;
        use std::{panic, panic::AssertUnwindSafe};

        #[test]
        fn live() {
            let mut registry = Registry::default();
            let key = registry.insert();
            registry.increment(key);
            assert_eq!(registry.strong_count(key), Some(2));
        }

        #[test]
        fn unlive() {
            let mut registry = Registry::default();
            let key = registry.insert();
            registry.decrement(key);

            let result = panic::catch_unwind(AssertUnwindSafe(|| {
                registry.increment(key);
            }));
            assert!(result.is_err());

            // Make sure count isn't modified.
            assert_eq!(registry.strong_count(key), Some(0));
            // And the slot is not revivable.
            assert!(!registry.upgrade(key));
        }

        #[test]
        #[should_panic(expected = "already released")]
        fn released() {
            let mut registry = Registry::default();
            let key = registry.insert();
            registry.decrement(key);
            registry.release(key);
            registry.increment(key);
        }
    }

    mod upgrade {
        use super::*;

        #[test]
        fn upgrade() {
            let mut registry = Registry::default();
            let key = registry.insert();
            assert!(registry.upgrade(key));
            assert_eq!(registry.strong_count(key), Some(2));
        }

        #[test]
        fn unlive() {
            let mut registry = Registry::default();
            let key = registry.insert();
            registry.decrement(key);
            assert!(!registry.upgrade(key));
        }

        #[test]
        fn released() {
            let mut registry = Registry::default();
            let key = registry.insert();
            registry.decrement(key);
            registry.release(key);
            assert!(!registry.upgrade(key));
        }
    }

    mod decrement {
        use super::*;
        use std::{panic, panic::AssertUnwindSafe};

        #[test]
        fn live() {
            let mut registry = Registry::default();
            let key = registry.insert();
            assert_eq!(registry.decrement(key), 1);
            assert!(registry.release_queue.is_empty());
        }

        #[test]
        fn unlive() {
            let mut registry = Registry::default();
            let key = registry.insert();
            registry.decrement(key);

            let result = panic::catch_unwind(AssertUnwindSafe(|| {
                registry.decrement(key);
            }));
            assert!(result.is_err());

            // Make sure count isn't modified.
            assert_eq!(registry.strong_count(key), Some(0));
            // And the slot is not revivable.
            assert!(!registry.upgrade(key));
        }

        #[test]
        #[should_panic(expected = "already released")]
        fn released() {
            let mut registry = Registry::default();
            let key = registry.insert();
            registry.decrement(key);
            registry.release(key);
            registry.decrement(key);
        }
    }

    mod release {
        use super::*;

        #[test]
        #[should_panic(expected = "released with count")]
        fn live() {
            let mut registry = Registry::default();
            let key = registry.insert();
            registry.release(key);
        }

        #[test]
        fn unlive() {
            let mut registry = Registry::default();
            let key = registry.insert();
            registry.decrement(key);
            registry.release(key);
            assert!(registry.strong_count(key).is_none());
            assert_eq!(registry.release_queue.len(), 1);
        }

        #[test]
        #[should_panic(expected = "already released")]
        fn released() {
            let mut registry = Registry::default();
            let key = registry.insert();
            registry.decrement(key);
            registry.release(key);
            registry.release(key);
        }
    }

    mod take_released {
        use super::*;

        #[test]
        fn drain() {
            let mut registry = Registry::default();
            let keys = (0..3).map(|_| registry.insert()).collect::<Vec<_>>();

            // Decrement and release keys 0 and 1, but keep key 2 as live.
            registry.decrement(keys[0]);
            registry.release(keys[0]);
            registry.decrement(keys[1]);
            registry.release(keys[1]);

            let released = registry.take_released().collect::<Vec<_>>();
            assert_eq!(released, [keys[0], keys[1]]);
            assert_eq!(registry.strong_count(keys[2]), Some(1));
        }

        #[test]
        fn clear() {
            let mut registry = Registry::default();
            let key = registry.insert();
            registry.decrement(key);
            registry.release(key);

            assert_eq!(registry.take_released().count(), 1);
            assert_eq!(registry.take_released().count(), 0);
        }

        #[test]
        fn empty() {
            let mut registry = Registry::default();
            registry.insert();
            assert_eq!(registry.take_released().count(), 0);
        }
    }

    #[test]
    fn stale_key() {
        let mut registry = Registry::default();
        let key1 = registry.insert();
        registry.decrement(key1);
        registry.release(key1);

        let key2 = registry.insert(); // might reuse the slot.
        assert_ne!(key1, key2);
        assert!(!registry.upgrade(key1));
        assert!(registry.strong_count(key1).is_none());
        assert_eq!(registry.strong_count(key2), Some(1))
    }
}