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
#![no_std]

use core::mem::ManuallyDrop;

/// Defer execution of a closure until the return value is dropped.
pub fn defer<F>(f: F) -> impl Drop
where
    F: FnOnce(),
{
    struct Defer<F: FnOnce()>(ManuallyDrop<F>);

    impl<F: FnOnce()> Drop for Defer<F> {
        fn drop(&mut self) {
            let f: F = unsafe { ManuallyDrop::take(&mut self.0) };
            f();
        }
    }

    Defer(ManuallyDrop::new(f))
}

/// Defer execution of a closure until the current scope end.
#[macro_export]
macro_rules! defer {
    ($e:expr) => {
        let _defer = $crate::defer(|| $e);
    };
}

#[test]
fn test() {
    use core::cell::RefCell;

    let i = RefCell::new(0);

    {
        let _d = defer(|| *i.borrow_mut() += 1);
        assert_eq!(*i.borrow(), 0);
    }

    assert_eq!(*i.borrow(), 1);
}

#[test]
fn test_macro() {
    use core::cell::RefCell;

    let i = RefCell::new(0);
    let k = RefCell::new(0);

    {
        defer!(*i.borrow_mut() += 1);
        defer!(*k.borrow_mut() += 1);
        assert_eq!(*i.borrow(), 0);
        assert_eq!(*k.borrow(), 0);
    }

    assert_eq!(*i.borrow(), 1);
    assert_eq!(*k.borrow(), 1);
}