src/doc/nomicon/src/vec-drain.md

   1 # Drain
   2
   3 Let's move on to Drain. Drain is largely the same as IntoIter, except that
   4 instead of consuming the Vec, it borrows the Vec and leaves its allocation
   5 untouched. For now we'll only implement the "basic" full-range version.
   6
   7 ```rust,ignore
   8 use std::marker::PhantomData;
   9
  10 struct Drain<'a, T: 'a> {
  11     // Need to bound the lifetime here, so we do it with `&'a mut Vec<T>`
  12     // because that's semantically what we contain. We're "just" calling
  13     // `pop()` and `remove(0)`.
  14     vec: PhantomData<&'a mut Vec<T>>
  15     start: *const T,
  16     end: *const T,
  17 }
  18
  19 impl<'a, T> Iterator for Drain<'a, T> {
  20     type Item = T;
  21     fn next(&mut self) -> Option<T> {
  22         if self.start == self.end {
  23             None
  24 ```
  25
  26 -- wait, this is seeming familiar. Let's do some more compression. Both
  27 IntoIter and Drain have the exact same structure, let's just factor it out.
  28
  29 ```rust
  30 struct RawValIter<T> {
  31     start: *const T,
  32     end: *const T,
  33 }
  34
  35 impl<T> RawValIter<T> {
  36     // unsafe to construct because it has no associated lifetimes.
  37     // This is necessary to store a RawValIter in the same struct as
  38     // its actual allocation. OK since it's a private implementation
  39     // detail.
  40     unsafe fn new(slice: &[T]) -> Self {
  41         RawValIter {
  42             start: slice.as_ptr(),
  43             end: if slice.len() == 0 {
  44                 // if `len = 0`, then this is not actually allocated memory.
  45                 // Need to avoid offsetting because that will give wrong
  46                 // information to LLVM via GEP.
  47                 slice.as_ptr()
  48             } else {
  49                 slice.as_ptr().offset(slice.len() as isize)
  50             }
  51         }
  52     }
  53 }
  54
  55 // Iterator and DoubleEndedIterator impls identical to IntoIter.
  56 ```
  57
  58 And IntoIter becomes the following:
  59
  60 ```rust,ignore
  61 pub struct IntoIter<T> {
  62     _buf: RawVec<T>, // we don't actually care about this. Just need it to live.
  63     iter: RawValIter<T>,
  64 }
  65
  66 impl<T> Iterator for IntoIter<T> {
  67     type Item = T;
  68     fn next(&mut self) -> Option<T> { self.iter.next() }
  69     fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
  70 }
  71
  72 impl<T> DoubleEndedIterator for IntoIter<T> {
  73     fn next_back(&mut self) -> Option<T> { self.iter.next_back() }
  74 }
  75
  76 impl<T> Drop for IntoIter<T> {
  77     fn drop(&mut self) {
  78         for _ in &mut self.iter {}
  79     }
  80 }
  81
  82 impl<T> Vec<T> {
  83     pub fn into_iter(self) -> IntoIter<T> {
  84         unsafe {
  85             let iter = RawValIter::new(&self);
  86
  87             let buf = ptr::read(&self.buf);
  88             mem::forget(self);
  89
  90             IntoIter {
  91                 iter: iter,
  92                 _buf: buf,
  93             }
  94         }
  95     }
  96 }
  97 ```
  98
  99 Note that I've left a few quirks in this design to make upgrading Drain to work
 100 with arbitrary subranges a bit easier. In particular we *could* have RawValIter
 101 drain itself on drop, but that won't work right for a more complex Drain.
 102 We also take a slice to simplify Drain initialization.
 103
 104 Alright, now Drain is really easy:
 105
 106 ```rust,ignore
 107 use std::marker::PhantomData;
 108
 109 pub struct Drain<'a, T: 'a> {
 110     vec: PhantomData<&'a mut Vec<T>>,
 111     iter: RawValIter<T>,
 112 }
 113
 114 impl<'a, T> Iterator for Drain<'a, T> {
 115     type Item = T;
 116     fn next(&mut self) -> Option<T> { self.iter.next() }
 117     fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
 118 }
 119
 120 impl<'a, T> DoubleEndedIterator for Drain<'a, T> {
 121     fn next_back(&mut self) -> Option<T> { self.iter.next_back() }
 122 }
 123
 124 impl<'a, T> Drop for Drain<'a, T> {
 125     fn drop(&mut self) {
 126         for _ in &mut self.iter {}
 127     }
 128 }
 129
 130 impl<T> Vec<T> {
 131     pub fn drain(&mut self) -> Drain<T> {
 132         unsafe {
 133             let iter = RawValIter::new(&self);
 134
 135             // this is a mem::forget safety thing. If Drain is forgotten, we just
 136             // leak the whole Vec's contents. Also we need to do this *eventually*
 137             // anyway, so why not do it now?
 138             self.len = 0;
 139
 140             Drain {
 141                 iter: iter,
 142                 vec: PhantomData,
 143             }
 144         }
 145     }
 146 }
 147 ```
 148
 149 For more details on the `mem::forget` problem, see the
 150 [section on leaks][leaks].
 151
 152 [leaks]: leaking.html