use std::cell::{Cell, RefCell};
use std::collections::BTreeMap;
thread_local! {
static CURRENT_PASS: Cell<&'static str> = const { Cell::new(UNATTRIBUTED) };
static STATS: RefCell<BTreeMap<(&'static str, &'static str), u64>> =
const { RefCell::new(BTreeMap::new()) };
}
pub const UNATTRIBUTED: &str = "?";
pub fn current_pass() -> &'static str {
CURRENT_PASS.get()
}
pub struct PassScope {
prev: &'static str,
}
pub fn enter(name: &'static str) -> PassScope {
PassScope {
prev: CURRENT_PASS.replace(name),
}
}
impl Drop for PassScope {
fn drop(&mut self) {
CURRENT_PASS.set(self.prev);
}
}
pub fn record_stat(key: &'static str, n: u64) {
STATS.with_borrow_mut(|stats| {
*stats.entry((current_pass(), key)).or_insert(0) += n;
});
}
pub fn drain_stats() -> Vec<((&'static str, &'static str), u64)> {
STATS.with_borrow_mut(std::mem::take).into_iter().collect()
}
pub fn absorb_stats(entries: impl IntoIterator<Item = ((&'static str, &'static str), u64)>) {
STATS.with_borrow_mut(|stats| {
for (key, n) in entries {
*stats.entry(key).or_insert(0) += n;
}
});
}
#[macro_export]
macro_rules! pass_log {
($lvl:ident, $($arg:tt)+) => {{
let __pass = $crate::pass_scope::current_pass();
$crate::__log::$lvl!(target: __pass, "[{}] {}", __pass, format_args!($($arg)+));
}};
}
#[macro_export]
macro_rules! stat {
($key:expr) => {
$crate::pass_scope::record_stat($key, 1)
};
($key:expr, $n:expr) => {
$crate::pass_scope::record_stat($key, $n)
};
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn scopes_nest_and_restore() {
assert_eq!(current_pass(), UNATTRIBUTED);
{
let _outer = enter("outer");
assert_eq!(current_pass(), "outer");
{
let _inner = enter("inner");
assert_eq!(current_pass(), "inner");
}
assert_eq!(current_pass(), "outer");
}
assert_eq!(current_pass(), UNATTRIBUTED);
}
#[test]
fn stats_accumulate_per_pass_and_drain() {
let _ = drain_stats();
{
let _s = enter("p1");
record_stat("k", 2);
record_stat("k", 3);
}
{
let _s = enter("p2");
record_stat("k", 1);
}
let drained = drain_stats();
assert_eq!(drained, vec![(("p1", "k"), 5), (("p2", "k"), 1)]);
assert!(drain_stats().is_empty());
}
}