solar_data_structures/
on_drop.rs

1use std::mem::ManuallyDrop;
2
3/// Returns a structure that calls `f` when dropped.
4#[inline(always)]
5pub fn defer<F: FnOnce()>(f: F) -> OnDrop<(), impl FnOnce(())> {
6    OnDrop::new((), move |()| f())
7}
8
9/// Runs `F` on `T` when the instance is dropped.
10pub struct OnDrop<T, F: FnOnce(T)>(pub ManuallyDrop<(T, Option<F>)>);
11
12impl<T, F: FnOnce(T)> OnDrop<T, F> {
13    /// Creates a new `OnDrop` instance.
14    #[inline(always)]
15    pub fn new(value: T, f: F) -> Self {
16        Self(ManuallyDrop::new((value, Some(f))))
17    }
18
19    /// Returns a reference to the inner value.
20    #[inline(always)]
21    pub fn inner(&self) -> &T {
22        &self.0.0
23    }
24
25    /// Returns a mutable reference to the inner value.
26    #[inline(always)]
27    pub fn inner_mut(&mut self) -> &mut T {
28        &mut self.0.0
29    }
30
31    /// Consumes the instance and returns the inner value.
32    #[inline(always)]
33    pub fn into_inner(mut self) -> T {
34        unsafe {
35            std::ptr::drop_in_place(&mut self.0.1);
36            std::ptr::read(&self.0.0)
37        }
38    }
39
40    /// Consumes the instance without running `F` on the inner value.
41    #[inline(always)]
42    pub fn disable(self) {
43        let _ = self.into_inner();
44    }
45}
46
47impl<T, F: FnOnce(T)> std::ops::Deref for OnDrop<T, F> {
48    type Target = T;
49
50    #[inline(always)]
51    fn deref(&self) -> &Self::Target {
52        &self.0.0
53    }
54}
55
56impl<T, F: FnOnce(T)> std::ops::DerefMut for OnDrop<T, F> {
57    #[inline(always)]
58    fn deref_mut(&mut self) -> &mut Self::Target {
59        &mut self.0.0
60    }
61}
62
63impl<T, F: FnOnce(T)> Drop for OnDrop<T, F> {
64    #[inline(always)]
65    fn drop(&mut self) {
66        unsafe {
67            if let Some(f) = self.0.1.take() {
68                f(std::ptr::read(&self.0.0));
69            } else {
70                ManuallyDrop::drop(&mut self.0);
71            }
72        }
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn test_defer() {
82        let mut x = 0usize;
83        let defer = defer(|| x += 1);
84        drop(defer);
85        assert_eq!(x, 1);
86    }
87
88    #[test]
89    fn test_on_drop() {
90        let mut x = 0usize;
91        let on_drop = OnDrop::new(&mut x, |x| *x += 1);
92        assert_eq!(**on_drop.inner(), 0);
93        drop(on_drop);
94        assert_eq!(x, 1);
95    }
96}