Skip to main content

runmat_runtime/
output_context.rs

1use runmat_thread_local::runmat_thread_local;
2use std::cell::RefCell;
3
4runmat_thread_local! {
5    static REQUESTED_OUTPUTS: RefCell<Option<usize>> = const { RefCell::new(None) };
6}
7
8pub struct OutputCountGuard {
9    prev: Option<usize>,
10}
11
12impl Drop for OutputCountGuard {
13    fn drop(&mut self) {
14        REQUESTED_OUTPUTS.with(|cell| {
15            *cell.borrow_mut() = self.prev;
16        });
17    }
18}
19
20pub fn push_output_count(count: usize) -> OutputCountGuard {
21    let prev = REQUESTED_OUTPUTS.with(|cell| {
22        let mut guard = cell.borrow_mut();
23        let prev = guard.take();
24        *guard = Some(count);
25        prev
26    });
27    OutputCountGuard { prev }
28}
29
30pub fn requested_output_count() -> Option<usize> {
31    REQUESTED_OUTPUTS.with(|cell| *cell.borrow())
32}