Merge branch 'harness' of https://github.com/wenyuzhao/MallocKit into harness

Author: wenyuzhao

mallockit/src/lib.rs

@@ -12,6 +12,7 @@
 #![feature(effects)]
 #![feature(asm_const)]
 #![feature(const_refs_to_cell)]
+#![feature(const_refs_to_static)]
 
 extern crate mallockit_macros;
 pub extern crate spin;

mallockit/src/stat.rs

@@ -1,104 +1,253 @@
 use std::{
     alloc::Layout,
     sync::{
-        atomic::{AtomicUsize, Ordering},
+        atomic::{AtomicBool, Ordering},
         Arc,
     },
 };
 
-use spin::Mutex;
+use atomic::Atomic;
+use spin::{Mutex, Once};
 
-use crate::{
-    space::meta::{Meta, Vec},
-    util::Lazy,
-};
+use crate::space::meta::{Meta, Vec};
 
-static COUNTERS: Mutex<Vec<Arc<Counter, Meta>>> = Mutex::new(Vec::new_in(Meta));
+pub static DEFAULT_COUNTER_GROUP: CounterGroup = CounterGroup::new("default");
+pub static ALL_GROUPS: Mutex<Vec<&'static CounterGroup>> = Mutex::new(Vec::new_in(Meta));
 
-pub type CounterRef = Lazy<Arc<Counter, Meta>>;
+pub struct CounterGroup {
+    name: &'static str,
+    counters: Mutex<Vec<Arc<dyn DynCounter, Meta>>>,
+    report_fn: Option<fn()>,
+    registered: AtomicBool,
+}
 
-pub const fn define_counter<const NAME: &'static str>() -> Lazy<Arc<Counter, Meta>> {
-    Lazy::new(|| {
-        let c = Arc::new_in(Counter::new(NAME), Meta);
-        COUNTERS.lock().push(c.clone());
-        c
-    })
+impl CounterGroup {
+    pub const fn new(name: &'static str) -> Self {
+        Self {
+            name,
+            counters: Mutex::new(Vec::new_in(Meta)),
+            report_fn: None,
+            registered: AtomicBool::new(false),
+        }
+    }
+
+    pub const fn with_report_fn(mut self, report_fn: fn()) -> Self {
+        self.report_fn = Some(report_fn);
+        self
+    }
+
+    pub const fn new_counter<T: Default + ToString + Copy + 'static>(
+        &'static self,
+        name: &'static str,
+    ) -> Counter<T> {
+        Counter::new_grouped(name, self)
+    }
+
+    fn add_counter(&'static self, counter: Arc<dyn DynCounter, Meta>) {
+        if self
+            .registered
+            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
+            .is_ok()
+        {
+            ALL_GROUPS.lock().push(self);
+        }
+        self.counters.lock().push(counter);
+    }
+
+    pub(crate) fn report(&self) {
+        if let Some(report_fn) = self.report_fn.as_ref() {
+            report_fn();
+        } else {
+            eprintln!("{}:", self.name);
+            while let Some(c) = self.counters.lock().pop() {
+                eprintln!("  {}: {}", c.name(), c.format_value());
+            }
+        }
+    }
+}
+
+#[allow(unused)]
+trait DynCounter: 'static + Sync + Send {
+    fn name(&self) -> &'static str;
+    fn format_value(&self) -> String;
+}
+
+struct CounterImpl<T: Default + ToString + Copy + 'static> {
+    name: &'static str,
+    value: Atomic<T>,
+}
+
+impl<T: Default + ToString + Copy + 'static> CounterImpl<T> {
+    pub fn new(name: &'static str) -> Self {
+        Self {
+            name,
+            value: Atomic::new(T::default()),
+        }
+    }
+}
+
+unsafe impl<T: Default + ToString + Copy + 'static> Send for CounterImpl<T> {}
+unsafe impl<T: Default + ToString + Copy + 'static> Sync for CounterImpl<T> {}
+
+impl<T: Default + ToString + Copy + 'static> DynCounter for CounterImpl<T> {
+    fn name(&self) -> &'static str {
+        self.name
+    }
+    fn format_value(&self) -> String {
+        self.value.load(Ordering::SeqCst).to_string()
+    }
+}
+
+pub struct Counter<T: Default + ToString + Copy + 'static = usize> {
+    name: &'static str,
+    inner: Once<Arc<CounterImpl<T>, Meta>>,
+    group: *const CounterGroup,
+}
+
+unsafe impl<T: Default + ToString + Copy + 'static> Send for Counter<T> {}
+unsafe impl<T: Default + ToString + Copy + 'static> Sync for Counter<T> {}
+
+impl<T: Default + ToString + Copy + 'static> Counter<T> {
+    pub const fn new(name: &'static str) -> Self {
+        Self {
+            name,
+            inner: Once::new(),
+            group: &DEFAULT_COUNTER_GROUP,
+        }
+    }
+
+    pub const fn new_grouped(name: &'static str, group: &'static CounterGroup) -> Self {
+        Self {
+            name,
+            inner: Once::new(),
+            group,
+        }
+    }
+
+    fn inner(&self) -> &Arc<CounterImpl<T>, Meta> {
+        self.inner.call_once(|| {
+            let c: Arc<CounterImpl<T>, Meta> = Arc::new_in(CounterImpl::new(self.name), Meta);
+            unsafe { &*self.group }.add_counter(c.clone());
+            c
+        })
+    }
+
+    pub fn get(&self) -> T {
+        assert!(cfg!(feature = "stat"));
+        self.inner().value.load(Ordering::SeqCst)
+    }
+
+    pub fn set(&self, value: T) {
+        assert!(cfg!(feature = "stat"));
+        self.inner().value.store(value, Ordering::SeqCst);
+    }
+
+    pub fn swap(&self, value: T) -> T {
+        assert!(cfg!(feature = "stat"));
+        self.inner().value.swap(value, Ordering::SeqCst)
+    }
+
+    pub fn fetch_update<F>(&self, f: impl FnMut(T) -> Option<T>) -> Result<T, T> {
+        assert!(cfg!(feature = "stat"));
+        self.inner()
+            .value
+            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, f)
+    }
+}
+
+macro_rules! impl_inc_dec {
+    ($t: ty) => {
+        impl Counter<$t> {
+            pub fn inc(&self, delta: $t) {
+                assert!(cfg!(feature = "stat"));
+                self.inner().value.fetch_add(delta, Ordering::SeqCst);
+            }
+
+            pub fn dec(&self, delta: $t) {
+                assert!(cfg!(feature = "stat"));
+                self.inner().value.fetch_sub(delta, Ordering::SeqCst);
+            }
+        }
+    };
 }
 
