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