From e2b953063d1c6302dd61a2284f8da8b770ddc11d Mon Sep 17 00:00:00 2001
From: Maybe Waffle <waffle.lapkin@gmail.com>
Date: Tue, 23 May 2023 14:53:36 +0000
Subject: [PATCH] Don't leak the function that is called on drop

---
 compiler/rustc_data_structures/src/lib.rs     | 22 ++++++++++++-------
 .../src/owned_slice/tests.rs                  |  4 ++--
 compiler/rustc_interface/src/interface.rs     |  4 ++--
 compiler/rustc_middle/src/ty/context/tls.rs   |  2 +-
 compiler/rustc_query_system/src/query/job.rs  |  4 ++--
 5 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/compiler/rustc_data_structures/src/lib.rs b/compiler/rustc_data_structures/src/lib.rs
index 5b9b0e106d254..859e384d8b529 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 1eb5378cd1ad8..520871a12be99 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, |_| &[]);
 
diff --git a/compiler/rustc_interface/src/interface.rs b/compiler/rustc_interface/src/interface.rs
index 681819703c298..39d56897999d9 100644
--- a/compiler/rustc_interface/src/interface.rs
+++ b/compiler/rustc_interface/src/interface.rs
@@ -3,9 +3,9 @@ use crate::util;
 use rustc_ast::token;
 use rustc_ast::{self as ast, LitKind, MetaItemKind};
 use rustc_codegen_ssa::traits::CodegenBackend;
+use rustc_data_structures::defer;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::sync::Lrc;
-use rustc_data_structures::OnDrop;
 use rustc_errors::registry::Registry;
 use rustc_errors::{ErrorGuaranteed, Handler};
 use rustc_lint::LintStore;
@@ -325,7 +325,7 @@ pub fn run_compiler<R: Send>(config: Config, f: impl FnOnce(&Compiler) -> R + Se
 
             rustc_span::set_source_map(compiler.sess.parse_sess.clone_source_map(), move || {
                 let r = {
-                    let _sess_abort_error = OnDrop(|| {
+                    let _sess_abort_error = defer(|| {
                         compiler.sess.finish_diagnostics(registry);
                     });
 
diff --git a/compiler/rustc_middle/src/ty/context/tls.rs b/compiler/rustc_middle/src/ty/context/tls.rs
index 5c1c419811e38..9de77b9fda11a 100644
--- a/compiler/rustc_middle/src/ty/context/tls.rs
+++ b/compiler/rustc_middle/src/ty/context/tls.rs
@@ -78,7 +78,7 @@ where
 {
     TLV.with(|tlv| {
         let old = tlv.replace(erase(context));
-        let _reset = rustc_data_structures::OnDrop(move || tlv.set(old));
+        let _reset = rustc_data_structures::defer(move || tlv.set(old));
         f()
     })
 }
diff --git a/compiler/rustc_query_system/src/query/job.rs b/compiler/rustc_query_system/src/query/job.rs
index 5f2ec656d1d1e..f45f7ca5da6dd 100644
--- a/compiler/rustc_query_system/src/query/job.rs
+++ b/compiler/rustc_query_system/src/query/job.rs
@@ -22,7 +22,7 @@ use {
     rustc_data_structures::fx::FxHashSet,
     rustc_data_structures::sync::Lock,
     rustc_data_structures::sync::Lrc,
-    rustc_data_structures::{jobserver, OnDrop},
+    rustc_data_structures::{defer, jobserver},
     rustc_span::DUMMY_SP,
     std::iter,
     std::process,
@@ -530,7 +530,7 @@ fn remove_cycle<D: DepKind>(
 /// all active queries for cycles before finally resuming all the waiters at once.
 #[cfg(parallel_compiler)]
 pub fn deadlock<D: DepKind>(query_map: QueryMap<D>, registry: &rayon_core::Registry) {
-    let on_panic = OnDrop(|| {
+    let on_panic = defer(|| {
         eprintln!("deadlock handler panicked, aborting process");
         process::abort();
     });