-static TOTAL_ALLOCATIONS: Lazy<Arc<Counter, Meta>> = define_counter::<"total-allocations">();
-static LARGE_ALLOCATIONS: Lazy<Arc<Counter, Meta>> = define_counter::<"large-allocations">();
-static TOTAL_DEALLOCATIONS: Lazy<Arc<Counter, Meta>> = define_counter::<"total-deallocations">();
-static LARGE_DEALLOCATIONS: Lazy<Arc<Counter, Meta>> = define_counter::<"large-deallocations">();
-
-static ALIGNMENTS: [Counter; 11] = [
-    Counter::new(""), // 1
-    Counter::new(""), // 2
-    Counter::new(""), // 4
-    Counter::new(""), // 8
-    Counter::new(""), // 16
-    Counter::new(""), // 32
-    Counter::new(""), // 64
-    Counter::new(""), // 128
-    Counter::new(""), // 256
-    Counter::new(""), // 512
-    Counter::new(""), // 1024
-];
-static OTHER_ALIGNMENT: Counter = Counter::new("");
-
-static SIZES: [Counter; 22] = [
-    Counter::new(""), // 1B
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""), // 1K
-    Counter::new(""),
-    Counter::new(""), // 4K
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""),
-    Counter::new(""), // 1M
-    Counter::new(""), // 2M
-];
-static OTHER_SIZE: Counter = Counter::new("");
+impl_inc_dec!(usize);
+impl_inc_dec!(u128);
+impl_inc_dec!(u64);
+impl_inc_dec!(u32);
+impl_inc_dec!(u16);
+impl_inc_dec!(u8);
+
+impl_inc_dec!(isize);
+impl_inc_dec!(i128);
+impl_inc_dec!(i64);
+impl_inc_dec!(i32);
+impl_inc_dec!(i16);
+impl_inc_dec!(i8);
 
