Drain

Let’s move on to Drain. Drain is largely the same as IntoIter, except thatinstead of consuming the Vec, it borrows the Vec and leaves its allocationuntouched. For now we’ll only implement the “basic” full-range version.

  1. use std::marker::PhantomData;
  3. struct Drain<'a, T: 'a> {
  4.     // Need to bound the lifetime here, so we do it with `&'a mut Vec<T>`
  5.     // because that's semantically what we contain. We're "just" calling
  6.     // `pop()` and `remove(0)`.
  7.     vec: PhantomData<&'a mut Vec<T>>,
  8.     start: *const T,
  9.     end: *const T,
  10. }
  12. impl<'a, T> Iterator for Drain<'a, T> {
  13.     type Item = T;
  14.     fn next(&mut self) -> Option<T> {
  15.         if self.start == self.end {
  16.             None

— wait, this is seeming familiar. Let’s do some more compression. BothIntoIter and Drain have the exact same structure, let’s just factor it out.

  1. struct RawValIter<T> {
  2.     start: *const T,
  3.     end: *const T,
  4. }
  6. impl<T> RawValIter<T> {
  7.     // unsafe to construct because it has no associated lifetimes.
  8.     // This is necessary to store a RawValIter in the same struct as
  9.     // its actual allocation. OK since it's a private implementation
  10.     // detail.
  11.     unsafe fn new(slice: &[T]) -> Self {
  12.         RawValIter {
  13.             start: slice.as_ptr(),
  14.             end: if slice.len() == 0 {
  15.                 // if `len = 0`, then this is not actually allocated memory.
  16.                 // Need to avoid offsetting because that will give wrong
  17.                 // information to LLVM via GEP.
  18.                 slice.as_ptr()
  19.             } else {
  20.                 slice.as_ptr().offset(slice.len() as isize)
  21.             }
  22.         }
  23.     }
  24. }
  26. // Iterator and DoubleEndedIterator impls identical to IntoIter.

And IntoIter becomes the following:

  1. pub struct IntoIter<T> {
  2.     _buf: RawVec<T>, // we don't actually care about this. Just need it to live.
  3.     iter: RawValIter<T>,
  4. }
  6. impl<T> Iterator for IntoIter<T> {
  7.     type Item = T;
  8.     fn next(&mut self) -> Option<T> { self.iter.next() }
  9.     fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
  10. }
  12. impl<T> DoubleEndedIterator for IntoIter<T> {
  13.     fn next_back(&mut self) -> Option<T> { self.iter.next_back() }
  14. }
  16. impl<T> Drop for IntoIter<T> {
  17.     fn drop(&mut self) {
  18.         for _ in &mut self.iter {}
  19.     }
  20. }
  22. impl<T> Vec<T> {
  23.     pub fn into_iter(self) -> IntoIter<T> {
  24.         unsafe {
  25.             let iter = RawValIter::new(&self);
  27.             let buf = ptr::read(&self.buf);
  28.             mem::forget(self);
  30.             IntoIter {
  31.                 iter: iter,
  32.                 _buf: buf,
  33.             }
  34.         }
  35.     }
  36. }

Note that I’ve left a few quirks in this design to make upgrading Drain to workwith arbitrary subranges a bit easier. In particular we could have RawValIterdrain itself on drop, but that won’t work right for a more complex Drain.We also take a slice to simplify Drain initialization.

Alright, now Drain is really easy:

  1. use std::marker::PhantomData;
  3. pub struct Drain<'a, T: 'a> {
  4.     vec: PhantomData<&'a mut Vec<T>>,
  5.     iter: RawValIter<T>,
  6. }
  8. impl<'a, T> Iterator for Drain<'a, T> {
  9.     type Item = T;
  10.     fn next(&mut self) -> Option<T> { self.iter.next() }
  11.     fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
  12. }
  14. impl<'a, T> DoubleEndedIterator for Drain<'a, T> {
  15.     fn next_back(&mut self) -> Option<T> { self.iter.next_back() }
  16. }
  18. impl<'a, T> Drop for Drain<'a, T> {
  19.     fn drop(&mut self) {
  20.         for _ in &mut self.iter {}
  21.     }
  22. }
  24. impl<T> Vec<T> {
  25.     pub fn drain(&mut self) -> Drain<T> {
  26.         unsafe {
  27.             let iter = RawValIter::new(&self);
  29.             // this is a mem::forget safety thing. If Drain is forgotten, we just
  30.             // leak the whole Vec's contents. Also we need to do this *eventually*
  31.             // anyway, so why not do it now?
  32.             self.len = 0;
  34.             Drain {
  35.                 iter: iter,
  36.                 vec: PhantomData,
  37.             }
  38.         }
  39.     }
  40. }

For more details on the mem::forget problem, see thesection on leaks.