1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
//! A library to quickly get the live/total/max counts of allocated instances.
//!
//! # Example
//!
//! ```
//! # if cfg!(not(feature = "enable")) { return; }
//!
//! #[derive(Default)]
//! struct Widget {
//!   _c: countme::Count<Self>,
//! }
//!
//! countme::enable(true);
//!
//! let w1 = Widget::default();
//! let w2 = Widget::default();
//! let w3 = Widget::default();
//! drop(w1);
//!
//! let counts = countme::get::<Widget>();
//! assert_eq!(counts.live, 2);
//! assert_eq!(counts.max_live, 3);
//! assert_eq!(counts.total, 3);
//!
//! eprintln!("{}", countme::get_all());
//! ```
//!
//! # Configuration
//!
//! By default, the implementation compiles to no-ops. Therefore it is possible
//! to include `Count` fields into library types.
//!
//! The `enable` cargo feature ungates the counting code. The feature can be
//! enabled anywhere in the crate graph.
//!
//! At run-time, the counters are controlled with [`enable`] function. Counting
//! is disabled by default. Call `enable(true)` early in `main` to enable:
//!
//! ```rust
//! fn main() {
//!     countme::enable(std::env::var("COUNTME").is_ok());
//! }
//! ```
//!
//! The code is optimized for the case where counting is not enabled at runtime
//! (counting is a relaxed load and a branch to a function call).
//!
//! The `print_at_exit` Cargo feature uses `atexit` call to print final counts
//! before the program exits (it also enables counting at runtime). Use it only
//! when you can't modify the main to print counts -- `atexit` is not guaranteed
//! to work with rust's runtime.
#[cfg(feature = "enable")]
mod imp;

use std::{fmt, marker::PhantomData};

#[derive(Debug, Clone, PartialEq, Eq, Default)]
#[non_exhaustive]
pub struct Counts {
    /// The total number of tokens created.
    pub total: usize,
    /// The historical maximum of the `live` count.
    pub max_live: usize,
    /// The number of tokens which were created, but are not destroyed yet.
    pub live: usize,
}

/// Store this inside your struct as `_c: countme::Count<Self>`.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Count<T> {
    ghost: PhantomData<fn(T)>,
}

impl<T> Default for Count<T> {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Clone for Count<T> {
    #[inline]
    fn clone(&self) -> Self {
        Self::new()
    }
}

impl<T> Count<T> {
    /// Create new `Count`, incrementing the corresponding count.
    #[inline]
    pub fn new() -> Count<T> {
        #[cfg(feature = "enable")]
        imp::inc::<T>();
        Count { ghost: PhantomData }
    }
}

impl<T> Drop for Count<T> {
    #[inline]
    fn drop(&mut self) {
        #[cfg(feature = "enable")]
        imp::dec::<T>();
    }
}

/// Enable or disable counting at runtime.
///
/// Counting is enabled by default.
pub fn enable(_yes: bool) {
    #[cfg(feature = "enable")]
    imp::enable(_yes);
}

/// Returns the counts for the `T` type.
#[inline]
pub fn get<T>() -> Counts {
    #[cfg(feature = "enable")]
    {
        return imp::get::<T>();
    }
    #[cfg(not(feature = "enable"))]
    {
        return Counts::default();
    }
}

/// Returns a collection of counts for all types.
pub fn get_all() -> AllCounts {
    #[cfg(feature = "enable")]
    {
        return imp::get_all();
    }
    #[cfg(not(feature = "enable"))]
    {
        return AllCounts::default();
    }
}

/// A collection of counts for all types.
#[derive(Default, Clone, Debug)]
pub struct AllCounts {
    entries: Vec<(&'static str, Counts)>,
}

impl fmt::Display for AllCounts {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fn sep(mut n: usize) -> String {
            let mut groups = Vec::new();
            while n >= 1000 {
                groups.push(format!("{:03}", n % 1000));
                n /= 1000;
            }
            groups.push(n.to_string());
            groups.reverse();
            groups.join("_")
        }

        if self.entries.is_empty() {
            return if cfg!(feature = "enable") {
                writeln!(f, "all counts are zero")
            } else {
                writeln!(f, "counts are disabled")
            };
        }
        let max_width =
            self.entries.iter().map(|(name, _count)| name.chars().count()).max().unwrap_or(0);
        for (name, counts) in &self.entries {
            writeln!(
                f,
                "{:<max_width$}  {:>12} {:>12} {:>12}",
                name,
                sep(counts.total),
                sep(counts.max_live),
                sep(counts.live),
                max_width = max_width
            )?;
        }
        writeln!(
            f,
            "{:<max_width$}  {:>12} {:>12} {:>12}",
            "",
            "total",
            "max_live",
            "live",
            max_width = max_width
        )?;
        Ok(())
    }
}