1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use std::cell::RefCell;
use std::thread::LocalKey;

use crate::util::finalizer::Finalizer;

pub trait Events {
    fn new() -> Self;
    fn take(&mut self) -> Self;
}

pub type Stack<E> = RefCell<Vec<E>>;

pub fn new_stack<E: Events>() -> Stack<E> {
    RefCell::new(Vec::new())
}

pub type Local<E> = LocalKey<Stack<E>>;

pub fn collect<E: Events, R>(local: &'static Local<E>, f: impl FnOnce() -> R) -> (R, E) {
    local.with(move |cell| {
        {
            let events = E::new();
            let mut events_stack = cell.borrow_mut();
            events_stack.push(events);
        }

        // Removes the events even in case of panic
        let finalizer = Finalizer::new(|| {
            local.with(move |cell| {
                let mut events_stack = cell.borrow_mut();
                events_stack.pop();
            })
        });

        // This function may panic
        let result = f();

        let events = {
            let mut events_stack = cell.borrow_mut();
            let events = events_stack.last_mut().unwrap();
            events.take()
        };

        drop(finalizer);

        (result, events)
    })
}

fn enabled_with_cell<E: Events>(cell: &RefCell<Vec<E>>) -> bool {
    !cell.borrow().is_empty()
}

pub fn enabled<E: Events>(local: &'static Local<E>) -> bool {
    local.with(move |cell| enabled_with_cell(cell))
}

pub fn modify<E: Events>(local: &'static Local<E>, f: impl FnOnce(&mut Vec<E>)) {
    local.with(move |cell| {
        if enabled_with_cell(cell) {
            let mut events_stack = cell.borrow_mut();
            f(&mut events_stack);
        }
    });
}

#[cfg(test)]
mod tests {
    use std::panic::catch_unwind;

    use crate::util::events::{self, Events, Stack};

    struct Counter(u32);

    impl Events for Counter {
        fn new() -> Self {
            Counter(0)
        }

        fn take(&mut self) -> Self {
            Counter(self.0)
        }
    }

    thread_local! {
        static LOCAL: Stack<Counter> = events::new_stack();
    }

    fn enabled() -> bool {
        events::enabled(&LOCAL)
    }

    fn collect(f: impl FnOnce()) -> u32 {
        let ((), Counter(n)) = events::collect(&LOCAL, f);
        n
    }

    fn inc() {
        events::modify(&LOCAL, |stack| stack.iter_mut().for_each(|c| c.0 += 1));
    }

    #[test]
    fn collect_and_modify() {
        let n = collect(|| {
            inc();
            inc();

            let m = collect(|| inc());

            assert_eq!(1, m);
        });

        assert_eq!(3, n);
    }

    #[test]
    fn is_only_enabled_during_collection() {
        assert!(!enabled());
        collect(|| {
            assert!(enabled());
        });
        assert!(!enabled());
    }

    #[test]
    fn logger_is_not_enabled_after_panic() {
        let _ = catch_unwind(|| {
            collect(|| {
                panic!();
            })
        });
        assert!(!enabled());
    }
}