solar_data_structures/
drop_guard.rs

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