src/librustc_data_structures/small_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 vector type intended to be used for collecting from iterators onto the stack.
  12 //!
  13 //! Space for up to N elements is provided on the stack.  If more elements are collected, Vec is
  14 //! used to store the values on the heap. SmallVec is similar to AccumulateVec, but adds
  15 //! the ability to push elements.
  16 //!
  17 //! The N above is determined by Array's implementor, by way of an associatated constant.
  18
  19 use std::ops::{Deref, DerefMut};
  20 use std::iter::{IntoIterator, FromIterator};
  21 use std::fmt::{self, Debug};
  22 use std::mem;
  23 use std::ptr;
  24
  25 use rustc_serialize::{Encodable, Encoder, Decodable, Decoder};
  26
  27 use accumulate_vec::{IntoIter, AccumulateVec};
  28 use array_vec::Array;
  29
  30 pub struct SmallVec<A: Array>(AccumulateVec<A>);
  31
  32 impl<A> Clone for SmallVec<A>
  33     where A: Array,
  34           A::Element: Clone {
  35     fn clone(&self) -> Self {
  36         SmallVec(self.0.clone())
  37     }
  38 }
  39
  40 impl<A> Debug for SmallVec<A>
  41     where A: Array + Debug,
  42           A::Element: Debug {
  43     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
  44         f.debug_tuple("SmallVec").field(&self.0).finish()
  45     }
  46 }
  47
  48 impl<A: Array> SmallVec<A> {
  49     pub fn new() -> Self {
  50         SmallVec(AccumulateVec::new())
  51     }
  52
  53     pub fn with_capacity(cap: usize) -> Self {
  54         let mut vec = SmallVec::new();
  55         vec.reserve(cap);
  56         vec
  57     }
  58
  59     pub fn one(el: A::Element) -> Self {
  60         SmallVec(AccumulateVec::one(el))
  61     }
  62
  63     pub fn many<I: IntoIterator<Item=A::Element>>(els: I) -> Self {
  64         SmallVec(AccumulateVec::many(els))
  65     }
  66
  67     pub fn expect_one(self, err: &'static str) -> A::Element {
  68         assert!(self.len() == 1, err);
  69         match self.0 {
  70             AccumulateVec::Array(arr) => arr.into_iter().next().unwrap(),
  71             AccumulateVec::Heap(vec) => vec.into_iter().next().unwrap(),
  72         }
  73     }
  74
  75     /// Will reallocate onto the heap if needed.
  76     pub fn push(&mut self, el: A::Element) {
  77         self.reserve(1);
  78         match self.0 {
  79             AccumulateVec::Array(ref mut array) => array.push(el),
  80             AccumulateVec::Heap(ref mut vec) => vec.push(el),
  81         }
  82     }
  83
  84     pub fn reserve(&mut self, n: usize) {
  85         match self.0 {
  86             AccumulateVec::Array(_) => {
  87                 if self.len() + n > A::LEN {
  88                     let len = self.len();
  89                     let array = mem::replace(&mut self.0,
  90                             AccumulateVec::Heap(Vec::with_capacity(len + n)));
  91                     if let AccumulateVec::Array(array) = array {
  92                         match self.0 {
  93                             AccumulateVec::Heap(ref mut vec) => vec.extend(array),
  94                             _ => unreachable!()
  95                         }
  96                     }
  97                 }
  98             }
  99             AccumulateVec::Heap(ref mut vec) => vec.reserve(n)
 100         }
 101     }
 102
 103     pub unsafe fn set_len(&mut self, len: usize) {
 104         match self.0 {
 105             AccumulateVec::Array(ref mut arr) => arr.set_len(len),
 106             AccumulateVec::Heap(ref mut vec) => vec.set_len(len),
 107         }
 108     }
 109
 110     pub fn insert(&mut self, index: usize, element: A::Element) {
 111         let len = self.len();
 112
 113         // Reserve space for shifting elements to the right
 114         self.reserve(1);
 115
 116         assert!(index <= len);
 117
 118         unsafe {
 119             // infallible
 120             // The spot to put the new value
 121             {
 122                 let p = self.as_mut_ptr().offset(index as isize);
 123                 // Shift everything over to make space. (Duplicating the
 124                 // `index`th element into two consecutive places.)
 125                 ptr::copy(p, p.offset(1), len - index);
 126                 // Write it in, overwriting the first copy of the `index`th
 127                 // element.
 128                 ptr::write(p, element);
 129             }
 130             self.set_len(len + 1);
 131         }
 132     }
 133
 134     pub fn truncate(&mut self, len: usize) {
 135         unsafe {
 136             while len < self.len() {
 137                 // Decrement len before the drop_in_place(), so a panic on Drop
 138                 // doesn't re-drop the just-failed value.
 139                 let newlen = self.len() - 1;
 140                 self.set_len(newlen);
 141                 ::std::ptr::drop_in_place(self.get_unchecked_mut(newlen));
 142             }
 143         }
 144     }
 145 }
 146
 147 impl<A: Array> Deref for SmallVec<A> {
 148     type Target = AccumulateVec<A>;
 149     fn deref(&self) -> &Self::Target {
 150         &self.0
 151     }
 152 }
 153
 154 impl<A: Array> DerefMut for SmallVec<A> {
 155     fn deref_mut(&mut self) -> &mut AccumulateVec<A> {
 156         &mut self.0
 157     }
 158 }
 159
 160 impl<A: Array> FromIterator<A::Element> for SmallVec<A> {
 161     fn from_iter<I>(iter: I) -> Self where I: IntoIterator<Item=A::Element> {
 162         SmallVec(iter.into_iter().collect())
 163     }
 164 }
 165
 166 impl<A: Array> Extend<A::Element> for SmallVec<A> {
 167     fn extend<I: IntoIterator<Item=A::Element>>(&mut self, iter: I) {
 168         let iter = iter.into_iter();
 169         self.reserve(iter.size_hint().0);
 170         for el in iter {
 171             self.push(el);
 172         }
 173     }
 174 }
 175
 176 impl<A: Array> IntoIterator for SmallVec<A> {
 177     type Item = A::Element;
 178     type IntoIter = IntoIter<A>;
 179     fn into_iter(self) -> Self::IntoIter {
 180         self.0.into_iter()
 181     }
 182 }
 183
 184 impl<A: Array> Default for SmallVec<A> {
 185     fn default() -> SmallVec<A> {
 186         SmallVec::new()
 187     }
 188 }
 189
 190 impl<A> Encodable for SmallVec<A>
 191     where A: Array,
 192           A::Element: Encodable {
 193     fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
 194         s.emit_seq(self.len(), |s| {
 195             for (i, e) in self.iter().enumerate() {
 196                 try!(s.emit_seq_elt(i, |s| e.encode(s)));
 197             }
 198             Ok(())
 199         })
 200     }
 201 }
 202
 203 impl<A> Decodable for SmallVec<A>
 204     where A: Array,
 205           A::Element: Decodable {
 206     fn decode<D: Decoder>(d: &mut D) -> Result<SmallVec<A>, D::Error> {
 207         d.read_seq(|d, len| {
 208             let mut vec = SmallVec::with_capacity(len);
 209             for i in 0..len {
 210                 vec.push(try!(d.read_seq_elt(i, |d| Decodable::decode(d))));
 211             }
 212             Ok(vec)
 213         })
 214     }
 215 }