Add drain method to AccumulateVec/ArrayVec

You can now call .drain(..) on SmallVec, AccumulateVec and ArrayVec

Author: canndrew

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<A: Array> AccumulateVec<A> {
             AccumulateVec::Heap(ref mut vec) => vec.pop(),
         }
     }
+
+    pub fn drain<R>(&mut self, range: R) -> Drain<A>
+        where R: RangeArgument<usize>
+    {
+        match *self {
+            AccumulateVec::Array(ref mut v) => {
+                Drain::Array(v.drain(range))
+            },
+            AccumulateVec::Heap(ref mut v) => {
+                Drain::Heap(v.drain(range))
+            },
+        }
+    }
 }
 
 impl<A: Array> Deref for AccumulateVec<A> {
@@ -132,6 +146,31 @@ impl<A: Array> Iterator for IntoIter<A> {
     }
 }
 
+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<A::Element> {
+        match *self {
+            Drain::Array(ref mut drain) => drain.next(),
+            Drain::Heap(ref mut drain) => drain.next(),
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        match *self {
+            Drain::Array(ref drain) => drain.size_hint(),
+            Drain::Heap(ref drain) => drain.size_hint(),
+        }
+    }
+}
+
 impl<A: Array> IntoIterator for AccumulateVec<A> {
     type Item = A::Element;
     type IntoIter = IntoIter<A>;

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<A: Array> ArrayVec<A> {
             None
         }
     }
+
+    pub fn drain<R>(&mut self, range: R) -> Drain<A>
+        where R: RangeArgument<usize>
+    {
+        // 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<A> Default for ArrayVec<A>
@@ -179,6 +218,51 @@ impl<A: Array> Iterator for Iter<A> {
     }
 }
 
+pub struct Drain<'a, A: Array>
+        where A::Element: 'a
+{
+    tail_start: usize,
+    tail_len: usize,
+    iter: slice::Iter<'a, ManuallyDrop<A::Element>>,
+    array_vec: Shared<ArrayVec<A>>,
+}
+
+impl<'a, A: Array> Iterator for Drain<'a, A> {
+    type Item = A::Element;
+
+    #[inline]
+    fn next(&mut self) -> Option<A::Element> {
+        self.iter.next().map(|elt| unsafe { ptr::read(elt as *const ManuallyDrop<_>).value })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        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<A: Array> IntoIterator for ArrayVec<A> {
     type Item = A::Element;
     type IntoIter = Iter<A>;

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)]