enumiter/
lib.rs

1use std::mem::{transmute_copy, size_of};
2
3/// May only be safely implemented for enums with:
4///
5///   1. No variant parameters whatsoever.
6///   2. No index assignment. For example: `enum T { A = 1, B = 2 }`.
7pub unsafe trait AllVariantsTakeNoParameters : Eq  + Copy {}
8
9pub struct EnumIter<T: AllVariantsTakeNoParameters> {
10  cur : T,
11  end : T,
12  done: bool,
13}
14
15impl<T: AllVariantsTakeNoParameters> EnumIter<T> {
16  #[inline]
17  pub fn new(start: T, end: T) -> EnumIter<T> {
18    assert!(size_of::<T>() == 1
19         || size_of::<T>() == 2
20         || size_of::<T>() == 4
21         || size_of::<T>() == 8);
22
23    EnumIter {
24      cur:  start,
25      end:  end,
26      done: false,
27    }
28  }
29}
30
31macro_rules! next {
32  ($ty : ty, $this: expr) => ({
33    let ret = $this.cur;
34    let cur: $ty = transmute_copy(&ret);
35    let next = cur + 1;
36    let next = transmute_copy(&next);
37    $this.cur = next;
38    ret
39  })
40}
41
42impl<T: AllVariantsTakeNoParameters> Iterator for EnumIter<T> {
43  type Item = T;
44
45  #[inline]
46  fn next(&mut self) -> Option<T> {
47    if self.done {
48      None
49    } else {
50      let ret = unsafe {
51        match size_of::<T>() {
52          1 => next!(u8 , self),
53          2 => next!(u16, self),
54          4 => next!(u32, self),
55          8 => next!(u64, self),
56          _ => unreachable!()
57        }
58      };
59      self.done = ret == self.end;
60      Some(ret)
61    }
62  }
63}
64
65/// `enum_iter(Enum::FirstElement, Enum::LastElement)` returns an iterator
66/// through the different variants of the enum. `T` must satisfy
67/// `AllVariantsTakeNoParameters`, which pretty much just means you built a
68/// simple rust enum with no variant parameters and no manual assignment of
69/// integer mappings.
70pub fn enum_iter<T: AllVariantsTakeNoParameters>(start: T, end: T) -> EnumIter<T> {
71  EnumIter::new(start, end)
72}
73
74#[cfg(test)]
75mod test {
76  use super::{enum_iter, AllVariantsTakeNoParameters};
77
78  #[derive(Clone, Copy, PartialEq, Eq, Debug)]
79  enum Test {
80    One,
81    Two,
82    Three,
83  }
84
85  unsafe impl AllVariantsTakeNoParameters for Test {}
86
87  #[test]
88  fn it_works() {
89    let mut vals = Vec::new();
90
91    let _ = Test::Two; // ignore dead code warning.
92
93    for x in enum_iter(Test::One, Test::Three) {
94      vals.push(x);
95    }
96
97    assert_eq!(format!("{:?}", vals), "[One, Two, Three]");
98  }
99}