diff --git a/src/libcollections/binary_heap.rs b/src/libcollections/binary_heap.rs index 2e575ddb00a32..a778948ed90e8 100644 --- a/src/libcollections/binary_heap.rs +++ b/src/libcollections/binary_heap.rs @@ -554,7 +554,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/string.rs b/src/libcollections/string.rs index 33189bd68bdc6..2e5b967a18a6f 100644 --- a/src/libcollections/string.rs +++ b/src/libcollections/string.rs @@ -22,14 +22,14 @@ use core::fmt; use core::hash; use core::iter::{IntoIterator, FromIterator}; use core::mem; -use core::ops::{self, Deref, Add, Index}; +use core::ops::{self, Deref, Add, Index, Range, RangeTo, RangeFrom, RangeFull}; use core::ptr; use core::raw::Slice as RawSlice; use unicode::str as unicode_str; use unicode::str::Utf16Item; use borrow::{Cow, IntoCow}; -use str::{self, CharRange, FromStr, Utf8Error}; +use str::{self, CharRange, FromStr, Utf8Error, Chars}; use vec::{DerefVec, Vec, as_vec}; /// A growable string stored as a UTF-8 encoded buffer. @@ -670,6 +670,123 @@ impl String { pub fn clear(&mut self) { self.vec.clear() } + + /// Creates a draining iterator that clears the specified byte-range in the String + /// and iterates over the characters contained in the range, backshifting the + /// remaining bytes. + /// + /// # Example + /// + /// ``` + /// let mut s = "Hello Wörld!".to_string(); + /// let s2: String = s.drain(6..12).collect(); + /// assert_eq!(s, "Hello !"); + /// assert_eq!(s2, "Wörld"); + /// ``` + /// + /// # Panics + /// + /// Panics if the range is decreasing, if the upper bound is larger than the + /// length of the String, or if the start and the end of the range don't lie on + /// character boundaries. + pub fn drain<'a, T: DrainRange>(&'a mut self, range: T) -> Drain<'a> { + range.drain(self) + } +} + +/// A trait for draining a string. +/// +/// See the documentation of `String::drain`. +pub trait DrainRange { + fn drain<'a>(&self, s: &'a mut String) -> Drain<'a>; +} + +impl DrainRange for Range { + fn drain<'a>(&self, s: &'a mut String) -> Drain<'a> { + assert!(self.start <= self.end, "Range not increasing"); + assert!(self.end <= s.len(), "Range out of bounds"); + unsafe { + let slice = mem::transmute::<&str, &'static str>(&s[self.start..self.end]); + let tail = s.len() - self.end; + s.as_mut_vec().set_len(tail + self.start); + let ptr = s.as_mut_vec().as_mut_ptr(); + Drain { + tail: tail, + start: ptr.offset(self.start as isize), + end: ptr.offset(self.end as isize), + chars: slice.chars(), + } + } + } +} + +impl DrainRange for RangeFrom { + fn drain<'a>(&self, s: &'a mut String) -> Drain<'a> { + assert!(self.start <= s.len(), "Range out of bounds"); + (self.start..s.len()).drain(s) + } +} + +impl DrainRange for RangeTo { + fn drain<'a>(&self, s: &'a mut String) -> Drain<'a> { + (0..self.end).drain(s) + } +} + +impl DrainRange for RangeFull { + fn drain<'a>(&self, s: &'a mut String) -> Drain<'a> { + (0..s.len()).drain(s) + } +} + +impl DrainRange for usize { + fn drain<'a>(&self, s: &'a mut String) -> Drain<'a> { + (*self..*self+1).drain(s) + } +} + +/// An iterator that drains part of string. +#[unsafe_no_drop_flag] +pub struct Drain<'a> { + tail: usize, + start: *const u8, + end: *const u8, + chars: Chars<'a>, +} + +unsafe impl<'a> Sync for Drain<'a> {} +unsafe impl<'a> Send for Drain<'a> {} + +impl<'a> Iterator for Drain<'a> { + type Item = char; + + #[inline] + fn next(&mut self) -> Option { + self.chars.next() + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.chars.size_hint() + } +} + +impl<'a> DoubleEndedIterator for Drain<'a> { + #[inline] + fn next_back(&mut self) -> Option { + self.chars.next_back() + } +} + +#[unsafe_destructor] +impl<'a> Drop for Drain<'a> { + fn drop(&mut self) { + // self.start == null if drop has already been called, so we can use + // #[unsafe_no_drop_flag]. + if !self.start.is_null() { + unsafe { ptr::copy(self.start as *mut _, self.end, self.tail); } + } + } } impl FromUtf8Error { @@ -1442,4 +1559,12 @@ mod tests { r }); } + + #[test] + fn test_drain() { + let mut s = "Hello Wörld!".to_string(); + let s2: String = s.drain(6..12).collect(); + assert_eq!(s, "Hello !"); + assert_eq!(s2, "Wörld"); + } } diff --git a/src/libcollections/vec.rs b/src/libcollections/vec.rs index 2e947ea24602a..f5511f2e60388 100644 --- a/src/libcollections/vec.rs +++ b/src/libcollections/vec.rs @@ -60,9 +60,10 @@ use core::iter::{repeat, FromIterator, IntoIterator}; use core::marker::PhantomData; use core::mem; use core::num::{Int, UnsignedInt}; -use core::ops::{Index, IndexMut, Deref, Add}; +use core::ops::{Index, IndexMut, Deref, Add, Range, RangeFrom, RangeFull, RangeTo}; use core::ops; use core::ptr; +use core::cmp; use core::ptr::Unique; use core::raw::Slice as RawSlice; use core::slice; @@ -731,37 +732,29 @@ 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. + /// Creates a draining iterator that clears the specified range in the `Vec` and + /// iterates over the removed items from start to end. /// /// # Examples /// /// ``` - /// let mut v = vec!["a".to_string(), "b".to_string()]; - /// for s in v.drain() { - /// // s has type String, not &String + /// let mut v = vec![0, 1, 2, 3, 4, 5, 6]; + /// for d in v.drain(1..6) { + /// // d has type i32, not &i32 /// println!("{}", s); /// } - /// assert!(v.is_empty()); + /// assert_eq!(v, vec![0, 6]); /// ``` + /// + /// # Panics + /// + /// Panics if the range is decreasing or if the upper bound is larger than the + /// length of the vector. #[inline] - #[unstable(feature = "collections", - reason = "matches collection reform specification, waiting for dust to settle")] - pub fn drain(&mut self) -> Drain { - 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); - Drain { - ptr: begin, - end: end, - marker: PhantomData, - } - } + #[unstable(feature = "collections", reason = "matches drain RFC, + waiting for dust to settle")] + pub fn drain<'a, U: DrainRange>(&'a mut self, range: U) -> Drain<'a, T> { + range.drain(self) } /// Clears the vector, removing all values. @@ -1748,14 +1741,86 @@ impl Drop for IntoIter { } } -/// An iterator that drains a vector. +/// A trait for draining a vector. +/// +/// See the documentation of `Vec::drain`. +pub trait DrainRange { + fn drain<'a, T>(&self, vec: &'a mut Vec) -> Drain<'a, T>; +} + +impl DrainRange for Range { + fn drain<'a, T>(&self, vec: &'a mut Vec) -> Drain<'a, T> { + assert!(self.start <= self.end, "Range not increasing"); + assert!(self.end <= vec.len(), "Range out of bounds"); + unsafe { + let tail = vec.len() - self.end; + vec.set_len(tail + self.start); + let ptr = vec.as_ptr(); + let (start, end) = if mem::size_of::() == 0 { + // make sure that start is not null + ((self.start + 1) as *const T, (self.end + 1) as *const T) + } else { + (ptr.offset(self.start as isize), ptr.offset(self.end as isize)) + }; + Drain { + tail: tail, + start: start, + end: end, + left: start, + right: end, + marker1: PhantomData, + marker2: PhantomData, + } + } + } +} + +impl DrainRange for RangeFrom { + fn drain<'a, T>(&self, vec: &'a mut Vec) -> Drain<'a, T> { + assert!(self.start <= vec.len(), "Range out of bounds"); + (self.start..vec.len()).drain(vec) + } +} + +impl DrainRange for RangeTo { + fn drain<'a, T>(&self, vec: &'a mut Vec) -> Drain<'a, T> { + (0..self.end).drain(vec) + } +} + +impl DrainRange for RangeFull { + fn drain<'a, T>(&self, vec: &'a mut Vec) -> Drain<'a, T> { + (0..vec.len()).drain(vec) + } +} + +impl DrainRange for usize { + fn drain<'a, T>(&self, vec: &'a mut Vec) -> Drain<'a, T> { + (*self..*self+1).drain(vec) + } +} + +/// An iterator that drains part of 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>, +pub struct Drain<'a, T> { + tail: usize, + start: *const T, + end: *const T, + left: *const T, + right: *const T, + marker1: PhantomData<&'a ()>, + // Drain contains functions to retrieve T but none to insert T. + marker2: PhantomData T>, + + // The straightforward marker would be &'a mut Vec. If we were writing this in safe + // code, that's what we would have after all. However, that would also induce + // invariance on T, which given that Drain only ever extracts values of T is stricter + // than necessary. Therefore, we use this more subtle formulation, which uses a &'a () + // marker to bind the lifetime securing the vector, and which uses a second marker to + // express that we have a way of producing T instances that we are going to employ. + // This gives covariance in T. } unsafe impl<'a, T: Sync> Sync for Drain<'a, T> {} @@ -1768,20 +1833,20 @@ impl<'a, T> Iterator for Drain<'a, T> { #[inline] fn next(&mut self) -> Option { unsafe { - if self.ptr == self.end { + if self.left == self.right { 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); + self.left = (self.left as usize + 1) as *mut T; // Use a non-null pointer value Some(ptr::read(EMPTY as *mut T)) } else { - let old = self.ptr; - self.ptr = self.ptr.offset(1); + let old = self.left; + self.left = self.left.offset(1); Some(ptr::read(old)) } @@ -1791,9 +1856,9 @@ impl<'a, T> Iterator for Drain<'a, T> { #[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}); + let diff = (self.right as usize) - (self.left as usize); + let size = cmp::max(mem::size_of::(), 1); + let exact = diff / size; (exact, Some(exact)) } } @@ -1803,19 +1868,19 @@ impl<'a, T> DoubleEndedIterator for Drain<'a, T> { #[inline] fn next_back(&mut self) -> Option { unsafe { - if self.end == self.ptr { + if self.left == self.right { 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); + self.right = (self.right as usize - 1) as *mut T; // Use a non-null pointer value Some(ptr::read(EMPTY as *mut T)) } else { - self.end = self.end.offset(-1); + self.right = self.right.offset(-1); - Some(ptr::read(self.end)) + Some(ptr::read(self.right)) } } } @@ -1829,11 +1894,16 @@ impl<'a, T> ExactSizeIterator for Drain<'a, T> {} #[stable(feature = "rust1", since = "1.0.0")] impl<'a, T> Drop for Drain<'a, T> { fn drop(&mut self) { - // self.ptr == self.end == null if drop has already been called, - // so we can use #[unsafe_no_drop_flag]. + // self.start == null if drop has already been called, so we can use + // #[unsafe_no_drop_flag]. + if !self.start.is_null() { + // destroy the remaining elements + for _x in self.by_ref() {} - // destroy the remaining elements - for _x in self.by_ref() {} + if mem::size_of::() > 0 { + unsafe { ptr::copy(self.start as *mut _, self.end, self.tail); } + } + } } } @@ -1958,6 +2028,7 @@ impl Drop for PartialVecZeroSized { mod tests { use prelude::*; use core::mem::size_of; + use core::atomic; use core::iter::repeat; use test::Bencher; use super::as_vec; @@ -2406,37 +2477,58 @@ mod tests { #[test] fn test_drain_items() { - let mut vec = vec![1, 2, 3]; + let mut vec = vec![0, 1, 2]; let mut vec2 = vec![]; - for i in vec.drain() { + for i in vec.drain(1..2) { vec2.push(i); } - assert_eq!(vec, []); - assert_eq!(vec2, [ 1, 2, 3 ]); + assert_eq!(vec, [0, 2]); + assert_eq!(vec2, [1]); } #[test] fn test_drain_items_reverse() { - let mut vec = vec![1, 2, 3]; + let mut vec = vec![0, 1, 2]; let mut vec2 = vec![]; - for i in vec.drain().rev() { + for i in vec.drain(1..).rev() { vec2.push(i); } - assert_eq!(vec, []); - assert_eq!(vec2, [3, 2, 1]); + assert_eq!(vec, [0]); + assert_eq!(vec2, [2, 1]); } #[test] 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] + fn test_drain_drop() { + // Assumes that every test is run only once. + static COUNTER: atomic::AtomicUsize = atomic::ATOMIC_USIZE_INIT; + + #[derive(Debug, PartialEq)] + struct X(u8); + + impl Drop for X { + fn drop(&mut self) { + COUNTER.fetch_add(1, atomic::Ordering::SeqCst); + } + } + + let mut vec = vec![X(0), X(1), X(2), X(3)]; + vec.drain(1..3); + let vec2 = vec![X(0), X(3)]; + assert_eq!(vec, vec2); + assert_eq!(COUNTER.load(atomic::Ordering::SeqCst), 2); + } + #[test] fn test_into_boxed_slice() { let xs = vec![1, 2, 3]; diff --git a/src/libcollections/vec_map.rs b/src/libcollections/vec_map.rs index 515de74e340ac..a3176bd83d82c 100644 --- a/src/libcollections/vec_map.rs +++ b/src/libcollections/vec_map.rs @@ -435,7 +435,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) } } /// Return the number of elements in the map. 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 6d29d5f2b3e91..9e6877199dec4 100644 --- a/src/test/run-pass/sync-send-iterators-in-libcollections.rs +++ b/src/test/run-pass/sync-send-iterators-in-libcollections.rs @@ -94,5 +94,13 @@ 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); + + let mut vec = Vec::::new(); + is_send(vec.drain(..)); + is_sync(vec.drain(..)); + + let mut string = String::new(); + is_send(string.drain(..)); + is_sync(string.drain(..)); }