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