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