diff --git a/library/core/src/hint.rs b/library/core/src/hint.rs index c72eeb9a9c976..1bb1c56950879 100644 --- a/library/core/src/hint.rs +++ b/library/core/src/hint.rs @@ -786,12 +786,42 @@ pub fn select_unpredictable(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 { + inner: *mut MaybeUninit, + } + + impl Drop for DropOnPanic { + 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() } } diff --git a/library/coretests/tests/hint.rs b/library/coretests/tests/hint.rs index 032bbc1dcc80f..9ef5567681e93 100644 --- a/library/coretests/tests/hint.rs +++ b/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, + 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); +}