diff --git a/src/libcollections/binary_heap.rs b/src/libcollections/binary_heap.rs index 0f05e5796aa..af7112a5cb4 100644 --- a/src/libcollections/binary_heap.rs +++ b/src/libcollections/binary_heap.rs @@ -546,7 +546,7 @@ impl BinaryHeap { #[unstable(feature = "collections", reason = "matches collection reform specification, waiting for dust to settle")] pub fn drain(&mut self) -> Drain { - Drain { iter: self.data.drain() } + Drain { iter: self.data.drain(..) } } /// Drops all items from the binary heap. diff --git a/src/libcollections/lib.rs b/src/libcollections/lib.rs index 992424858e1..ae2ad7751aa 100644 --- a/src/libcollections/lib.rs +++ b/src/libcollections/lib.rs @@ -42,7 +42,8 @@ #![feature(slice_patterns)] #![feature(debug_builders)] #![feature(utf8_error)] -#![cfg_attr(test, feature(rand, rustc_private, test, hash, collections))] +#![cfg_attr(test, feature(rand, rustc_private, test, hash, collections, + collections_drain, collections_range))] #![cfg_attr(test, allow(deprecated))] // rand #![feature(no_std)] diff --git a/src/libcollections/vec.rs b/src/libcollections/vec.rs index 526150915a7..ff3703731e0 100644 --- a/src/libcollections/vec.rs +++ b/src/libcollections/vec.rs @@ -69,6 +69,8 @@ use core::usize; use borrow::{Cow, IntoCow}; +use super::range::RangeArgument; + // FIXME- fix places which assume the max vector allowed has memory usize::MAX. static MAX_MEMORY_SIZE: usize = isize::MAX as usize; @@ -718,36 +720,61 @@ impl Vec { unsafe { other.set_len(0); } } - /// Creates a draining iterator that clears the `Vec` and iterates over - /// the removed items from start to end. + /// Create a draining iterator that removes the specified range in the vector + /// and yields the removed items from start to end. The element range is + /// removed even if the iterator is not consumed until the end. + /// + /// Note: It is unspecified how many elements are removed from the vector, + /// if the `Drain` value is leaked. + /// + /// # Panics + /// + /// Panics if the starting point is greater than the end point or if + /// the end point is greater than the length of the vector. /// /// # Examples /// /// ``` - /// # #![feature(collections)] - /// let mut v = vec!["a".to_string(), "b".to_string()]; - /// for s in v.drain() { - /// // s has type String, not &String - /// println!("{}", s); - /// } - /// assert!(v.is_empty()); + /// # #![feature(collections_drain, collections_range)] + /// + /// // Draining using `..` clears the whole vector. + /// let mut v = vec![1, 2, 3]; + /// let u: Vec<_> = v.drain(..).collect(); + /// assert_eq!(v, &[]); + /// assert_eq!(u, &[1, 2, 3]); /// ``` - #[inline] - #[unstable(feature = "collections", - reason = "matches collection reform specification, waiting for dust to settle")] - pub fn drain(&mut self) -> Drain { + #[unstable(feature = "collections_drain", + reason = "recently added, matches RFC")] + pub fn drain(&mut self, range: R) -> Drain where R: RangeArgument { + // Memory safety + // + // When the Drain is first created, it shortens the length of + // the source vector to make sure no uninitalized or moved-from elements + // are accessible at all if the Drain's destructor never gets to run. + // + // Drain will ptr::read out the values to remove. + // When finished, remaining tail of the vec is copied back to cover + // the hole, and the vector length is restored to the new length. + // + let len = self.len(); + let start = *range.start().unwrap_or(&0); + let end = *range.end().unwrap_or(&len); + assert!(start <= end); + assert!(end <= len); + unsafe { - let begin = *self.ptr as *const T; - let end = if mem::size_of::() == 0 { - (*self.ptr as usize + self.len()) as *const T - } else { - (*self.ptr).offset(self.len() as isize) as *const T - }; - self.set_len(0); + // set self.vec length's to start, to be safe in case Drain is leaked + self.set_len(start); + // Use the borrow in the IterMut to indicate borrowing behavior of the + // whole Drain iterator (like &mut T). + let range_slice = slice::from_raw_parts_mut( + self.as_mut_ptr().offset(start as isize), + end - start); Drain { - ptr: begin, - end: end, - marker: PhantomData, + tail_start: end, + tail_len: len - end, + iter: range_slice.iter_mut(), + vec: self as *mut _, } } } @@ -1799,14 +1826,16 @@ impl Drop for IntoIter { } } -/// An iterator that drains a vector. -#[unsafe_no_drop_flag] -#[unstable(feature = "collections", - reason = "recently added as part of collections reform 2")] -pub struct Drain<'a, T:'a> { - ptr: *const T, - end: *const T, - marker: PhantomData<&'a T>, +/// A draining iterator for `Vec`. +#[unstable(feature = "collections_drain", reason = "recently added")] +pub struct Drain<'a, T: 'a> { + /// Index of tail to preserve + tail_start: usize, + /// Length of tail + tail_len: usize, + /// Current remaining range to remove + iter: slice::IterMut<'a, T>, + vec: *mut Vec, } unsafe impl<'a, T: Sync> Sync for Drain<'a, T> {} @@ -1818,34 +1847,15 @@ impl<'a, T> Iterator for Drain<'a, T> { #[inline] fn next(&mut self) -> Option { - unsafe { - if self.ptr == self.end { - None - } else { - if mem::size_of::() == 0 { - // purposefully don't use 'ptr.offset' because for - // vectors with 0-size elements this would return the - // same pointer. - self.ptr = mem::transmute(self.ptr as usize + 1); - - // Use a non-null pointer value - Some(ptr::read(EMPTY as *mut T)) - } else { - let old = self.ptr; - self.ptr = self.ptr.offset(1); - - Some(ptr::read(old)) - } + self.iter.next().map(|elt| + unsafe { + ptr::read(elt as *const _) } - } + ) } - #[inline] fn size_hint(&self) -> (usize, Option) { - let diff = (self.end as usize) - (self.ptr as usize); - let size = mem::size_of::(); - let exact = diff / (if size == 0 {1} else {size}); - (exact, Some(exact)) + self.iter.size_hint() } } @@ -1853,41 +1863,40 @@ impl<'a, T> Iterator for Drain<'a, T> { impl<'a, T> DoubleEndedIterator for Drain<'a, T> { #[inline] fn next_back(&mut self) -> Option { - unsafe { - if self.end == self.ptr { - None - } else { - if mem::size_of::() == 0 { - // See above for why 'ptr.offset' isn't used - self.end = mem::transmute(self.end as usize - 1); - - // Use a non-null pointer value - Some(ptr::read(EMPTY as *mut T)) - } else { - self.end = self.end.offset(-1); - - Some(ptr::read(self.end)) - } + self.iter.next_back().map(|elt| + unsafe { + ptr::read(elt as *const _) } - } + ) } } -#[stable(feature = "rust1", since = "1.0.0")] -impl<'a, T> ExactSizeIterator for Drain<'a, T> {} - #[unsafe_destructor] #[stable(feature = "rust1", since = "1.0.0")] impl<'a, T> Drop for Drain<'a, T> { fn drop(&mut self) { - // self.ptr == self.end == mem::POST_DROP_USIZE if drop has already been called, - // so we can use #[unsafe_no_drop_flag]. + // exhaust self first + while let Some(_) = self.next() { } - // destroy the remaining elements - for _x in self.by_ref() {} + if self.tail_len > 0 { + unsafe { + let source_vec = &mut *self.vec; + // memmove back untouched tail, update to new length + let start = source_vec.len(); + let tail = self.tail_start; + let src = source_vec.as_ptr().offset(tail as isize); + let dst = source_vec.as_mut_ptr().offset(start as isize); + ptr::copy(src, dst, self.tail_len); + source_vec.set_len(start + self.tail_len); + } + } } } + +#[stable(feature = "rust1", since = "1.0.0")] +impl<'a, T> ExactSizeIterator for Drain<'a, T> {} + //////////////////////////////////////////////////////////////////////////////// // Conversion from &[T] to &Vec //////////////////////////////////////////////////////////////////////////////// diff --git a/src/libcollections/vec_map.rs b/src/libcollections/vec_map.rs index d473504d544..6ff819fc87c 100644 --- a/src/libcollections/vec_map.rs +++ b/src/libcollections/vec_map.rs @@ -418,7 +418,7 @@ impl VecMap { } let filter: fn((usize, Option)) -> Option<(usize, V)> = filter; // coerce to fn ptr - Drain { iter: self.v.drain().enumerate().filter_map(filter) } + Drain { iter: self.v.drain(..).enumerate().filter_map(filter) } } /// Returns the number of elements in the map. diff --git a/src/libcollectionstest/lib.rs b/src/libcollectionstest/lib.rs index 5c109dc8104..935cfb762c2 100644 --- a/src/libcollectionstest/lib.rs +++ b/src/libcollectionstest/lib.rs @@ -10,6 +10,7 @@ #![feature(box_syntax)] #![feature(collections)] +#![feature(collections_drain)] #![feature(core)] #![feature(hash)] #![feature(rand)] diff --git a/src/libcollectionstest/vec.rs b/src/libcollectionstest/vec.rs index 2923bea9828..7d3b156a7cc 100644 --- a/src/libcollectionstest/vec.rs +++ b/src/libcollectionstest/vec.rs @@ -460,7 +460,7 @@ fn test_move_items_zero_sized() { fn test_drain_items() { let mut vec = vec![1, 2, 3]; let mut vec2 = vec![]; - for i in vec.drain() { + for i in vec.drain(..) { vec2.push(i); } assert_eq!(vec, []); @@ -471,7 +471,7 @@ fn test_drain_items() { fn test_drain_items_reverse() { let mut vec = vec![1, 2, 3]; let mut vec2 = vec![]; - for i in vec.drain().rev() { + for i in vec.drain(..).rev() { vec2.push(i); } assert_eq!(vec, []); @@ -482,13 +482,43 @@ fn test_drain_items_reverse() { fn test_drain_items_zero_sized() { let mut vec = vec![(), (), ()]; let mut vec2 = vec![]; - for i in vec.drain() { + for i in vec.drain(..) { vec2.push(i); } assert_eq!(vec, []); assert_eq!(vec2, [(), (), ()]); } +#[test] +#[should_panic] +fn test_drain_out_of_bounds() { + let mut v = vec![1, 2, 3, 4, 5]; + v.drain(5..6); +} + +#[test] +fn test_drain_range() { + let mut v = vec![1, 2, 3, 4, 5]; + for _ in v.drain(4..) { + } + assert_eq!(v, &[1, 2, 3, 4]); + + let mut v: Vec<_> = (1..6).map(|x| x.to_string()).collect(); + for _ in v.drain(1..4) { + } + assert_eq!(v, &[1.to_string(), 5.to_string()]); + + let mut v: Vec<_> = (1..6).map(|x| x.to_string()).collect(); + for _ in v.drain(1..4).rev() { + } + assert_eq!(v, &[1.to_string(), 5.to_string()]); + + let mut v: Vec<_> = vec![(); 5]; + for _ in v.drain(1..4).rev() { + } + assert_eq!(v, &[(), ()]); +} + #[test] fn test_into_boxed_slice() { let xs = vec![1, 2, 3]; diff --git a/src/test/run-pass/sync-send-iterators-in-libcollections.rs b/src/test/run-pass/sync-send-iterators-in-libcollections.rs index 447b4de450b..d7877bff0cb 100644 --- a/src/test/run-pass/sync-send-iterators-in-libcollections.rs +++ b/src/test/run-pass/sync-send-iterators-in-libcollections.rs @@ -12,6 +12,7 @@ #![allow(unused_mut)] #![feature(collections)] +#![feature(collections_drain)] extern crate collections; @@ -96,5 +97,6 @@ fn main() { all_sync_send!(VecMap::::new(), iter, iter_mut, drain, into_iter, keys, values); - all_sync_send!(Vec::::new(), into_iter, drain); + all_sync_send!(Vec::::new(), into_iter); + is_sync_send!(Vec::::new(), drain(..)); }