src/libsyntax/util/small_vector.rs

   1 // Copyright 2013-2014 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 use self::SmallVectorRepr::*;
  12 use self::IntoIterRepr::*;
  13
  14 use std::iter::{IntoIterator, FromIterator};
  15 use std::mem;
  16 use std::slice;
  17 use std::vec;
  18
  19 use fold::MoveMap;
  20
  21 /// A vector type optimized for cases where the size is almost always 0 or 1
  22 pub struct SmallVector<T> {
  23     repr: SmallVectorRepr<T>,
  24 }
  25
  26 enum SmallVectorRepr<T> {
  27     Zero,
  28     One(T),
  29     Many(Vec<T>),
  30 }
  31
  32 impl<T> FromIterator<T> for SmallVector<T> {
  33     fn from_iter<I: IntoIterator<Item=T>>(iter: I) -> SmallVector<T> {
  34         let mut v = SmallVector::zero();
  35         v.extend(iter);
  36         v
  37     }
  38 }
  39
  40 impl<T> Extend<T> for SmallVector<T> {
  41     fn extend<I: IntoIterator<Item=T>>(&mut self, iter: I) {
  42         for val in iter {
  43             self.push(val);
  44         }
  45     }
  46 }
  47
  48 impl<T> SmallVector<T> {
  49     pub fn zero() -> SmallVector<T> {
  50         SmallVector { repr: Zero }
  51     }
  52
  53     pub fn one(v: T) -> SmallVector<T> {
  54         SmallVector { repr: One(v) }
  55     }
  56
  57     pub fn many(vs: Vec<T>) -> SmallVector<T> {
  58         SmallVector { repr: Many(vs) }
  59     }
  60
  61     pub fn as_slice<'a>(&'a self) -> &'a [T] {
  62         match self.repr {
  63             Zero => {
  64                 let result: &[T] = &[];
  65                 result
  66             }
  67             One(ref v) => {
  68                 unsafe { slice::from_raw_parts(v, 1) }
  69             }
  70             Many(ref vs) => vs
  71         }
  72     }
  73
  74     pub fn pop(&mut self) -> Option<T> {
  75         match self.repr {
  76             Zero => None,
  77             One(..) => {
  78                 let one = mem::replace(&mut self.repr, Zero);
  79                 match one {
  80                     One(v1) => Some(v1),
  81                     _ => unreachable!()
  82                 }
  83             }
  84             Many(ref mut vs) => vs.pop(),
  85         }
  86     }
  87
  88     pub fn push(&mut self, v: T) {
  89         match self.repr {
  90             Zero => self.repr = One(v),
  91             One(..) => {
  92                 let one = mem::replace(&mut self.repr, Zero);
  93                 match one {
  94                     One(v1) => mem::replace(&mut self.repr, Many(vec!(v1, v))),
  95                     _ => unreachable!()
  96                 };
  97             }
  98             Many(ref mut vs) => vs.push(v)
  99         }
 100     }
 101
 102     pub fn push_all(&mut self, other: SmallVector<T>) {
 103         for v in other.into_iter() {
 104             self.push(v);
 105         }
 106     }
 107
 108     pub fn get<'a>(&'a self, idx: usize) -> &'a T {
 109         match self.repr {
 110             One(ref v) if idx == 0 => v,
 111             Many(ref vs) => &vs[idx],
 112             _ => panic!("out of bounds access")
 113         }
 114     }
 115
 116     pub fn expect_one(self, err: &'static str) -> T {
 117         match self.repr {
 118             One(v) => v,
 119             Many(v) => {
 120                 if v.len() == 1 {
 121                     v.into_iter().next().unwrap()
 122                 } else {
 123                     panic!(err)
 124                 }
 125             }
 126             _ => panic!(err)
 127         }
 128     }
 129
 130     /// Deprecated: use `into_iter`.
 131     #[unstable(feature = "rustc_private", issue = "0")]
 132     #[deprecated(since = "1.0.0", reason = "use into_iter")]
 133     pub fn move_iter(self) -> IntoIter<T> {
 134         self.into_iter()
 135     }
 136
 137     pub fn into_iter(self) -> IntoIter<T> {
 138         let repr = match self.repr {
 139             Zero => ZeroIterator,
 140             One(v) => OneIterator(v),
 141             Many(vs) => ManyIterator(vs.into_iter())
 142         };
 143         IntoIter { repr: repr }
 144     }
 145
 146     pub fn len(&self) -> usize {
 147         match self.repr {
 148             Zero => 0,
 149             One(..) => 1,
 150             Many(ref vals) => vals.len()
 151         }
 152     }
 153
 154     pub fn is_empty(&self) -> bool { self.len() == 0 }
 155 }
 156
 157 pub struct IntoIter<T> {
 158     repr: IntoIterRepr<T>,
 159 }
 160
 161 enum IntoIterRepr<T> {
 162     ZeroIterator,
 163     OneIterator(T),
 164     ManyIterator(vec::IntoIter<T>),
 165 }
 166
 167 impl<T> Iterator for IntoIter<T> {
 168     type Item = T;
 169
 170     fn next(&mut self) -> Option<T> {
 171         match self.repr {
 172             ZeroIterator => None,
 173             OneIterator(..) => {
 174                 let mut replacement = ZeroIterator;
 175                 mem::swap(&mut self.repr, &mut replacement);
 176                 match replacement {
 177                     OneIterator(v) => Some(v),
 178                     _ => unreachable!()
 179                 }
 180             }
 181             ManyIterator(ref mut inner) => inner.next()
 182         }
 183     }
 184
 185     fn size_hint(&self) -> (usize, Option<usize>) {
 186         match self.repr {
 187             ZeroIterator => (0, Some(0)),
 188             OneIterator(..) => (1, Some(1)),
 189             ManyIterator(ref inner) => inner.size_hint()
 190         }
 191     }
 192 }
 193
 194 impl<T> MoveMap<T> for SmallVector<T> {
 195     fn move_map<F>(self, mut f: F) -> SmallVector<T> where F: FnMut(T) -> T {
 196         let repr = match self.repr {
 197             Zero => Zero,
 198             One(v) => One(f(v)),
 199             Many(vs) => Many(vs.move_map(f))
 200         };
 201         SmallVector { repr: repr }
 202     }
 203 }
 204
 205 #[cfg(test)]
 206 mod tests {
 207     use super::*;
 208
 209     #[test]
 210     fn test_len() {
 211         let v: SmallVector<isize> = SmallVector::zero();
 212         assert_eq!(0, v.len());
 213
 214         assert_eq!(1, SmallVector::one(1).len());
 215         assert_eq!(5, SmallVector::many(vec![1, 2, 3, 4, 5]).len());
 216     }
 217
 218     #[test]
 219     fn test_push_get() {
 220         let mut v = SmallVector::zero();
 221         v.push(1);
 222         assert_eq!(1, v.len());
 223         assert_eq!(&1, v.get(0));
 224         v.push(2);
 225         assert_eq!(2, v.len());
 226         assert_eq!(&2, v.get(1));
 227         v.push(3);
 228         assert_eq!(3, v.len());
 229         assert_eq!(&3, v.get(2));
 230     }
 231
 232     #[test]
 233     fn test_from_iter() {
 234         let v: SmallVector<isize> = (vec![1, 2, 3]).into_iter().collect();
 235         assert_eq!(3, v.len());
 236         assert_eq!(&1, v.get(0));
 237         assert_eq!(&2, v.get(1));
 238         assert_eq!(&3, v.get(2));
 239     }
 240
 241     #[test]
 242     fn test_move_iter() {
 243         let v = SmallVector::zero();
 244         let v: Vec<isize> = v.into_iter().collect();
 245         assert_eq!(v, Vec::new());
 246
 247         let v = SmallVector::one(1);
 248         assert_eq!(v.into_iter().collect::<Vec<_>>(), [1]);
 249
 250         let v = SmallVector::many(vec![1, 2, 3]);
 251         assert_eq!(v.into_iter().collect::<Vec<_>>(), [1, 2, 3]);
 252     }
 253
 254     #[test]
 255     #[should_panic]
 256     fn test_expect_one_zero() {
 257         let _: isize = SmallVector::zero().expect_one("");
 258     }
 259
 260     #[test]
 261     #[should_panic]
 262     fn test_expect_one_many() {
 263         SmallVector::many(vec!(1, 2)).expect_one("");
 264     }
 265
 266     #[test]
 267     fn test_expect_one_one() {
 268         assert_eq!(1, SmallVector::one(1).expect_one(""));
 269         assert_eq!(1, SmallVector::many(vec!(1)).expect_one(""));
 270     }
 271 }