diff --git a/src/liballoc/collections/vec_deque.rs b/src/liballoc/collections/vec_deque.rs index 7307aad8a9e..de92c4997e3 100644 --- a/src/liballoc/collections/vec_deque.rs +++ b/src/liballoc/collections/vec_deque.rs @@ -2575,47 +2575,61 @@ unsafe impl Send for Drain<'_, T> {} #[stable(feature = "drain", since = "1.6.0")] impl Drop for Drain<'_, T> { fn drop(&mut self) { - self.for_each(drop); + struct DropGuard<'r, 'a, T>(&'r mut Drain<'a, T>); - let source_deque = unsafe { self.deque.as_mut() }; + impl<'r, 'a, T> Drop for DropGuard<'r, 'a, T> { + fn drop(&mut self) { + self.0.for_each(drop); - // T = source_deque_tail; H = source_deque_head; t = drain_tail; h = drain_head - // - // T t h H - // [. . . o o x x o o . . .] - // - let orig_tail = source_deque.tail; - let drain_tail = source_deque.head; - let drain_head = self.after_tail; - let orig_head = self.after_head; + let source_deque = unsafe { self.0.deque.as_mut() }; - let tail_len = count(orig_tail, drain_tail, source_deque.cap()); - let head_len = count(drain_head, orig_head, source_deque.cap()); + // T = source_deque_tail; H = source_deque_head; t = drain_tail; h = drain_head + // + // T t h H + // [. . . o o x x o o . . .] + // + let orig_tail = source_deque.tail; + let drain_tail = source_deque.head; + let drain_head = self.0.after_tail; + let orig_head = self.0.after_head; - // Restore the original head value - source_deque.head = orig_head; + let tail_len = count(orig_tail, drain_tail, source_deque.cap()); + let head_len = count(drain_head, orig_head, source_deque.cap()); - match (tail_len, head_len) { - (0, 0) => { - source_deque.head = 0; - source_deque.tail = 0; - } - (0, _) => { - source_deque.tail = drain_head; - } - (_, 0) => { - source_deque.head = drain_tail; - } - _ => unsafe { - if tail_len <= head_len { - source_deque.tail = source_deque.wrap_sub(drain_head, tail_len); - source_deque.wrap_copy(source_deque.tail, orig_tail, tail_len); - } else { - source_deque.head = source_deque.wrap_add(drain_tail, head_len); - source_deque.wrap_copy(drain_tail, drain_head, head_len); + // Restore the original head value + source_deque.head = orig_head; + + match (tail_len, head_len) { + (0, 0) => { + source_deque.head = 0; + source_deque.tail = 0; + } + (0, _) => { + source_deque.tail = drain_head; + } + (_, 0) => { + source_deque.head = drain_tail; + } + _ => unsafe { + if tail_len <= head_len { + source_deque.tail = source_deque.wrap_sub(drain_head, tail_len); + source_deque.wrap_copy(source_deque.tail, orig_tail, tail_len); + } else { + source_deque.head = source_deque.wrap_add(drain_tail, head_len); + source_deque.wrap_copy(drain_tail, drain_head, head_len); + } + }, } - }, + } } + + while let Some(item) = self.next() { + let guard = DropGuard(self); + drop(item); + mem::forget(guard); + } + + DropGuard(self); } } diff --git a/src/liballoc/tests/vec_deque.rs b/src/liballoc/tests/vec_deque.rs index 2dc50d0c70e..07f1f098954 100644 --- a/src/liballoc/tests/vec_deque.rs +++ b/src/liballoc/tests/vec_deque.rs @@ -1606,3 +1606,41 @@ fn truncate_leak() { assert_eq!(unsafe { DROPS }, 7); } + +#[test] +fn test_drain_leak() { + static mut DROPS: i32 = 0; + + #[derive(Debug, PartialEq)] + struct D(u32, bool); + + impl Drop for D { + fn drop(&mut self) { + unsafe { + DROPS += 1; + } + + if self.1 { + panic!("panic in `drop`"); + } + } + } + + let mut v = VecDeque::new(); + v.push_back(D(4, false)); + v.push_back(D(5, false)); + v.push_back(D(6, false)); + v.push_front(D(3, false)); + v.push_front(D(2, true)); + v.push_front(D(1, false)); + v.push_front(D(0, false)); + + catch_unwind(AssertUnwindSafe(|| { + v.drain(1..=4); + })).ok(); + + assert_eq!(unsafe { DROPS }, 4); + assert_eq!(v.len(), 3); + drop(v); + assert_eq!(unsafe { DROPS }, 7); +}