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

   1 # The Final Code
   2
   3 ```rust
   4 #![feature(unique)]
   5 #![feature(alloc, heap_api)]
   6
   7 extern crate alloc;
   8
   9 use std::ptr::{Unique, self};
  10 use std::mem;
  11 use std::ops::{Deref, DerefMut};
  12 use std::marker::PhantomData;
  13
  14 use alloc::heap;
  15
  16 struct RawVec<T> {
  17     ptr: Unique<T>,
  18     cap: usize,
  19 }
  20
  21 impl<T> RawVec<T> {
  22     fn new() -> Self {
  23         unsafe {
  24             // !0 is usize::MAX. This branch should be stripped at compile time.
  25             let cap = if mem::size_of::<T>() == 0 { !0 } else { 0 };
  26
  27             // heap::EMPTY doubles as "unallocated" and "zero-sized allocation"
  28             RawVec { ptr: Unique::new(heap::EMPTY as *mut T), cap: cap }
  29         }
  30     }
  31
  32     fn grow(&mut self) {
  33         unsafe {
  34             let elem_size = mem::size_of::<T>();
  35
  36             // since we set the capacity to usize::MAX when elem_size is
  37             // 0, getting to here necessarily means the Vec is overfull.
  38             assert!(elem_size != 0, "capacity overflow");
  39
  40             let align = mem::align_of::<T>();
  41
  42             let (new_cap, ptr) = if self.cap == 0 {
  43                 let ptr = heap::allocate(elem_size, align);
  44                 (1, ptr)
  45             } else {
  46                 let new_cap = 2 * self.cap;
  47                 let ptr = heap::reallocate(*self.ptr as *mut _,
  48                                             self.cap * elem_size,
  49                                             new_cap * elem_size,
  50                                             align);
  51                 (new_cap, ptr)
  52             };
  53
  54             // If allocate or reallocate fail, we'll get `null` back
  55             if ptr.is_null() { oom() }
  56
  57             self.ptr = Unique::new(ptr as *mut _);
  58             self.cap = new_cap;
  59         }
  60     }
  61 }
  62
  63 impl<T> Drop for RawVec<T> {
  64     fn drop(&mut self) {
  65         let elem_size = mem::size_of::<T>();
  66         if self.cap != 0 && elem_size != 0 {
  67             let align = mem::align_of::<T>();
  68
  69             let num_bytes = elem_size * self.cap;
  70             unsafe {
  71                 heap::deallocate(*self.ptr as *mut _, num_bytes, align);
  72             }
  73         }
  74     }
  75 }
  76
  77
  78
  79
  80
  81 pub struct Vec<T> {
  82     buf: RawVec<T>,
  83     len: usize,
  84 }
  85
  86 impl<T> Vec<T> {
  87     fn ptr(&self) -> *mut T { *self.buf.ptr }
  88
  89     fn cap(&self) -> usize { self.buf.cap }
  90
  91     pub fn new() -> Self {
  92         Vec { buf: RawVec::new(), len: 0 }
  93     }
  94     pub fn push(&mut self, elem: T) {
  95         if self.len == self.cap() { self.buf.grow(); }
  96
  97         unsafe {
  98             ptr::write(self.ptr().offset(self.len as isize), elem);
  99         }
 100
 101         // Can't fail, we'll OOM first.
 102         self.len += 1;
 103     }
 104
 105     pub fn pop(&mut self) -> Option<T> {
 106         if self.len == 0 {
 107             None
 108         } else {
 109             self.len -= 1;
 110             unsafe {
 111                 Some(ptr::read(self.ptr().offset(self.len as isize)))
 112             }
 113         }
 114     }
 115
 116     pub fn insert(&mut self, index: usize, elem: T) {
 117         assert!(index <= self.len, "index out of bounds");
 118         if self.cap() == self.len { self.buf.grow(); }
 119
 120         unsafe {
 121             if index < self.len {
 122                 ptr::copy(self.ptr().offset(index as isize),
 123                           self.ptr().offset(index as isize + 1),
 124                           self.len - index);
 125             }
 126             ptr::write(self.ptr().offset(index as isize), elem);
 127             self.len += 1;
 128         }
 129     }
 130
 131     pub fn remove(&mut self, index: usize) -> T {
 132         assert!(index < self.len, "index out of bounds");
 133         unsafe {
 134             self.len -= 1;
 135             let result = ptr::read(self.ptr().offset(index as isize));
 136             ptr::copy(self.ptr().offset(index as isize + 1),
 137                       self.ptr().offset(index as isize),
 138                       self.len - index);
 139             result
 140         }
 141     }
 142
 143     pub fn into_iter(self) -> IntoIter<T> {
 144         unsafe {
 145             let iter = RawValIter::new(&self);
 146             let buf = ptr::read(&self.buf);
 147             mem::forget(self);
 148
 149             IntoIter {
 150                 iter: iter,
 151                 _buf: buf,
 152             }
 153         }
 154     }
 155
 156     pub fn drain(&mut self) -> Drain<T> {
 157         unsafe {
 158             let iter = RawValIter::new(&self);
 159
 160             // this is a mem::forget safety thing. If Drain is forgotten, we just
 161             // leak the whole Vec's contents. Also we need to do this *eventually*
 162             // anyway, so why not do it now?
 163             self.len = 0;
 164
 165             Drain {
 166                 iter: iter,
 167                 vec: PhantomData,
 168             }
 169         }
 170     }
 171 }
 172
 173 impl<T> Drop for Vec<T> {
 174     fn drop(&mut self) {
 175         while let Some(_) = self.pop() {}
 176         // allocation is handled by RawVec
 177     }
 178 }
 179
 180 impl<T> Deref for Vec<T> {
 181     type Target = [T];
 182     fn deref(&self) -> &[T] {
 183         unsafe {
 184             ::std::slice::from_raw_parts(self.ptr(), self.len)
 185         }
 186     }
 187 }
 188
 189 impl<T> DerefMut for Vec<T> {
 190     fn deref_mut(&mut self) -> &mut [T] {
 191         unsafe {
 192             ::std::slice::from_raw_parts_mut(self.ptr(), self.len)
 193         }
 194     }
 195 }
 196
 197
 198
 199
 200
 201 struct RawValIter<T> {
 202     start: *const T,
 203     end: *const T,
 204 }
 205
 206 impl<T> RawValIter<T> {
 207     unsafe fn new(slice: &[T]) -> Self {
 208         RawValIter {
 209             start: slice.as_ptr(),
 210             end: if mem::size_of::<T>() == 0 {
 211                 ((slice.as_ptr() as usize) + slice.len()) as *const _
 212             } else if slice.len() == 0 {
 213                 slice.as_ptr()
 214             } else {
 215                 slice.as_ptr().offset(slice.len() as isize)
 216             }
 217         }
 218     }
 219 }
 220
 221 impl<T> Iterator for RawValIter<T> {
 222     type Item = T;
 223     fn next(&mut self) -> Option<T> {
 224         if self.start == self.end {
 225             None
 226         } else {
 227             unsafe {
 228                 let result = ptr::read(self.start);
 229                 self.start = if mem::size_of::<T>() == 0 {
 230                     (self.start as usize + 1) as *const _
 231                 } else {
 232                     self.start.offset(1)
 233                 };
 234                 Some(result)
 235             }
 236         }
 237     }
 238
 239     fn size_hint(&self) -> (usize, Option<usize>) {
 240         let elem_size = mem::size_of::<T>();
 241         let len = (self.end as usize - self.start as usize)
 242                   / if elem_size == 0 { 1 } else { elem_size };
 243         (len, Some(len))
 244     }
 245 }
 246
 247 impl<T> DoubleEndedIterator for RawValIter<T> {
 248     fn next_back(&mut self) -> Option<T> {
 249         if self.start == self.end {
 250             None
 251         } else {
 252             unsafe {
 253                 self.end = if mem::size_of::<T>() == 0 {
 254                     (self.end as usize - 1) as *const _
 255                 } else {
 256                     self.end.offset(-1)
 257                 };
 258                 Some(ptr::read(self.end))
 259             }
 260         }
 261     }
 262 }
 263
 264
 265
 266
 267 pub struct IntoIter<T> {
 268     _buf: RawVec<T>, // we don't actually care about this. Just need it to live.
 269     iter: RawValIter<T>,
 270 }
 271
 272 impl<T> Iterator for IntoIter<T> {
 273     type Item = T;
 274     fn next(&mut self) -> Option<T> { self.iter.next() }
 275     fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
 276 }
 277
 278 impl<T> DoubleEndedIterator for IntoIter<T> {
 279     fn next_back(&mut self) -> Option<T> { self.iter.next_back() }
 280 }
 281
 282 impl<T> Drop for IntoIter<T> {
 283     fn drop(&mut self) {
 284         for _ in &mut *self {}
 285     }
 286 }
 287
 288
 289
 290
 291 pub struct Drain<'a, T: 'a> {
 292     vec: PhantomData<&'a mut Vec<T>>,
 293     iter: RawValIter<T>,
 294 }
 295
 296 impl<'a, T> Iterator for Drain<'a, T> {
 297     type Item = T;
 298     fn next(&mut self) -> Option<T> { self.iter.next_back() }
 299     fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
 300 }
 301
 302 impl<'a, T> DoubleEndedIterator for Drain<'a, T> {
 303     fn next_back(&mut self) -> Option<T> { self.iter.next_back() }
 304 }
 305
 306 impl<'a, T> Drop for Drain<'a, T> {
 307     fn drop(&mut self) {
 308         // pre-drain the iter
 309         for _ in &mut self.iter {}
 310     }
 311 }
 312
 313 /// Abort the process, we're out of memory!
 314 ///
 315 /// In practice this is probably dead code on most OSes
 316 fn oom() {
 317     ::std::process::exit(-9999);
 318 }
 319
 320 # fn main() {}
 321 ```