use alloc::vec::Vec;
use std::{
cell::{Cell, RefCell},
thread_local,
};
#[derive(Debug)]
struct TagContext {
namespace: &'static str,
ids: Vec<usize>,
next: usize,
}
thread_local! {
static TAGGING_ENABLED: Cell<bool> = const { Cell::new(false) };
static TAG_STACK: RefCell<Vec<TagContext>> = const { RefCell::new(Vec::new()) };
}
pub fn is_enabled() -> bool {
TAGGING_ENABLED.with(|flag| flag.get())
}
#[cfg(test)]
pub fn set_enabled(enabled: bool) {
TAGGING_ENABLED.with(|flag| flag.set(enabled));
}
pub fn with_tag<R>(ids: Vec<usize>, namespace: &'static str, f: impl FnOnce() -> R) -> R {
if ids.is_empty() {
panic!("tagged block '{namespace}' must include at least one id");
}
TAG_STACK.with(|stack| {
let mut stack = stack.borrow_mut();
if !stack.is_empty() {
panic!("nested tagged blocks are not allowed");
}
stack.push(TagContext { namespace, ids, next: 0 });
});
let result = f();
TAG_STACK.with(|stack| {
let mut stack = stack.borrow_mut();
let ctx = stack.pop().expect("tag stack underflow");
if ctx.next != ctx.ids.len() {
panic!(
"tagged block '{}' expected {} asserts, saw {}",
ctx.namespace,
ctx.ids.len(),
ctx.next
);
}
});
result
}
#[cfg(test)]
pub fn consume_tag() -> (usize, &'static str) {
TAG_STACK.with(|stack| {
let mut stack = stack.borrow_mut();
let ctx = stack.last_mut().expect("assertion made without an active tagged block");
if ctx.next >= ctx.ids.len() {
panic!(
"tagged block '{}' exceeded expected asserts ({})",
ctx.namespace,
ctx.ids.len()
);
}
let id = ctx.ids[ctx.next];
ctx.next += 1;
(id, ctx.namespace)
})
}