src/librustc_data_structures/array_vec.rs

   1 // Copyright 2016 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 //! A stack-allocated vector, allowing storage of N elements on the stack.
  12
  13 use std::marker::Unsize;
  14 use std::iter::Extend;
  15 use std::ptr::{self, drop_in_place, Shared};
  16 use std::ops::{Deref, DerefMut, Range};
  17 use std::hash::{Hash, Hasher};
  18 use std::slice;
  19 use std::fmt;
  20 use std::mem;
  21 use std::collections::range::RangeArgument;
  22 use std::collections::Bound::{Excluded, Included, Unbounded};
  23 use std::mem::ManuallyDrop;
  24
  25 pub unsafe trait Array {
  26     type Element;
  27     type PartialStorage: Unsize<[ManuallyDrop<Self::Element>]>;
  28     const LEN: usize;
  29 }
  30
  31 unsafe impl<T> Array for [T; 1] {
  32     type Element = T;
  33     type PartialStorage = [ManuallyDrop<T>; 1];
  34     const LEN: usize = 1;
  35 }
  36
  37 unsafe impl<T> Array for [T; 8] {
  38     type Element = T;
  39     type PartialStorage = [ManuallyDrop<T>; 8];
  40     const LEN: usize = 8;
  41 }
  42
  43 unsafe impl<T> Array for [T; 32] {
  44     type Element = T;
  45     type PartialStorage = [ManuallyDrop<T>; 32];
  46     const LEN: usize = 32;
  47 }
  48
  49 pub struct ArrayVec<A: Array> {
  50     count: usize,
  51     values: A::PartialStorage
  52 }
  53
  54 impl<A> Hash for ArrayVec<A>
  55     where A: Array,
  56           A::Element: Hash {
  57     fn hash<H>(&self, state: &mut H) where H: Hasher {
  58         (&self[..]).hash(state);
  59     }
  60 }
  61
  62 impl<A> Clone for ArrayVec<A>
  63     where A: Array,
  64           A::Element: Clone {
  65     fn clone(&self) -> Self {
  66         let mut v = ArrayVec::new();
  67         v.extend(self.iter().cloned());
  68         v
  69     }
  70 }
  71
  72 impl<A: Array> ArrayVec<A> {
  73     pub fn new() -> Self {
  74         ArrayVec {
  75             count: 0,
  76             values: unsafe { ::std::mem::uninitialized() },
  77         }
  78     }
  79
  80     pub fn len(&self) -> usize {
  81         self.count
  82     }
  83
  84     pub unsafe fn set_len(&mut self, len: usize) {
  85         self.count = len;
  86     }
  87
  88     /// Panics when the stack vector is full.
  89     pub fn push(&mut self, el: A::Element) {
  90         let arr = &mut self.values as &mut [ManuallyDrop<_>];
  91         arr[self.count] = ManuallyDrop::new(el);
  92         self.count += 1;
  93     }
  94
  95     pub fn pop(&mut self) -> Option<A::Element> {
  96         if self.count > 0 {
  97             let arr = &mut self.values as &mut [ManuallyDrop<_>];
  98             self.count -= 1;
  99             unsafe {
 100                 let value = ptr::read(&*arr[self.count]);
 101                 Some(value)
 102             }
 103         } else {
 104             None
 105         }
 106     }
 107
 108     pub fn drain<R>(&mut self, range: R) -> Drain<A>
 109         where R: RangeArgument<usize>
 110     {
 111         // Memory safety
 112         //
 113         // When the Drain is first created, it shortens the length of
 114         // the source vector to make sure no uninitalized or moved-from elements
 115         // are accessible at all if the Drain's destructor never gets to run.
 116         //
 117         // Drain will ptr::read out the values to remove.
 118         // When finished, remaining tail of the vec is copied back to cover
 119         // the hole, and the vector length is restored to the new length.
 120         //
 121         let len = self.len();
 122         let start = match range.start() {
 123             Included(&n) => n,
 124             Excluded(&n) => n + 1,
 125             Unbounded    => 0,
 126         };
 127         let end = match range.end() {
 128             Included(&n) => n + 1,
 129             Excluded(&n) => n,
 130             Unbounded    => len,
 131         };
 132         assert!(start <= end);
 133         assert!(end <= len);
 134
 135         unsafe {
 136             // set self.vec length's to start, to be safe in case Drain is leaked
 137             self.set_len(start);
 138             // Use the borrow in the IterMut to indicate borrowing behavior of the
 139             // whole Drain iterator (like &mut T).
 140             let range_slice = {
 141                 let arr = &mut self.values as &mut [ManuallyDrop<_>];
 142                 slice::from_raw_parts_mut(arr.as_mut_ptr().offset(start as isize),
 143                                           end - start)
 144             };
 145             Drain {
 146                 tail_start: end,
 147                 tail_len: len - end,
 148                 iter: range_slice.iter(),
 149                 array_vec: Shared::new(self as *mut _),
 150             }
 151         }
 152     }
 153 }
 154
 155 impl<A> Default for ArrayVec<A>
 156     where A: Array {
 157     fn default() -> Self {
 158         ArrayVec::new()
 159     }
 160 }
 161
 162 impl<A> fmt::Debug for ArrayVec<A>
 163     where A: Array,
 164           A::Element: fmt::Debug {
 165     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 166         self[..].fmt(f)
 167     }
 168 }
 169
 170 impl<A: Array> Deref for ArrayVec<A> {
 171     type Target = [A::Element];
 172     fn deref(&self) -> &Self::Target {
 173         unsafe {
 174             slice::from_raw_parts(&self.values as *const _ as *const A::Element, self.count)
 175         }
 176     }
 177 }
 178
 179 impl<A: Array> DerefMut for ArrayVec<A> {
 180     fn deref_mut(&mut self) -> &mut [A::Element] {
 181         unsafe {
 182             slice::from_raw_parts_mut(&mut self.values as *mut _ as *mut A::Element, self.count)
 183         }
 184     }
 185 }
 186
 187 impl<A: Array> Drop for ArrayVec<A> {
 188     fn drop(&mut self) {
 189         unsafe {
 190             drop_in_place(&mut self[..])
 191         }
 192     }
 193 }
 194
 195 impl<A: Array> Extend<A::Element> for ArrayVec<A> {
 196     fn extend<I>(&mut self, iter: I) where I: IntoIterator<Item=A::Element> {
 197         for el in iter {
 198             self.push(el);
 199         }
 200     }
 201 }
 202
 203 pub struct Iter<A: Array> {
 204     indices: Range<usize>,
 205     store: A::PartialStorage,
 206 }
 207
 208 impl<A: Array> Drop for Iter<A> {
 209     fn drop(&mut self) {
 210         for _ in self {}
 211     }
 212 }
 213
 214 impl<A: Array> Iterator for Iter<A> {
 215     type Item = A::Element;
 216
 217     fn next(&mut self) -> Option<A::Element> {
 218         let arr = &self.store as &[ManuallyDrop<_>];
 219         unsafe {
 220             self.indices.next().map(|i| ptr::read(&*arr[i]))
 221         }
 222     }
 223
 224     fn size_hint(&self) -> (usize, Option<usize>) {
 225         self.indices.size_hint()
 226     }
 227 }
 228
 229 pub struct Drain<'a, A: Array>
 230         where A::Element: 'a
 231 {
 232     tail_start: usize,
 233     tail_len: usize,
 234     iter: slice::Iter<'a, ManuallyDrop<A::Element>>,
 235     array_vec: Shared<ArrayVec<A>>,
 236 }
 237
 238 impl<'a, A: Array> Iterator for Drain<'a, A> {
 239     type Item = A::Element;
 240
 241     #[inline]
 242     fn next(&mut self) -> Option<A::Element> {
 243         self.iter.next().map(|elt| unsafe { ptr::read(&**elt) })
 244     }
 245
 246     fn size_hint(&self) -> (usize, Option<usize>) {
 247         self.iter.size_hint()
 248     }
 249 }
 250
 251 impl<'a, A: Array> Drop for Drain<'a, A> {
 252     fn drop(&mut self) {
 253         // exhaust self first
 254         while let Some(_) = self.next() {}
 255
 256         if self.tail_len > 0 {
 257             unsafe {
 258                 let source_array_vec = &mut *self.array_vec.as_mut_ptr();
 259                 // memmove back untouched tail, update to new length
 260                 let start = source_array_vec.len();
 261                 let tail = self.tail_start;
 262                 {
 263                     let mut arr = &mut source_array_vec.values as &mut [ManuallyDrop<_>];
 264                     let src = arr.as_ptr().offset(tail as isize);
 265                     let dst = arr.as_mut_ptr().offset(start as isize);
 266                     ptr::copy(src, dst, self.tail_len);
 267                 };
 268                 source_array_vec.set_len(start + self.tail_len);
 269             }
 270         }
 271     }
 272 }
 273
 274 impl<A: Array> IntoIterator for ArrayVec<A> {
 275     type Item = A::Element;
 276     type IntoIter = Iter<A>;
 277     fn into_iter(self) -> Self::IntoIter {
 278         let store = unsafe {
 279             ptr::read(&self.values)
 280         };
 281         let indices = 0..self.count;
 282         mem::forget(self);
 283         Iter {
 284             indices: indices,
 285             store: store,
 286         }
 287     }
 288 }
 289
 290 impl<'a, A: Array> IntoIterator for &'a ArrayVec<A> {
 291     type Item = &'a A::Element;
 292     type IntoIter = slice::Iter<'a, A::Element>;
 293     fn into_iter(self) -> Self::IntoIter {
 294         self.iter()
 295     }
 296 }
 297
 298 impl<'a, A: Array> IntoIterator for &'a mut ArrayVec<A> {
 299     type Item = &'a mut A::Element;
 300     type IntoIter = slice::IterMut<'a, A::Element>;
 301     fn into_iter(self) -> Self::IntoIter {
 302         self.iter_mut()
 303     }
 304 }