+// impl_inc_dec!(f32);
+// impl_inc_dec!(f64);
+// impl_inc_dec!(u32);
+// impl_inc_dec!(u64);
+
+pub static ALLOC_COUNTERS: CounterGroup = CounterGroup::new("alloc").with_report_fn(|| {
+    eprintln!("alloc:");
+    eprintln!(" total-allocations: {}", TOTAL_ALLOCATIONS.get());
+    eprintln!(" large-allocations: {}", LARGE_ALLOCATIONS.get());
+    eprintln!(" total-deallocations: {}", TOTAL_DEALLOCATIONS.get());
+    eprintln!(" large-deallocations: {}", LARGE_DEALLOCATIONS.get());
+    eprintln!("alignment:");
+    for (i, c) in ALIGNMENTS.iter().enumerate().take(ALIGNMENTS.len() - 1) {
+        eprintln!(" - {} = {}", i, c.get());
+    }
+    eprintln!(" - others = {}", ALIGNMENTS[ALIGNMENTS.len() - 1].get());
+    eprintln!("size:");
+    for (i, c) in SIZES.iter().enumerate().take(SIZES.len() - 1) {
+        eprintln!(" - {} = {}", i, c.get());
+    }
+    eprintln!(" - others = {}", SIZES[SIZES.len() - 1].get());
+});
+
+static TOTAL_ALLOCATIONS: Counter = ALLOC_COUNTERS.new_counter("total-allocations");
+static LARGE_ALLOCATIONS: Counter = ALLOC_COUNTERS.new_counter("large-allocations");
+static TOTAL_DEALLOCATIONS: Counter = ALLOC_COUNTERS.new_counter("total-deallocations");
+static LARGE_DEALLOCATIONS: Counter = ALLOC_COUNTERS.new_counter("large-deallocations");
+
+/// Power of two alignments from 1 to 1024, plus others.
+static ALIGNMENTS: [Counter; 12] = [const { ALLOC_COUNTERS.new_counter("align") }; 12];
+
+/// Power of two sizes from 1b to 2M, plus others.
+static SIZES: [Counter; 23] = [const { ALLOC_COUNTERS.new_counter("size") }; 23];
+
+#[inline(always)]
 pub fn run(block: impl Fn()) {
     if cfg!(not(feature = "stat")) {
         return;
     }
     block()
 }
 
+#[inline(always)]
 pub fn track_allocation(layout: Layout, is_large: bool) {
     run(|| {
-        let i = layout.align().trailing_zeros() as usize;
-        if i < ALIGNMENTS.len() {
-            ALIGNMENTS[i].inc(1);
-        } else {
-            OTHER_ALIGNMENT.inc(1);
+        let mut i = layout.align().trailing_zeros() as usize;
+        if i >= ALIGNMENTS.len() {
+            i = ALIGNMENTS.len() - 1;
         }
-        let i = layout.size().next_power_of_two().trailing_zeros() as usize;
-        if i < SIZES.len() {
-            SIZES[i].inc(1);
-        } else {
-            OTHER_SIZE.inc(1);
+        ALIGNMENTS[i].inc(1);
+        let mut i = layout.size().next_power_of_two().trailing_zeros() as usize;
+        if i >= SIZES.len() {
+            i = SIZES.len() - 1;
         }
+        SIZES[i].inc(1);
         if is_large {
             LARGE_ALLOCATIONS.inc(1);
         }
         TOTAL_ALLOCATIONS.inc(1);
     })
 }
 
+#[inline(always)]
 pub fn track_deallocation(is_large: bool) {
     run(|| {
         if is_large {
@@ -108,40 +257,11 @@ pub fn track_deallocation(is_large: bool) {
     })
 }
 
-pub struct Counter(#[allow(unused)] &'static str, AtomicUsize);
-
-impl Counter {
-    pub const fn new(name: &'static str) -> Self {
-        Self(name, AtomicUsize::new(0))
-    }
-    pub fn get(&self) -> usize {
-        assert!(cfg!(feature = "stat"));
-        self.1.load(Ordering::SeqCst)
-    }
-    pub fn inc(&self, delta: usize) {
-        assert!(cfg!(feature = "stat"));
-        self.1.fetch_add(delta, Ordering::SeqCst);
-    }
-}
-
-#[cfg(not(feature = "stat"))]
-pub(crate) fn report() {}
-
-#[cfg(feature = "stat")]
 pub(crate) fn report() {
-    eprintln!("alignment:");
-    for i in 0..ALIGNMENTS.len() {
-        eprintln!(" - {} = {}", i, ALIGNMENTS[i].get());
-    }
-    eprintln!(" - others = {}", OTHER_ALIGNMENT.get());
-    eprintln!("");
-    eprintln!("size:");
-    for i in 0..SIZES.len() {
-        eprintln!(" - {} = {}", i, SIZES[i].get());
+    if cfg!(not(feature = "stat")) {
+        return;
     }
-    eprintln!(" - others = {}", OTHER_SIZE.get());
-    eprintln!("");
-    while let Some(c) = COUNTERS.lock().pop() {
-        eprintln!("{}: {}", c.0, c.get());
+    for group in ALL_GROUPS.lock().iter() {
+        group.report();
     }
 }