mt-debug-counters 0.1.4

use lazy_static::lazy_static;
use parking_lot::Mutex;
use std::cmp::{max, min};
use std::collections::HashMap;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::Weak;

#[doc(hidden)]
#[macro_export]
macro_rules! declare_counter_i64_impl {
    ($name:expr, $mode:ty, $reset:expr, $extra:expr) => {
        $crate::counter::AtomicCounter::<$mode> {
            __get_counter: || {
                thread_local! {
                    static COUNTER: std::sync::Arc<std::sync::atomic::AtomicI64> = {
                        let arc = std::sync::Arc::new(std::sync::atomic::AtomicI64::new(0));
                        let mut list = $crate::counter::__COUNTERS_LIST.lock();
                        let cvec = list.entry($name.to_string()).or_insert((Vec::new(), 0, <$crate::counter::AtomicCounter<$mode> as $crate::counter::__CounterType>::MODE, $reset));
                        cvec.0.push(std::sync::Arc::downgrade(&arc));
                        arc
                    }
                }
                use std::ops::Deref;
                COUNTER.with(|c| {
                    unsafe {
                        &*(c.deref() as *const std::sync::atomic::AtomicI64)
                    }
                })
            },
            __extra: $extra,
            _phantom: std::marker::PhantomData
        }
    }
}

#[macro_export]
macro_rules! declare_counter_i64 {
    ($name:literal, $mode:ty, $reset:expr) => {
        $crate::declare_counter_i64_impl!($name, $mode, $reset, ())
    };
}

pub(crate) static COUNTER_SUFFIX: &str = "$COUNTER_83uRij";

#[macro_export]
macro_rules! declare_avg_counter_i64 {
    ($name:literal, $reset:expr) => {
        $crate::declare_counter_i64_impl!(
            $name,
            $crate::counter::AvgMode,
            $reset,
            $crate::declare_counter_i64_impl!(
                concat!($name, "$COUNTER_83uRij"),
                $crate::counter::SumMode,
                $reset,
                ()
            )
        )
    };
}

lazy_static! {
    #[doc(hidden)]
    pub static ref __COUNTERS_LIST: Mutex<HashMap<String, (Vec<Weak<AtomicI64>>, i64, __AcMode, bool)>> = Mutex::new(HashMap::new());
}

#[doc(hidden)]
pub trait AtomicCounterMode {
    type Extra;
}

pub struct SumMode {}
impl AtomicCounterMode for SumMode {
    type Extra = ();
}
pub struct MaxMode {}
impl AtomicCounterMode for MaxMode {
    type Extra = ();
}
pub struct MinMode {}
impl AtomicCounterMode for MinMode {
    type Extra = ();
}
pub struct AvgMode {}
impl AtomicCounterMode for AvgMode {
    type Extra = AtomicCounter<SumMode>;
}

pub struct AtomicCounter<MODE: AtomicCounterMode> {
    #[doc(hidden)]
    pub __get_counter: fn() -> &'static AtomicI64,
    #[doc(hidden)]
    pub __extra: MODE::Extra,
    #[doc(hidden)]
    pub _phantom: PhantomData<MODE>,
}

#[doc(hidden)]
#[derive(Eq, PartialEq)]
pub enum __AcMode {
    SUM,
    MAX,
    MIN,
    AVG,
}

#[doc(hidden)]
pub trait __CounterType {
    const MODE: __AcMode;
}

impl AtomicCounter<SumMode> {
    #[inline(always)]
    pub fn inc(&self) {
        self.inc_by(1);
    }

    #[inline(always)]
    pub fn inc_by(&self, value: i64) {
        (self.__get_counter)().fetch_add(value, Ordering::Relaxed);
    }

    #[inline(always)]
    pub fn sub(&self, value: i64) {
        (self.__get_counter)().fetch_sub(value, Ordering::Relaxed);
    }
}

impl AtomicCounter<AvgMode> {
    #[inline(always)]
    pub fn add_value(&self, value: i64) {
        (self.__get_counter)().fetch_add(value, Ordering::Relaxed);
        self.__extra.inc();
    }
}

impl __CounterType for AtomicCounter<SumMode> {
    const MODE: __AcMode = __AcMode::SUM;
}

impl AtomicCounter<MaxMode> {
    #[inline(always)]
    pub fn max(&self, val: i64) {
        (self.__get_counter)().fetch_max(val, Ordering::Relaxed);
    }
}

impl __CounterType for AtomicCounter<MaxMode> {
    const MODE: __AcMode = __AcMode::MAX;
}

impl AtomicCounter<MinMode> {
    #[inline(always)]
    pub fn min(&self, val: i64) {
        (self.__get_counter)().fetch_min(val, Ordering::Relaxed);
    }
}

impl __CounterType for AtomicCounter<MinMode> {
    const MODE: __AcMode = __AcMode::MIN;
}

impl __CounterType for AtomicCounter<AvgMode> {
    const MODE: __AcMode = __AcMode::AVG;
}

pub struct AtomicCounterGuardSum<'a> {
    value: i64,
    counter: &'a AtomicCounter<SumMode>,
}

impl<'a> AtomicCounterGuardSum<'a> {
    pub fn new(counter: &'a AtomicCounter<SumMode>, value: i64) -> Self {
        counter.inc_by(value);
        Self { value, counter }
    }
}

impl<'a> Drop for AtomicCounterGuardSum<'a> {
    fn drop(&mut self) {
        self.counter.sub(self.value);
    }
}

pub fn get_counter_value(name: &str) -> (i64, i64) {
    let mut counters = __COUNTERS_LIST.lock();

    let (ref mut vec, part_value, mode, reset) = if let Some(val) = counters.get_mut(name) {
        val
    } else {
        return (0, 0);
    };

    let reset_value = match mode {
        __AcMode::SUM => 0,
        __AcMode::MAX => 0,
        __AcMode::MIN => i64::MAX,
        __AcMode::AVG => 0,
    };

    if *reset {
        *part_value = reset_value;
    }

    let mut result = *part_value;

    vec.retain(|val| {
        if val.strong_count() > 0 {
            if let Some(value) = val.upgrade() {
                let value = value.swap(reset_value, Ordering::Relaxed);

                match mode {
                    __AcMode::SUM => {
                        result += value;
                    }
                    __AcMode::MAX => {
                        result = max(result, value);
                    }
                    __AcMode::MIN => {
                        result = min(result, value);
                    }
                    __AcMode::AVG => {
                        result += value;
                    }
                }
                true
            } else {
                false
            }
        } else {
            false
        }
    });

    *part_value = result;

    let is_average = __AcMode::AVG == *mode;
    drop(counters);

    let counter = if is_average {
        get_counter_value(&(name.to_string() + COUNTER_SUFFIX)).0
    } else {
        0
    };

    (result, counter)
}

#[cfg(test)]
mod tests {
    use super::SumMode;
    #[test]
    fn alloc_test() {
        let _sum_counter = declare_counter_i64!("test_counter", SumMode, false);
    }
}