Ensure consistent drop for panicking drop in hint::select_unpredictable

library/core/src/hint.rs

@@ -786,12 +786,45 @@ 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> {
+        // Invariant: valid pointer and points to an initialized value that is not further used,
+        // i.e. it can be dropped by this guard.
+        inner: *mut T,
+    }
+
+    impl<T> Drop for DropOnPanic<T> {
+        fn drop(&mut self) {
+            // SAFETY: Must be guaranteed on construction of local type `DropOnPanic`.
+            unsafe { self.inner.drop_in_place() }
+        }
+    }
+
+    let true_ptr = true_val.as_mut_ptr();
+    let false_ptr = false_val.as_mut_ptr();
+
     // 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. We construct a temporary by-pointer guard around the selected value while
+        // dropping the unselected value. Arguments overlap here, so we can not use mutable
+        // reference for these arguments.
+        let guard = crate::intrinsics::select_unpredictable(condition, true_ptr, false_ptr);
+        let drop = crate::intrinsics::select_unpredictable(condition, false_ptr, true_ptr);
+
+        // SAFETY: both pointers are well-aligned and point to initialized values inside a
+        // `MaybeUninit` each. In both possible values for `condition` the pointer `guard` and
+        // `drop` do not alias (even though the two argument pairs we have selected from did alias
+        // each other).
+        let guard = DropOnPanic { inner: guard };
+        drop.drop_in_place();
+        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,39 @@ fn select_unpredictable_drop() {
     assert!(a_dropped.get());
     assert!(b_dropped.get());
 }
+
+#[test]
+#[should_panic = "message canary"]
+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, "message canary");
+        }
+    }
+
+    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(true, selected, unselected);
+}