Don't leak the function that is called on drop

2 files changed, 16 insertions, 10 deletions

diff --git a/compiler/rustc_data_structures/src/lib.rs b/compiler/rustc_data_structures/src/lib.rs
index 5b9b0e106d2..859e384d8b5 100644
--- a/compiler/rustc_data_structures/src/lib.rs
+++ b/compiler/rustc_data_structures/src/lib.rs
@@ -102,21 +102,27 @@ pub mod unord;
 pub use ena::undo_log;
 pub use ena::unify;
 
-pub struct OnDrop<F: Fn()>(pub F);
+/// Returns a structure that calls `f` when dropped.
+pub fn defer<F: FnOnce()>(f: F) -> OnDrop<F> {
+    OnDrop(Some(f))
+}
+
+pub struct OnDrop<F: FnOnce()>(Option<F>);
 
-impl<F: Fn()> OnDrop<F> {
-    /// Forgets the function which prevents it from running.
-    /// Ensure that the function owns no memory, otherwise it will be leaked.
+impl<F: FnOnce()> OnDrop<F> {
+    /// Disables on-drop call.
     #[inline]
-    pub fn disable(self) {
-        std::mem::forget(self);
+    pub fn disable(mut self) {
+        self.0.take();
     }
 }
 
-impl<F: Fn()> Drop for OnDrop<F> {
+impl<F: FnOnce()> Drop for OnDrop<F> {
     #[inline]
     fn drop(&mut self) {
-        (self.0)();
+        if let Some(f) = self.0.take() {
+            f();
+        }
     }
 }
 
diff --git a/compiler/rustc_data_structures/src/owned_slice/tests.rs b/compiler/rustc_data_structures/src/owned_slice/tests.rs
index 1eb5378cd1a..520871a12be 100644
--- a/compiler/rustc_data_structures/src/owned_slice/tests.rs
+++ b/compiler/rustc_data_structures/src/owned_slice/tests.rs
@@ -7,8 +7,8 @@ use std::{
 };
 
 use crate::{
+    defer,
     owned_slice::{slice_owned, try_slice_owned, OwnedSlice},
-    OnDrop,
 };
 
 #[test]
@@ -66,7 +66,7 @@ fn boxed() {
 fn drop_drops() {
     let flag = Arc::new(AtomicBool::new(false));
     let flag_prime = Arc::clone(&flag);
-    let d = OnDrop(move || flag_prime.store(true, atomic::Ordering::Relaxed));
+    let d = defer(move || flag_prime.store(true, atomic::Ordering::Relaxed));
 
     let slice = slice_owned(d, |_| &[]);