captrack 0.1.0

Capacity telemetry for Rust collections — call-site macros that record peak capacity, with zero overhead when disabled.
use std::collections::BTreeSet;

use crate::registry;
use crate::IntoInner;

/// A `BTreeSet<T>` wrapper that records creation count and capacity samples.
///
/// Like `TrackedBTreeMap`, the capacity hint is accepted but ignored for the
/// inner allocation.
///
/// # Samples record `len()`, NOT peak occupancy
///
/// B-tree capacity is not observable.  On every `Drop` (or `IntoIterator` /
/// `From` conversion), `inner.len()` is pushed as the sample — this is the
/// element count **at the moment of Drop**, not the maximum ever observed.
///
/// For sets that may shrink before Drop (e.g. via `BTreeSet::clear` or
/// repeated `remove` calls), the recorded sample **undercounts** the true
/// peak.  Use `BTreeSet::len()` at the known peak point if accurate peak
/// tracking is required.
pub struct TrackedBTreeSet<T: Ord> {
    inner: BTreeSet<T>,
    #[allow(dead_code)]
    name: &'static str,
    file: &'static str,
    line: u32,
    column: u32,
}

impl<T: Ord> TrackedBTreeSet<T> {
    pub fn new_named(
        _cap_hint: usize,
        name: &'static str,
        file: &'static str,
        line: u32,
        column: u32,
    ) -> Self {
        registry::record_creation(name, file, line, column);
        Self {
            inner: BTreeSet::new(),
            name,
            file,
            line,
            column,
        }
    }

    /// Wrap an already-constructed `BTreeSet<T>` for capacity telemetry.
    ///
    /// Records creation in the registry; `inner` is moved as-is.
    /// The sample metric at `Drop` is `inner.len()` (BTreeSet has no capacity).
    #[inline]
    pub fn wrap_from(
        inner: BTreeSet<T>,
        name: &'static str,
        file: &'static str,
        line: u32,
        column: u32,
    ) -> Self {
        registry::record_creation(name, file, line, column);
        Self {
            inner,
            name,
            file,
            line,
            column,
        }
    }
}

impl<T: Ord> std::ops::Deref for TrackedBTreeSet<T> {
    type Target = BTreeSet<T>;
    fn deref(&self) -> &BTreeSet<T> {
        &self.inner
    }
}

impl<T: Ord> std::ops::DerefMut for TrackedBTreeSet<T> {
    fn deref_mut(&mut self) -> &mut BTreeSet<T> {
        &mut self.inner
    }
}

impl<T: Ord> Drop for TrackedBTreeSet<T> {
    fn drop(&mut self) {
        registry::record_sample(self.file, self.line, self.column, self.inner.len());
    }
}

impl<T: Ord> From<TrackedBTreeSet<T>> for BTreeSet<T> {
    fn from(tracked: TrackedBTreeSet<T>) -> BTreeSet<T> {
        registry::record_sample(
            tracked.file,
            tracked.line,
            tracked.column,
            tracked.inner.len(),
        );
        // SAFETY: `tracked` is owned and forgotten below; ptr::read bit-copies `inner`.
        let inner = unsafe { std::ptr::read(&tracked.inner) };
        std::mem::forget(tracked);
        inner
    }
}

impl<T: Ord> IntoInner for TrackedBTreeSet<T> {
    type Inner = BTreeSet<T>;
    #[inline]
    fn into_inner(self) -> BTreeSet<T> {
        BTreeSet::from(self)
    }
}

impl<T: Ord> IntoIterator for TrackedBTreeSet<T> {
    type Item = T;
    type IntoIter = std::collections::btree_set::IntoIter<T>;

    fn into_iter(self) -> Self::IntoIter {
        registry::record_sample(self.file, self.line, self.column, self.inner.len());
        // SAFETY: `self` is owned and forgotten below; ptr::read bit-copies `inner`.
        let inner = unsafe { std::ptr::read(&self.inner) };
        std::mem::forget(self);
        inner.into_iter()
    }
}