1
//! Defer a closure until later.
2

            
3
/// Runs a closure when dropped.
4
pub(crate) struct Defer<T, F: FnOnce(T)>(Option<DeferInner<T, F>>);
5

            
6
/// Everything contained by a [`Defer`].
7
struct 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

            
14
impl<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
36
    pub(crate) fn new(arg: T, f: F) -> Self {
18
36
        Self(Some(DeferInner { arg, f }))
19
36
    }
20

            
21
    /// Return the provided `T` and drop the provided closure without running it.
22
24
    pub(crate) fn cancel(mut self) -> T {
23
24
        // other than the drop handler, there are no other places that mutate the `Option`, so it
24
24
        // should always be `Some` here
25
24
        self.0.take().expect("`Defer` is missing a value").arg
26
24
    }
27
}
28

            
29
impl<T, F: FnOnce(T)> std::ops::Drop for Defer<T, F> {
30
36
    fn drop(&mut self) {
31
36
        if let Some(DeferInner { arg, f }) = self.0.take() {
32
12
            f(arg);
33
24
        }
34
36
    }
35
}
36

            
37
#[cfg(test)]
38
mod 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
}