diff --git a/src/librustc_data_structures/accumulate_vec.rs b/src/librustc_data_structures/accumulate_vec.rs index 937cb3f6007..78af655852d 100644 --- a/src/librustc_data_structures/accumulate_vec.rs +++ b/src/librustc_data_structures/accumulate_vec.rs @@ -19,6 +19,7 @@ use std::ops::{Deref, DerefMut}; use std::iter::{self, IntoIterator, FromIterator}; use std::slice; use std::vec; +use std::collections::range::RangeArgument; use rustc_serialize::{Encodable, Encoder, Decodable, Decoder}; @@ -71,6 +72,19 @@ impl AccumulateVec { AccumulateVec::Heap(ref mut vec) => vec.pop(), } } + + pub fn drain(&mut self, range: R) -> Drain + where R: RangeArgument + { + match *self { + AccumulateVec::Array(ref mut v) => { + Drain::Array(v.drain(range)) + }, + AccumulateVec::Heap(ref mut v) => { + Drain::Heap(v.drain(range)) + }, + } + } } impl Deref for AccumulateVec { @@ -132,6 +146,31 @@ impl Iterator for IntoIter { } } +pub enum Drain<'a, A: Array> + where A::Element: 'a +{ + Array(array_vec::Drain<'a, A>), + Heap(vec::Drain<'a, A::Element>), +} + +impl<'a, A: Array> Iterator for Drain<'a, A> { + type Item = A::Element; + + fn next(&mut self) -> Option { + match *self { + Drain::Array(ref mut drain) => drain.next(), + Drain::Heap(ref mut drain) => drain.next(), + } + } + + fn size_hint(&self) -> (usize, Option) { + match *self { + Drain::Array(ref drain) => drain.size_hint(), + Drain::Heap(ref drain) => drain.size_hint(), + } + } +} + impl IntoIterator for AccumulateVec { type Item = A::Element; type IntoIter = IntoIter; diff --git a/src/librustc_data_structures/array_vec.rs b/src/librustc_data_structures/array_vec.rs index 631cf2cfcf6..844e9041d20 100644 --- a/src/librustc_data_structures/array_vec.rs +++ b/src/librustc_data_structures/array_vec.rs @@ -12,12 +12,13 @@ use std::marker::Unsize; use std::iter::Extend; -use std::ptr::{self, drop_in_place}; +use std::ptr::{self, drop_in_place, Shared}; use std::ops::{Deref, DerefMut, Range}; use std::hash::{Hash, Hasher}; use std::slice; use std::fmt; use std::mem; +use std::collections::range::RangeArgument; pub unsafe trait Array { type Element; @@ -103,6 +104,44 @@ impl ArrayVec { None } } + + 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 { + // 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 = { + let arr = &mut self.values as &mut [ManuallyDrop<_>]; + slice::from_raw_parts_mut(arr.as_mut_ptr().offset(start as isize), + end - start) + }; + Drain { + tail_start: end, + tail_len: len - end, + iter: range_slice.iter(), + array_vec: Shared::new(self as *mut _), + } + } + } } impl Default for ArrayVec @@ -179,6 +218,51 @@ impl Iterator for Iter { } } +pub struct Drain<'a, A: Array> + where A::Element: 'a +{ + tail_start: usize, + tail_len: usize, + iter: slice::Iter<'a, ManuallyDrop>, + array_vec: Shared>, +} + +impl<'a, A: Array> Iterator for Drain<'a, A> { + type Item = A::Element; + + #[inline] + fn next(&mut self) -> Option { + self.iter.next().map(|elt| unsafe { ptr::read(elt as *const ManuallyDrop<_>).value }) + } + + fn size_hint(&self) -> (usize, Option) { + self.iter.size_hint() + } +} + +impl<'a, A: Array> Drop for Drain<'a, A> { + fn drop(&mut self) { + // exhaust self first + while let Some(_) = self.next() {} + + if self.tail_len > 0 { + unsafe { + let source_array_vec = &mut **self.array_vec; + // memmove back untouched tail, update to new length + let start = source_array_vec.len(); + let tail = self.tail_start; + { + let mut arr = &mut source_array_vec.values as &mut [ManuallyDrop<_>]; + let src = arr.as_ptr().offset(tail as isize); + let dst = arr.as_mut_ptr().offset(start as isize); + ptr::copy(src, dst, self.tail_len); + }; + source_array_vec.set_len(start + self.tail_len); + } + } + } +} + impl IntoIterator for ArrayVec { type Item = A::Element; type IntoIter = Iter; diff --git a/src/librustc_data_structures/lib.rs b/src/librustc_data_structures/lib.rs index d3ec674daed..ee75a3596e1 100644 --- a/src/librustc_data_structures/lib.rs +++ b/src/librustc_data_structures/lib.rs @@ -25,6 +25,8 @@ html_root_url = "https://doc.rust-lang.org/nightly/")] #![cfg_attr(not(stage0), deny(warnings))] +#![feature(shared)] +#![feature(collections_range)] #![feature(nonzero)] #![feature(rustc_private)] #![feature(staged_api)]