tor_chanmgr/util/
defer.rs

1//! Defer a closure until later.
2
3/// Runs a closure when dropped.
4pub(crate) struct Defer<T, F: FnOnce(T)>(Option<DeferInner<T, F>>);
5
6/// Everything contained by a [`Defer`].
7struct DeferInner<T, F: FnOnce(T)> {
8    /// The argument `f` should be called with when [`Defer`] is dropped.
9    arg: T,
10    /// The function to call.
11    f: F,
12}
13
14impl<T, F: FnOnce(T)> Defer<T, F> {
15    /// Defer running the provided closure `f` with `arg` until the returned [`Defer`] is dropped.
16    #[must_use]
17    pub(crate) fn new(arg: T, f: F) -> Self {
18        Self(Some(DeferInner { arg, f }))
19    }
20
21    /// Return the provided `T` and drop the provided closure without running it.
22    pub(crate) fn cancel(mut self) -> T {
23        // other than the drop handler, there are no other places that mutate the `Option`, so it
24        // should always be `Some` here
25        self.0.take().expect("`Defer` is missing a value").arg
26    }
27}
28
29impl<T, F: FnOnce(T)> std::ops::Drop for Defer<T, F> {
30    fn drop(&mut self) {
31        if let Some(DeferInner { arg, f }) = self.0.take() {
32            f(arg);
33        }
34    }
35}
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40
41    use std::sync::atomic::{AtomicU32, Ordering};
42
43    #[test]
44    fn test_drop() {
45        let x = AtomicU32::new(0);
46        {
47            let _defer = Defer::new(5, |inc| {
48                x.fetch_add(inc, Ordering::Relaxed);
49            });
50            assert_eq!(x.load(Ordering::Relaxed), 0);
51        }
52        assert_eq!(x.load(Ordering::Relaxed), 5);
53    }
54
55    #[test]
56    fn test_cancel() {
57        let x = AtomicU32::new(0);
58        {
59            let defer = Defer::new(5, |inc| {
60                x.fetch_add(inc, Ordering::Relaxed);
61            });
62            assert_eq!(defer.cancel(), 5);
63            assert_eq!(x.load(Ordering::Relaxed), 0);
64        }
65        assert_eq!(x.load(Ordering::Relaxed), 0);
66    }
67
68    #[test]
69    #[should_panic]
70    fn test_panic() {
71        let _ = Defer::new((), |()| {
72            panic!("intentional panic");
73        });
74    }
75}