Fix panic safety when reserving during splice

src/lib.rs

@@ -383,11 +383,13 @@ impl<T, const N: usize> Drain<'_, T, N> {
         let vec = unsafe { self.vec.as_mut() };
         let len = self.tail_start + self.tail_len;
 
-        // Test
+        // Include the tail when reserving so it survives a reallocation.
         let old_len = vec.len();
         unsafe { vec.set_len(len) }
-        vec.reserve(additional);
+        let result = vec.try_reserve(additional);
+        // Restore the prefix length before a reservation error can panic.
         unsafe { vec.set_len(old_len) };
+        infallible(result);
 
         let new_tail_start = self.tail_start + additional;
         unsafe {

tests/main.rs

@@ -227,6 +227,52 @@ fn splice_inline_fill_then_move_tail_ub_test() {
     assert!(!v.spilled());
 }
 
+#[test]
+fn splice_reserve_panic() {
+    struct CountDrop<'a>(&'a Cell<usize>);
+
+    impl Drop for CountDrop<'_> {
+        fn drop(&mut self) {
+            self.0.set(self.0.get() + 1);
+        }
+    }
+
+    for capacity in [4, 8] {
+        for additional in [usize::MAX, isize::MAX as usize] {
+            let drops = Cell::new(0);
+            let mut v: SmallVec<CountDrop<'_>, 4> = SmallVec::with_capacity(capacity);
+            v.push(CountDrop(&drops));
+
+            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+                drop(v.splice(
+                    0..0,
+                    std::iter::repeat_with(|| CountDrop(&drops)).take(additional)
+                ));
+            }));
+
+            assert!(result.is_err());
+            assert_eq!(v.len(), 1);
+            assert_eq!(drops.get(), 0);
+            drop(v);
+            assert_eq!(drops.get(), 1);
+        }
+    }
+}
+
+#[test]
+fn splice_spill_preserves_tail() {
+    let mut v: SmallVec<Box<usize>, 4> = (0..4).map(Box::new).collect();
+    assert!(!v.spilled());
+
+    drop(v.splice(1..2, (10..15).map(Box::new)));
+
+    assert!(v.spilled());
+    assert_eq!(
+        v.iter().map(|value| **value).collect::<Vec<_>>(),
+        [0, 10, 11, 12, 13, 14, 2, 3]
+    );
+}
+
 #[test]
 fn into_iter() {
     let mut v: SmallVec<u8, 2> = SmallVec::new();