swansong 0.3.6

use super::guard::GuardInfo;
use event_listener::{Event, EventListener, IntoNotification};
use std::sync::{
    atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
    Arc, Mutex, Weak,
};

/// Registry of provenance records for the guards created on this node.
///
/// Entries are `Weak`: a dropped guard's entry simply stops upgrading, so
/// `Guard::drop` performs no accounting here. Dead entries are swept
/// opportunistically during registration (amortized O(1) per registration,
/// keeping memory bounded by the live guard count) and at read time.
#[derive(Debug, Default)]
struct GuardRegistry {
    guards: Vec<Weak<GuardInfo>>,
    sweep_threshold: usize,
}

impl GuardRegistry {
    const MIN_SWEEP_THRESHOLD: usize = 16;

    fn register(&mut self, guard: Weak<GuardInfo>) {
        if self.guards.len() >= self.sweep_threshold.max(Self::MIN_SWEEP_THRESHOLD) {
            self.sweep();
        }
        self.guards.push(guard);
    }

    fn sweep(&mut self) {
        self.guards.retain(|guard| guard.strong_count() > 0);
        self.sweep_threshold = self.guards.len() * 2;
    }

    fn live(&self) -> impl Iterator<Item = Arc<GuardInfo>> + '_ {
        self.guards.iter().filter_map(Weak::upgrade)
    }
}

/// Packed subset nonemptiness: low 32 bits = `local_guards`, high 32 bits = `nonempty_kids`.
/// The subset (this node + descendants) is nonempty iff `packed != 0`.
const LOCAL_ONE: u64 = 1;
const KIDS_ONE: u64 = 1 << 32;
const LOCAL_MASK: u64 = 0xFFFF_FFFF;

#[derive(Debug, Default)]
pub(crate) struct Inner {
    stop_event: Event,
    zero_event: Event,
    packed: AtomicU64,
    stopped: AtomicBool,
    swansong_count: AtomicUsize,
    /// Strong refs to every ancestor, deepest first: `[parent, grandparent, ..., root]`.
    /// Empty for a root. Strong (not Weak) so that a guard's decrement walk always
    /// has valid ancestors to update, even if every user-held Swansong to an
    /// intermediate ancestor has been dropped.
    ancestors: Vec<Arc<Inner>>,
    children: Mutex<Vec<Weak<Inner>>>,
    guard_registry: Mutex<GuardRegistry>,
}

impl Inner {
    pub(crate) fn new_root() -> Arc<Self> {
        Arc::new(Self {
            swansong_count: AtomicUsize::new(1),
            ..Self::default()
        })
    }

    pub(crate) fn new_child(parent: &Arc<Self>) -> Arc<Self> {
        let mut ancestors = Vec::with_capacity(parent.ancestors.len() + 1);
        ancestors.push(Arc::clone(parent));
        ancestors.extend(parent.ancestors.iter().cloned());

        let inner = Arc::new(Self {
            swansong_count: AtomicUsize::new(1),
            ancestors,
            ..Self::default()
        });
        parent.add_child(Arc::downgrade(&inner));
        if parent.is_stopped() {
            inner.stop();
        }
        inner
    }

    fn add_child(&self, child: Weak<Inner>) {
        let mut children = self.children.lock().unwrap();
        children.retain(|c| c.upgrade().is_some_and(|c| !c.is_complete()));
        children.push(child);
    }

    pub(crate) fn is_stopped(&self) -> bool {
        self.stopped.load(Ordering::SeqCst)
    }

    pub(crate) fn is_stopped_relaxed(&self) -> bool {
        self.stopped.load(Ordering::Relaxed)
    }

    pub(crate) fn is_zero(&self) -> bool {
        self.packed.load(Ordering::SeqCst) == 0
    }

    pub(crate) fn is_zero_relaxed(&self) -> bool {
        self.packed.load(Ordering::Relaxed) == 0
    }

    pub(crate) fn is_complete(&self) -> bool {
        self.is_stopped() && self.is_zero()
    }

    pub(crate) fn is_root(&self) -> bool {
        self.ancestors.is_empty()
    }

    pub(crate) fn listen_zero(&self) -> EventListener {
        self.zero_event.listen()
    }

    pub(crate) fn listen_stop(&self) -> EventListener {
        self.stop_event.listen()
    }

    pub(crate) fn swansong_clone(&self) {
        self.swansong_count.fetch_add(1, Ordering::SeqCst);
    }

