use std::collections::{BTreeMap, btree_map::Entry};
#[cfg(feature = "stats")]
use crate::util::events;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Counter {
Value(u64),
Overflow,
}
impl Counter {
pub fn new() -> Self {
Counter::Value(0)
}
pub fn inc(self) -> Self {
match self {
Counter::Overflow => Counter::Overflow,
Counter::Value(n) => {
let incremented = n.checked_add(1);
incremented.map_or(Counter::Overflow, Counter::Value)
}
}
}
pub fn merge(self, other: Self) -> Self {
match (self, other) {
(Counter::Value(left), Counter::Value(right)) => {
let sum = left.checked_add(right);
sum.map_or(Counter::Overflow, Counter::Value)
}
_ => Counter::Overflow,
}
}
pub fn value(self) -> Option<u64> {
match self {
Counter::Overflow => None,
Counter::Value(n) => Some(n),
}
}
}
impl Default for Counter {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Stat(pub BTreeMap<String, Counter>);
impl Stat {
pub fn new() -> Self {
let counters = BTreeMap::new();
Stat(counters)
}
pub fn inc(&mut self, value: String) {
let counter_entry = self.0.entry(value).or_default();
*counter_entry = counter_entry.inc();
}
pub fn merge(mut self, other: Self) -> Self {
if self.0.len() < other.0.len() {
other.merge(self)
} else {
for (value, right_counter) in other.0.into_iter() {
match self.0.entry(value) {
Entry::Vacant(entry) => {
entry.insert(right_counter);
}
Entry::Occupied(entry) => {
let left_counter = *entry.get();
*entry.into_mut() = left_counter.merge(right_counter);
}
}
}
self
}
}
pub fn total_counter(&self) -> Counter {
self.0
.values()
.fold(Counter::new(), |left, &right| left.merge(right))
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Stats(pub BTreeMap<&'static str, Stat>);
impl Stats {
pub fn new() -> Self {
Stats(BTreeMap::new())
}
pub fn inc(&mut self, key: &'static str, value: String) {
let stat_entry = self.0.entry(key).or_default();
stat_entry.inc(value);
}
pub fn merge(mut self, other: Self) -> Self {
if self.0.len() < other.0.len() {
other.merge(self)
} else {
for (key, right_stat) in other.0.into_iter() {
let left_stat = self.0.remove(key);
let stat = match left_stat {
None => right_stat,
Some(left_stat) => left_stat.merge(right_stat),
};
self.0.insert(key, stat);
}
self
}
}
}
#[cfg(feature = "stats")]
impl events::Events for Stats {
fn new() -> Self {
Stats::new()
}
fn take(&mut self) -> Self {
let first_key = self.0.keys().next().cloned();
match first_key {
None => Stats::new(),
Some(first_key) => {
let stats = self.0.split_off(first_key);
Stats(stats)
}
}
}
}
#[cfg(feature = "stats")]
thread_local! {
static LOCAL: events::Stack<Stats> = events::new_stack();
}
pub fn collect<R>(f: impl FnOnce() -> R) -> (R, Stats) {
#[cfg(feature = "stats")]
{
events::collect(&LOCAL, f)
}
#[cfg(not(feature = "stats"))]
{
(f(), Stats::new())
}
}
pub fn enabled() -> bool {
#[cfg(feature = "stats")]
{
events::enabled(&LOCAL)
}
#[cfg(not(feature = "stats"))]
{
false
}
}
pub fn inc(key: &'static str, value: impl FnOnce() -> String) {
#[cfg(feature = "stats")]
{
events::modify(&LOCAL, move |stack| {
let value = value();
let len = stack.len();
stack[0..len - 1]
.iter_mut()
.for_each(|stats| stats.inc(key, value.clone()));
stack[len - 1].inc(key, value);
});
}
#[cfg(not(feature = "stats"))]
{
let _ = key;
let _ = value;
}
}
#[cfg(test)]
mod tests {
use crate::stats::Counter::{self, Overflow, Value};
#[test]
fn counter_inc_examples() {
assert_eq!(Counter::new().inc(), Value(1));
assert_eq!(Value(u64::MAX).inc(), Overflow);
}
#[test]
fn counter_merge_examples() {
assert_eq!(Overflow.merge(Overflow), Overflow);
assert_eq!(Overflow.merge(Counter::new()), Overflow);
assert_eq!(Counter::new().merge(Overflow), Overflow);
assert_eq!(Value(1).merge(Value(1)), Value(2));
}
#[cfg(feature = "stats")]
#[test]
fn stats_take_takes_all_elements() {
use crate::stats::{Stat, Stats};
use crate::util::events::Events;
let mut stat1 = Stat::new();
let mut stat2 = Stat::new();
let mut stats = Stats::new();
stat1.0.insert("foofoo".to_string(), Value(1));
stat2.0.insert("barbar".to_string(), Overflow);
stats.0.insert("foo", stat1);
stats.0.insert("bar", stat2);
assert_eq!(stats, stats.clone().take());
}
}