Ensure consistent drop for hint::select_unpredictable

There are a few alternatives to the implementation. The principal
problem is that the selected value must be owned (in the sense of having
any drop flag of sorts) when the unselected value is dropped, such that
panic unwind goes through the drop of both. This ownership must then be
passed on in return when the drop went smoothly. The basic way of
achieving this is by extracting the selected value first, at the cost of
relying on the optimizer a little more for detecting the copy as
constructing the return value despite having a place in the body.

2 files changed, 42 insertions, 1 deletions

diff --git a/library/core/src/hint.rs b/library/core/src/hint.rs
index c72eeb9a9c9..53c17d6c887 100644
--- a/library/core/src/hint.rs
+++ b/library/core/src/hint.rs
@@ -790,8 +790,12 @@ pub fn select_unpredictable<T>(condition: bool, true_val: T, false_val: T) -> T
     // is returned. This is necessary because the intrinsic doesn't drop the
     // value that is  not selected.
     unsafe {
+        // Extract the selected value first, ensure it is dropped as well if dropping the unselected
+        // value panics.
+        let ret = crate::intrinsics::select_unpredictable(condition, &true_val, &false_val)
+            .assume_init_read();
         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()
+        ret
     }
 }
diff --git a/library/coretests/tests/hint.rs b/library/coretests/tests/hint.rs
index 032bbc1dcc8..9ef5567681e 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<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);
+}