From e5e5fb9d805dacdd5ceaf73254d61726d75f455a Mon Sep 17 00:00:00 2001
From: Amanieu d'Antras <amanieu@gmail.com>
Date: Thu, 17 Apr 2025 17:29:32 +0200
Subject: [PATCH] Fix drop handling in `hint::select_unpredictable`

This intrinsic doesn't drop the value that is not selected so this is
manually done in the public function that wraps the intrinsic.
---
 library/core/src/hint.rs           | 30 +++++++++++++++++++++---------
 library/core/src/intrinsics/mod.rs |  2 ++
 library/coretests/tests/hint.rs    | 23 +++++++++++++++++++++++
 library/coretests/tests/lib.rs     |  2 ++
 4 files changed, 48 insertions(+), 9 deletions(-)
 create mode 100644 library/coretests/tests/hint.rs

diff --git a/library/core/src/hint.rs b/library/core/src/hint.rs
index f6708cc4bc912..1ca23ab6eea66 100644
--- a/library/core/src/hint.rs
+++ b/library/core/src/hint.rs
@@ -4,6 +4,7 @@
 //!
 //! Hints may be compile time or runtime.
 
+use crate::mem::MaybeUninit;
 use crate::{intrinsics, ub_checks};
 
 /// Informs the compiler that the site which is calling this function is not
@@ -735,9 +736,9 @@ pub const fn cold_path() {
     crate::intrinsics::cold_path()
 }
 
-/// Returns either `true_val` or `false_val` depending on the value of `b`,
-/// with a hint to the compiler that `b` is unlikely to be correctly
-/// predicted by a CPU’s branch predictor.
+/// Returns either `true_val` or `false_val` depending on the value of
+/// `condition`, with a hint to the compiler that `condition` is unlikely to be
+/// correctly predicted by a CPU’s branch predictor.
 ///
 /// This method is functionally equivalent to
 /// ```ignore (this is just for illustrative purposes)
@@ -753,10 +754,10 @@ pub const fn cold_path() {
 /// search.
 ///
 /// Note however that this lowering is not guaranteed (on any platform) and
-/// should not be relied upon when trying to write constant-time code. Also
-/// be aware that this lowering might *decrease* performance if `condition`
-/// is well-predictable. It is advisable to perform benchmarks to tell if
-/// this function is useful.
+/// should not be relied upon when trying to write cryptographic constant-time
+/// code. Also be aware that this lowering might *decrease* performance if
+/// `condition` is well-predictable. It is advisable to perform benchmarks to
+/// tell if this function is useful.
 ///
 /// # Examples
 ///
@@ -780,6 +781,17 @@ pub const fn cold_path() {
 /// ```
 #[inline(always)]
 #[unstable(feature = "select_unpredictable", issue = "133962")]
-pub fn select_unpredictable<T>(b: bool, true_val: T, false_val: T) -> T {
-    crate::intrinsics::select_unpredictable(b, true_val, false_val)
+pub fn select_unpredictable<T>(condition: bool, true_val: T, false_val: T) -> T {
+    // FIXME(https://github.com/rust-lang/unsafe-code-guidelines/issues/245):
+    // Change this to use ManuallyDrop instead.
+    let mut true_val = MaybeUninit::new(true_val);
+    let mut false_val = MaybeUninit::new(false_val);
+    // SAFETY: The value that is not selected is dropped, and the selected one
+    // is returned. This is necessary because the intrinsic doesn't drop the
+    // value that is  not selected.
+    unsafe {
+        crate::intrinsics::select_unpredictable(!condition, &mut true_val, &mut false_val)
+            .assume_init_drop();
+        crate::intrinsics::select_unpredictable(condition, true_val, false_val).assume_init()
+    }
 }
diff --git a/library/core/src/intrinsics/mod.rs b/library/core/src/intrinsics/mod.rs
index e5604d277f57a..a01efb2adebe3 100644
--- a/library/core/src/intrinsics/mod.rs
+++ b/library/core/src/intrinsics/mod.rs
@@ -1327,6 +1327,8 @@ pub const fn unlikely(b: bool) -> bool {
 /// any safety invariants.
 ///
 /// The public form of this instrinsic is [`core::hint::select_unpredictable`].
+/// However unlike the public form, the intrinsic will not drop the value that
+/// is not selected.
 #[unstable(feature = "core_intrinsics", issue = "none")]
 #[rustc_intrinsic]
 #[rustc_nounwind]
diff --git a/library/coretests/tests/hint.rs b/library/coretests/tests/hint.rs
new file mode 100644
index 0000000000000..032bbc1dcc80f
--- /dev/null
+++ b/library/coretests/tests/hint.rs
@@ -0,0 +1,23 @@
+#[test]
+fn select_unpredictable_drop() {
+    use core::cell::Cell;
+    struct X<'a>(&'a Cell<bool>);
+    impl Drop for X<'_> {
+        fn drop(&mut self) {
+            self.0.set(true);
+        }
+    }
+
+    let a_dropped = Cell::new(false);
+    let b_dropped = Cell::new(false);
+    let a = X(&a_dropped);
+    let b = X(&b_dropped);
+    assert!(!a_dropped.get());
+    assert!(!b_dropped.get());
+    let selected = core::hint::select_unpredictable(core::hint::black_box(true), a, b);
+    assert!(!a_dropped.get());
+    assert!(b_dropped.get());
+    drop(selected);
+    assert!(a_dropped.get());
+    assert!(b_dropped.get());
+}
diff --git a/library/coretests/tests/lib.rs b/library/coretests/tests/lib.rs
index 7ad154796f649..1c43bfe0ed4a9 100644
--- a/library/coretests/tests/lib.rs
+++ b/library/coretests/tests/lib.rs
@@ -68,6 +68,7 @@
 #![feature(pointer_is_aligned_to)]
 #![feature(portable_simd)]
 #![feature(ptr_metadata)]
+#![feature(select_unpredictable)]
 #![feature(slice_from_ptr_range)]
 #![feature(slice_internals)]
 #![feature(slice_partition_dedup)]
@@ -147,6 +148,7 @@ mod ffi;
 mod fmt;
 mod future;
 mod hash;
+mod hint;
 mod intrinsics;
 mod io;
 mod iter;