solar_data_structures/
drop_guard.rs

1use std::mem::ManuallyDrop;
2
3/// Returns a structure that calls `f` when dropped.
4#[inline]
5pub const fn defer<F: FnOnce()>(f: F) -> DropGuard<(), impl FnOnce(())> {
6    DropGuard::new((), move |()| f())
7}
8
9/// Runs `F` on `T` when the instance is dropped.
10///
11/// Equivalent of `std::mem::DropGuard`.
12#[must_use]
13pub struct DropGuard<T, F: FnOnce(T)> {
14    inner: ManuallyDrop<T>,
15    f: ManuallyDrop<F>,
16}
17
18impl<T, F: FnOnce(T)> DropGuard<T, F> {
19    /// Creates a new `OnDrop` instance.
20    #[inline]
21    pub const fn new(value: T, f: F) -> Self {
22        Self { inner: ManuallyDrop::new(value), f: ManuallyDrop::new(f) }
23    }
24
25    /// Consumes the `DropGuard`, returning the wrapped value.
26    ///
27    /// This will not execute the closure. This is implemented as an associated
28    /// function to prevent any potential conflicts with any other methods called
29    /// `into_inner` from the `Deref` and `DerefMut` impls.
30    ///
31    /// It is typically preferred to call this function instead of `mem::forget`
32    /// because it will return the stored value and drop variables captured
33    /// by the closure instead of leaking their owned resources.
34    #[inline]
35    #[must_use]
36    pub fn into_inner(guard: Self) -> T {
37        // First we ensure that dropping the guard will not trigger
38        // its destructor
39        let mut guard = ManuallyDrop::new(guard);
40
41        // Next we manually read the stored value from the guard.
42        //
43        // SAFETY: this is safe because we've taken ownership of the guard.
44        let value = unsafe { ManuallyDrop::take(&mut guard.inner) };
45
46        // Finally we drop the stored closure. We do this *after* having read
47        // the value, so that even if the closure's `drop` function panics,
48        // unwinding still tries to drop the value.
49        //
50        // SAFETY: this is safe because we've taken ownership of the guard.
51        unsafe { ManuallyDrop::drop(&mut guard.f) };
52        value
53    }
54}
55
56impl<T, F: FnOnce(T)> std::ops::Deref for DropGuard<T, F> {
57    type Target = T;
58
59    #[inline]
60    fn deref(&self) -> &Self::Target {
61        &self.inner
62    }
63}
64
65impl<T, F: FnOnce(T)> std::ops::DerefMut for DropGuard<T, F> {
66    #[inline]
67    fn deref_mut(&mut self) -> &mut Self::Target {
68        &mut self.inner
69    }
70}
71
72impl<T, F: FnOnce(T)> Drop for DropGuard<T, F> {
73    #[inline]
74    fn drop(&mut self) {
75        // SAFETY: `DropGuard` is in the process of being dropped.
76        let inner = unsafe { ManuallyDrop::take(&mut self.inner) };
77
78        // SAFETY: `DropGuard` is in the process of being dropped.
79        let f = unsafe { ManuallyDrop::take(&mut self.f) };
80
81        f(inner);
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    struct DropCheck<'a>(&'a mut usize);
90    impl Drop for DropCheck<'_> {
91        fn drop(&mut self) {
92            *self.0 += 1;
93        }
94    }
95
96    #[test]
97    fn test_defer() {
98        let mut x = 0usize;
99        let defer = defer(|| x += 1);
100        drop(defer);
101        assert_eq!(x, 1);
102    }
103
104    #[test]
105    fn drop_guard() {
106        let mut x = 0usize;
107        let guard = DropGuard::new(&mut x, |x| *x += 1);
108        assert_eq!(**guard, 0);
109        drop(guard);
110        assert_eq!(x, 1);
111    }
112
113    #[test]
114    fn normal() {
115        let mut dropped = 0;
116        let mut closure_called = 0;
117        let mut closure_dropped = 0;
118        let closure_drop_check = DropCheck(&mut closure_dropped);
119        let guard = DropGuard::new(DropCheck(&mut dropped), |_s| {
120            drop(closure_drop_check);
121            closure_called += 1;
122        });
123        drop(guard);
124        assert_eq!(dropped, 1);
125        assert_eq!(closure_called, 1);
126        assert_eq!(closure_dropped, 1);
127    }
128
129    #[test]
130    fn disable() {
131        let mut dropped = 0;
132        let mut closure_called = 0;
133        let mut closure_dropped = 0;
134        let closure_drop_check = DropCheck(&mut closure_dropped);
135        let guard = DropGuard::new(DropCheck(&mut dropped), |_s| {
136            drop(closure_drop_check);
137            closure_called += 1;
138        });
139        let value = DropGuard::into_inner(guard);
140        assert_eq!(*value.0, 0);
141        assert_eq!(closure_called, 0);
142        assert_eq!(closure_dropped, 1);
143
144        drop(value);
145        assert_eq!(dropped, 1);
146    }
147}
solar_data_structures/drop_guard.rs

solar_data_structures/
drop_guard.rs
