src/libextra/enum_set.rs

   1 // Copyright 2012 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 structure for holding a set of enum variants
  12 //!
  13 //! This module defines a container which uses an efficient bit mask
  14 //! representation to hold C-like enum variants.
  15
  16 #[deriving(Clone, Eq, IterBytes, ToStr, Encodable, Decodable)]
  17 /// A specialized Set implementation to use enum types.
  18 pub struct EnumSet<E> {
  19     // We must maintain the invariant that no bits are set
  20     // for which no variant exists
  21     priv bits: uint
  22 }
  23
  24 /// An interface for casting C-like enum to uint and back.
  25 pub trait CLike {
  26     /// Converts C-like enum to uint.
  27     fn to_uint(&self) -> uint;
  28     /// Converts uint to C-like enum.
  29     fn from_uint(uint) -> Self;
  30 }
  31
  32 fn bit<E:CLike>(e: E) -> uint {
  33     1 << e.to_uint()
  34 }
  35
  36 impl<E:CLike> EnumSet<E> {
  37     /// Returns an empty EnumSet.
  38     pub fn empty() -> EnumSet<E> {
  39         EnumSet {bits: 0}
  40     }
  41
  42     /// Returns true if an EnumSet is empty.
  43     pub fn is_empty(&self) -> bool {
  44         self.bits == 0
  45     }
  46
  47     /// Returns true if an EnumSet contains any enum of a given EnumSet
  48     pub fn intersects(&self, e: EnumSet<E>) -> bool {
  49         (self.bits & e.bits) != 0
  50     }
  51
  52     /// Returns an intersection of both EnumSets.
  53     pub fn intersection(&self, e: EnumSet<E>) -> EnumSet<E> {
  54         EnumSet {bits: self.bits & e.bits}
  55     }
  56
  57     /// Returns true if a given EnumSet is included in an EnumSet.
  58     pub fn contains(&self, e: EnumSet<E>) -> bool {
  59         (self.bits & e.bits) == e.bits
  60     }
  61
  62     /// Returns a union of both EnumSets.
  63     pub fn union(&self, e: EnumSet<E>) -> EnumSet<E> {
  64         EnumSet {bits: self.bits | e.bits}
  65     }
  66
  67     /// Add an enum to an EnumSet
  68     pub fn add(&mut self, e: E) {
  69         self.bits |= bit(e);
  70     }
  71
  72     /// Returns true if an EnumSet contains a given enum
  73     pub fn contains_elem(&self, e: E) -> bool {
  74         (self.bits & bit(e)) != 0
  75     }
  76
  77     /// Returns an iterator over an EnumSet
  78     pub fn iter(&self) -> EnumSetIterator<E> {
  79         EnumSetIterator::new(self.bits)
  80     }
  81 }
  82
  83 impl<E:CLike> Sub<EnumSet<E>, EnumSet<E>> for EnumSet<E> {
  84     fn sub(&self, e: &EnumSet<E>) -> EnumSet<E> {
  85         EnumSet {bits: self.bits & !e.bits}
  86     }
  87 }
  88
  89 impl<E:CLike> BitOr<EnumSet<E>, EnumSet<E>> for EnumSet<E> {
  90     fn bitor(&self, e: &EnumSet<E>) -> EnumSet<E> {
  91         EnumSet {bits: self.bits | e.bits}
  92     }
  93 }
  94
  95 impl<E:CLike> BitAnd<EnumSet<E>, EnumSet<E>> for EnumSet<E> {
  96     fn bitand(&self, e: &EnumSet<E>) -> EnumSet<E> {
  97         EnumSet {bits: self.bits & e.bits}
  98     }
  99 }
 100
 101 /// An iterator over an EnumSet
 102 pub struct EnumSetIterator<E> {
 103     priv index: uint,
 104     priv bits: uint,
 105 }
 106
 107 impl<E:CLike> EnumSetIterator<E> {
 108     fn new(bits: uint) -> EnumSetIterator<E> {
 109         EnumSetIterator { index: 0, bits: bits }
 110     }
 111 }
 112
 113 impl<E:CLike> Iterator<E> for EnumSetIterator<E> {
 114     fn next(&mut self) -> Option<E> {
 115         if (self.bits == 0) {
 116             return None;
 117         }
 118
 119         while (self.bits & 1) == 0 {
 120             self.index += 1;
 121             self.bits >>= 1;
 122         }
 123         let elem = CLike::from_uint(self.index);
 124         self.index += 1;
 125         self.bits >>= 1;
 126         Some(elem)
 127     }
 128
 129     fn size_hint(&self) -> (uint, Option<uint>) {
 130         let exact = self.bits.population_count();
 131         (exact, Some(exact))
 132     }
 133 }
 134
 135 #[cfg(test)]
 136 mod test {
 137
 138     use std::cast;
 139
 140     use enum_set::*;
 141
 142     #[deriving(Eq)]
 143     #[repr(uint)]
 144     enum Foo {
 145         A, B, C
 146     }
 147
 148     impl CLike for Foo {
 149         fn to_uint(&self) -> uint {
 150             *self as uint
 151         }
 152
 153         fn from_uint(v: uint) -> Foo {
 154             unsafe { cast::transmute(v) }
 155         }
 156     }
 157
 158     #[test]
 159     fn test_empty() {
 160         let e: EnumSet<Foo> = EnumSet::empty();
 161         assert!(e.is_empty());
 162     }
 163
 164     ///////////////////////////////////////////////////////////////////////////
 165     // intersect
 166
 167     #[test]
 168     fn test_two_empties_do_not_intersect() {
 169         let e1: EnumSet<Foo> = EnumSet::empty();
 170         let e2: EnumSet<Foo> = EnumSet::empty();
 171         assert!(!e1.intersects(e2));
 172     }
 173
 174     #[test]
 175     fn test_empty_does_not_intersect_with_full() {
 176         let e1: EnumSet<Foo> = EnumSet::empty();
 177
 178         let mut e2: EnumSet<Foo> = EnumSet::empty();
 179         e2.add(A);
 180         e2.add(B);
 181         e2.add(C);
 182
 183         assert!(!e1.intersects(e2));
 184     }
 185
 186     #[test]
 187     fn test_disjoint_intersects() {
 188         let mut e1: EnumSet<Foo> = EnumSet::empty();
 189         e1.add(A);
 190
 191         let mut e2: EnumSet<Foo> = EnumSet::empty();
 192         e2.add(B);
 193
 194         assert!(!e1.intersects(e2));
 195     }
 196
 197     #[test]
 198     fn test_overlapping_intersects() {
 199         let mut e1: EnumSet<Foo> = EnumSet::empty();
 200         e1.add(A);
 201
 202         let mut e2: EnumSet<Foo> = EnumSet::empty();
 203         e2.add(A);
 204         e2.add(B);
 205
 206         assert!(e1.intersects(e2));
 207     }
 208
 209     ///////////////////////////////////////////////////////////////////////////
 210     // contains and contains_elem
 211
 212     #[test]
 213     fn test_contains() {
 214         let mut e1: EnumSet<Foo> = EnumSet::empty();
 215         e1.add(A);
 216
 217         let mut e2: EnumSet<Foo> = EnumSet::empty();
 218         e2.add(A);
 219         e2.add(B);
 220
 221         assert!(!e1.contains(e2));
 222         assert!(e2.contains(e1));
 223     }
 224
 225     #[test]
 226     fn test_contains_elem() {
 227         let mut e1: EnumSet<Foo> = EnumSet::empty();
 228         e1.add(A);
 229         assert!(e1.contains_elem(A));
 230         assert!(!e1.contains_elem(B));
 231         assert!(!e1.contains_elem(C));
 232
 233         e1.add(A);
 234         e1.add(B);
 235         assert!(e1.contains_elem(A));
 236         assert!(e1.contains_elem(B));
 237         assert!(!e1.contains_elem(C));
 238     }
 239
 240     ///////////////////////////////////////////////////////////////////////////
 241     // iter
 242
 243     #[test]
 244     fn test_iterator() {
 245         let mut e1: EnumSet<Foo> = EnumSet::empty();
 246
 247         let elems: ~[Foo] = e1.iter().collect();
 248         assert_eq!(~[], elems)
 249
 250         e1.add(A);
 251         let elems: ~[Foo] = e1.iter().collect();
 252         assert_eq!(~[A], elems)
 253
 254         e1.add(C);
 255         let elems: ~[Foo] = e1.iter().collect();
 256         assert_eq!(~[A,C], elems)
 257
 258         e1.add(C);
 259         let elems: ~[Foo] = e1.iter().collect();
 260         assert_eq!(~[A,C], elems)
 261
 262         e1.add(B);
 263         let elems: ~[Foo] = e1.iter().collect();
 264         assert_eq!(~[A,B,C], elems)
 265     }
 266
 267     ///////////////////////////////////////////////////////////////////////////
 268     // operators
 269
 270     #[test]
 271     fn test_operators() {
 272         let mut e1: EnumSet<Foo> = EnumSet::empty();
 273         e1.add(A);
 274         e1.add(C);
 275
 276         let mut e2: EnumSet<Foo> = EnumSet::empty();
 277         e2.add(B);
 278         e2.add(C);
 279
 280         let e_union = e1 | e2;
 281         let elems: ~[Foo] = e_union.iter().collect();
 282         assert_eq!(~[A,B,C], elems)
 283
 284         let e_intersection = e1 & e2;
 285         let elems: ~[Foo] = e_intersection.iter().collect();
 286         assert_eq!(~[C], elems)
 287
 288         let e_subtract = e1 - e2;
 289         let elems: ~[Foo] = e_subtract.iter().collect();
 290         assert_eq!(~[A], elems)
 291     }
 292 }