vor 0.2.0

//! Runtime-named profiling scopes: the dynamic counterpart to the
//! compile-time [`profile_scope!`](crate::profile_scope) macro.
//!
//! [`scope`] opens a puffin scope whose name is chosen at run time, for
//! callers that can't use the attribute macros - notably the Python
//! bindings, where a `@vor.profile` decorator wraps each call in
//! `scope(qualname)` to mirror `#[vor::profile]`.

use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};

use puffin::{ScopeId, ThreadProfiler};

// One ScopeId per distinct name, registered once and reused: puffin
// mints a fresh id and stores the name's details on each registration,
// so registering per call would leak ids and re-store the name.
static SCOPE_IDS: OnceLock<Mutex<HashMap<String, ScopeId>>> = OnceLock::new();

/// Open guard for a runtime-named scope; the scope closes when it drops.
#[must_use = "bind the guard so the scope stays open for the block"]
pub struct ScopeGuard(Option<usize>);

/// Open a profiling scope named at run time, the dynamic counterpart to
/// [`profile_scope!`](crate::profile_scope). The returned guard closes
/// the scope on drop, and the scope shows in the flame chart under
/// `name`.
///
/// A no-op while scope collection is off, so it costs nothing until
/// [`enable`](crate::enable) is called. Begin and end run on the
/// calling thread, so the guard must drop on the thread that opened it.
pub fn scope(name: &str) -> ScopeGuard {
    if !puffin::are_scopes_on() {
        return ScopeGuard(None);
    }
    let ids = SCOPE_IDS.get_or_init(|| Mutex::new(HashMap::new()));
    let mut map = ids.lock().unwrap();
    let id = *map
        .entry(name.to_owned())
        .or_insert_with(|| ThreadProfiler::call(|tp| tp.register_function_scope(name.to_owned(), "", 0)));
    drop(map);
    let offset = ThreadProfiler::call(|tp| tp.begin_scope(id, ""));
    ScopeGuard(Some(offset))
}

impl Drop for ScopeGuard {
    fn drop(&mut self) {
        if let Some(offset) = self.0 {
            ThreadProfiler::call(|tp| tp.end_scope(offset));
        }
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn dynamic_scope_lands_in_the_flame_collection() {
        crate::enable();
        let view = puffin::GlobalFrameView::default();
        {
            let _s = super::scope("py::train_step");
            std::thread::sleep(std::time::Duration::from_micros(50));
        }
        crate::frame_mark();

        let view = view.lock();
        let found = view
            .scope_collection()
            .scopes_by_id()
            .values()
            .any(|d| d.name().as_ref() == "py::train_step");
        assert!(found, "runtime scope name should reach the flame collection");
    }
}