Ensure consistent drop for panicking drop in hint::select_unpredictable

library/core/src/hint.rs

@@ -786,12 +786,42 @@ pub fn select_unpredictable<T>(condition: bool, true_val: T, false_val: T) -> T
     // Change this to use ManuallyDrop instead.
     let mut true_val = MaybeUninit::new(true_val);
     let mut false_val = MaybeUninit::new(false_val);
+
+    struct DropOnPanic<T> {
+        inner: *mut MaybeUninit<T>,
+    }
+
+    impl<T> Drop for DropOnPanic<T> {
+        fn drop(&mut self) {
+            // SAFETY: Must be guaranteed on construction of local type `DropOnPanic`.
+            unsafe { (*self.inner).assume_init_drop() }
+        }
+    }
+
+    let true_ptr = (&mut true_val) as *mut _;
+    let false_ptr = (&mut false_val) as *mut _;
+
     // 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();
+        // Extract the selected value first, ensure it is dropped as well if dropping the unselected
+        // value panics.
+        let (guard, drop) = crate::intrinsics::select_unpredictable(
+            condition,
+            (true_ptr, false_ptr),
+            (false_ptr, true_ptr),
+        );
+
+        // SAFETY: both pointers are to valid `MaybeUninit`, in both variants they do not alias but
+        // the two arguments we have selected from did alias each other.
+        let guard = DropOnPanic { inner: guard };
+        (*drop).assume_init_drop();
+        crate::mem::forget(guard);
+
+        // Note that it is important to use the values here. Reading from the pointer we got makes
+        // LLVM forget the !unpredictable annotation sometimes (in tests, integer sized values in
+        // particular seemed to confuse it, also observed in llvm/llvm-project #82340).
         crate::intrinsics::select_unpredictable(condition, true_val, false_val).assume_init()
     }
 }

library/coretests/tests/hint.rs

@@ -21,3 +21,40 @@ fn select_unpredictable_drop() {
     assert!(a_dropped.get());
     assert!(b_dropped.get());
 }
+
+#[test]
+#[should_panic]
+fn select_unpredictable_drop_on_panic() {
+    use core::cell::Cell;
+
+    struct X<'a> {
+        cell: &'a Cell<u16>,
+        expect: u16,
+        write: u16,
+    }
+
+    impl Drop for X<'_> {
+        fn drop(&mut self) {
+            let value = self.cell.get();
+            self.cell.set(self.write);
+            assert_eq!(value, self.expect);
+        }
+    }
+
+    let cell = Cell::new(0);
+
+    // Trigger a double-panic if the selected cell was not dropped during panic.
+    let _armed = X { cell: &cell, expect: 0xdead, write: 0 };
+    let selected = X { cell: &cell, write: 0xdead, expect: 1 };
+    let unselected = X { cell: &cell, write: 1, expect: 0xff };
+
+    // The correct drop order is:
+    //
+    // 1. `unselected` drops, writes 1, and panics as 0 != 0xff
+    // 2. `selected` drops during unwind, writes 0xdead and does not panic as 1 == 1
+    // 3. `armed` drops during unwind, writes 0 and does not panic as 0xdead == 0xdead
+    //
+    // If `selected` is not dropped, `armed` panics as 1 != 0xdead
+    let _unreachable =
+        core::hint::select_unpredictable(core::hint::black_box(true), selected, unselected);
+}