    /// Decrement the swansong handle count. Returns `true` if this was the last handle.
    pub(crate) fn swansong_drop(&self) -> bool {
        self.swansong_count.fetch_sub(1, Ordering::SeqCst) == 1
    }

    pub(crate) fn stop(&self) {
        log::trace!("intending to stop");

        if self.stopped.swap(true, Ordering::SeqCst) {
            log::trace!("was already stopped");
            return;
        }

        log::trace!("stopped");
        self.stop_event.notify(usize::MAX.relaxed());

        // Snapshot children under the lock, then release before recursing to avoid
        // holding the lock across child.stop() calls.
        let children: Vec<Arc<Inner>> = {
            let guard = self.children.lock().unwrap();
            guard.iter().filter_map(Weak::upgrade).collect()
        };
        for child in children {
            child.stop();
        }

        if self.is_zero() {
            self.zero_event.notify(usize::MAX.relaxed());
        }
    }

    /// Called by `Guard::new`. Records the guard's provenance, increments this
    /// node's local guard count and, if the subset transitioned from empty to
    /// nonempty, propagates the transition up the ancestor chain.
    pub(crate) fn increment_guard(&self, info: Weak<GuardInfo>) {
        self.guard_registry.lock().unwrap().register(info);
        let prev = self.packed.fetch_add(LOCAL_ONE, Ordering::SeqCst);
        log::trace!("incremented local: {prev:#x} -> {:#x}", prev + LOCAL_ONE);
        if prev == 0 {
            for ancestor in &self.ancestors {
                let prev = ancestor.packed.fetch_add(KIDS_ONE, Ordering::SeqCst);
                if prev != 0 {
                    break;
                }
            }
        }
    }

    /// Called by `Guard::drop`. Decrements this node's local guard count and, if the
    /// subset transitioned from nonempty to empty, notifies this node's zero event
    /// (if stopped) and propagates the transition up the ancestor chain.
    pub(crate) fn decrement_guard(&self) {
        let prev = self.packed.fetch_sub(LOCAL_ONE, Ordering::SeqCst);
        let new = prev - LOCAL_ONE;
        log::trace!("decremented local: {prev:#x} -> {new:#x}");
        if new == 0 {
            if self.is_stopped() {
                self.zero_event.notify(usize::MAX.relaxed());
            }
            for ancestor in &self.ancestors {
                let prev = ancestor.packed.fetch_sub(KIDS_ONE, Ordering::SeqCst);
                let new = prev - KIDS_ONE;
                if new == 0 {
                    if ancestor.is_stopped() {
                        ancestor.zero_event.notify(usize::MAX.relaxed());
                    }
                } else {
                    break;
                }
            }
        }
    }

    /// O(subtree) walk summing real guards across self and all live descendants.
    pub(crate) fn guard_count_subtree(&self) -> usize {
        let local = usize::try_from(self.packed.load(Ordering::Relaxed) & LOCAL_MASK)
            .expect("local_guards fits in usize");
        let children_sum: usize = self
            .children
            .lock()
            .unwrap()
            .iter()
            .filter_map(Weak::upgrade)
            .map(|c| c.guard_count_subtree())
            .sum();
        local + children_sum
    }

    /// O(subtree) walk collecting provenance records for all live guards in
    /// self and all live descendants.
    pub(crate) fn guard_infos_subtree(&self, out: &mut Vec<Arc<GuardInfo>>) {
        {
            let mut registry = self.guard_registry.lock().unwrap();
            registry.sweep();
            out.extend(registry.live());
        }
        let children: Vec<Arc<Inner>> = {
            let children = self.children.lock().unwrap();
            children.iter().filter_map(Weak::upgrade).collect()
        };
        for child in children {
            child.guard_infos_subtree(out);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::panic::Location;

    #[test]
    fn registry_memory_stays_bounded_without_reads() {
        let mut registry = GuardRegistry::default();
        for _ in 0..1000 {
            let info = GuardInfo::new(Location::caller());
            registry.register(Arc::downgrade(&info));
            drop(info);
        }
        assert!(registry.guards.len() <= GuardRegistry::MIN_SWEEP_THRESHOLD);
    }

    #[test]
    fn registry_tracks_live_guards() {
        let mut registry = GuardRegistry::default();
        let infos: Vec<_> = (0..100)
            .map(|_| {
                let info = GuardInfo::new(Location::caller());
                registry.register(Arc::downgrade(&info));
                info
            })
            .collect();
        assert_eq!(registry.live().count(), 100);
        drop(infos);
        assert_eq!(registry.live().count(), 0);
    }
}