src/librustc/util/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 use std::iterator::Iterator;
  12
  13 #[deriving(Clone, Eq, IterBytes, ToStr)]
  14 pub struct EnumSet<E> {
  15     // We must maintain the invariant that no bits are set
  16     // for which no variant exists
  17     priv bits: uint
  18 }
  19
  20 pub trait CLike {
  21     pub fn to_uint(&self) -> uint;
  22     pub fn from_uint(uint) -> Self;
  23 }
  24
  25 fn bit<E:CLike>(e: E) -> uint {
  26     1 << e.to_uint()
  27 }
  28
  29 impl<E:CLike> EnumSet<E> {
  30     pub fn empty() -> EnumSet<E> {
  31         EnumSet {bits: 0}
  32     }
  33
  34     pub fn is_empty(&self) -> bool {
  35         self.bits == 0
  36     }
  37
  38     pub fn intersects(&self, e: EnumSet<E>) -> bool {
  39         (self.bits & e.bits) != 0
  40     }
  41
  42     pub fn intersection(&self, e: EnumSet<E>) -> EnumSet<E> {
  43         EnumSet {bits: self.bits & e.bits}
  44     }
  45
  46     pub fn contains(&self, e: EnumSet<E>) -> bool {
  47         (self.bits & e.bits) == e.bits
  48     }
  49
  50     pub fn union(&self, e: EnumSet<E>) -> EnumSet<E> {
  51         EnumSet {bits: self.bits | e.bits}
  52     }
  53
  54     pub fn add(&mut self, e: E) {
  55         self.bits |= bit(e);
  56     }
  57
  58     pub fn contains_elem(&self, e: E) -> bool {
  59         (self.bits & bit(e)) != 0
  60     }
  61
  62     pub fn each(&self, f: &fn(E) -> bool) -> bool {
  63         let mut bits = self.bits;
  64         let mut index = 0;
  65         while bits != 0 {
  66             if (bits & 1) != 0 {
  67                 let e = CLike::from_uint(index);
  68                 if !f(e) {
  69                     return false;
  70                 }
  71             }
  72             index += 1;
  73             bits >>= 1;
  74         }
  75         return true;
  76     }
  77
  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 pub struct EnumSetIterator<E> {
 102     priv index: uint,
 103     priv bits: uint,
 104 }
 105
 106 impl<E:CLike> EnumSetIterator<E> {
 107     fn new(bits: uint) -> EnumSetIterator<E> {
 108         EnumSetIterator { index: 0, bits: bits }
 109     }
 110 }
 111
 112 impl<E:CLike> Iterator<E> for EnumSetIterator<E> {
 113     fn next(&mut self) -> Option<E> {
 114         if (self.bits == 0) {
 115             return None;
 116         }
 117
 118         while (self.bits & 1) == 0 {
 119             self.index += 1;
 120             self.bits >>= 1;
 121         }
 122         let elem = CLike::from_uint(self.index);
 123         self.index += 1;
 124         self.bits >>= 1;
 125         Some(elem)
 126     }
 127
 128     fn size_hint(&self) -> (uint, Option<uint>) {
 129         let exact = self.bits.population_count();
 130         (exact, Some(exact))
 131     }
 132 }
 133
 134 #[cfg(test)]
 135 mod test {
 136
 137     use std::cast;
 138
 139     use util::enum_set::*;
 140
 141     #[deriving(Eq)]
 142     enum Foo {
 143         A, B, C
 144     }
 145
 146     impl CLike for Foo {
 147         pub fn to_uint(&self) -> uint {
 148             *self as uint
 149         }
 150
 151         pub fn from_uint(v: uint) -> Foo {
 152             unsafe { cast::transmute(v) }
 153         }
 154     }
 155
 156     #[test]
 157     fn test_empty() {
 158         let e: EnumSet<Foo> = EnumSet::empty();
 159         assert!(e.is_empty());
 160     }
 161
 162     ///////////////////////////////////////////////////////////////////////////
 163     // intersect
 164
 165     #[test]
 166     fn test_two_empties_do_not_intersect() {
 167         let e1: EnumSet<Foo> = EnumSet::empty();
 168         let e2: EnumSet<Foo> = EnumSet::empty();
 169         assert!(!e1.intersects(e2));
 170     }
 171
 172     #[test]
 173     fn test_empty_does_not_intersect_with_full() {
 174         let e1: EnumSet<Foo> = EnumSet::empty();
 175
 176         let mut e2: EnumSet<Foo> = EnumSet::empty();
 177         e2.add(A);
 178         e2.add(B);
 179         e2.add(C);
 180
 181         assert!(!e1.intersects(e2));
 182     }
 183
 184     #[test]
 185     fn test_disjoint_intersects() {
 186         let mut e1: EnumSet<Foo> = EnumSet::empty();
 187         e1.add(A);
 188
 189         let mut e2: EnumSet<Foo> = EnumSet::empty();
 190         e2.add(B);
 191
 192         assert!(!e1.intersects(e2));
 193     }
 194
 195     #[test]
 196     fn test_overlapping_intersects() {
 197         let mut e1: EnumSet<Foo> = EnumSet::empty();
 198         e1.add(A);
 199
 200         let mut e2: EnumSet<Foo> = EnumSet::empty();
 201         e2.add(A);
 202         e2.add(B);
 203
 204         assert!(e1.intersects(e2));
 205     }
 206
 207     ///////////////////////////////////////////////////////////////////////////
 208     // contains and contains_elem
 209
 210     #[test]
 211     fn test_contains() {
 212         let mut e1: EnumSet<Foo> = EnumSet::empty();
 213         e1.add(A);
 214
 215         let mut e2: EnumSet<Foo> = EnumSet::empty();
 216         e2.add(A);
 217         e2.add(B);
 218
 219         assert!(!e1.contains(e2));
 220         assert!(e2.contains(e1));
 221     }
 222
 223     #[test]
 224     fn test_contains_elem() {
 225         let mut e1: EnumSet<Foo> = EnumSet::empty();
 226         e1.add(A);
 227         assert!(e1.contains_elem(A));
 228         assert!(!e1.contains_elem(B));
 229         assert!(!e1.contains_elem(C));
 230
 231         e1.add(A);
 232         e1.add(B);
 233         assert!(e1.contains_elem(A));
 234         assert!(e1.contains_elem(B));
 235         assert!(!e1.contains_elem(C));
 236     }
 237
 238     ///////////////////////////////////////////////////////////////////////////
 239     // iter / each
 240
 241     #[test]
 242     fn test_iterator() {
 243         let mut e1: EnumSet<Foo> = EnumSet::empty();
 244
 245         let elems: ~[Foo] = e1.iter().collect();
 246         assert_eq!(~[], elems)
 247
 248         e1.add(A);
 249         let elems: ~[Foo] = e1.iter().collect();
 250         assert_eq!(~[A], elems)
 251
 252         e1.add(C);
 253         let elems: ~[Foo] = e1.iter().collect();
 254         assert_eq!(~[A,C], elems)
 255
 256         e1.add(C);
 257         let elems: ~[Foo] = e1.iter().collect();
 258         assert_eq!(~[A,C], elems)
 259
 260         e1.add(B);
 261         let elems: ~[Foo] = e1.iter().collect();
 262         assert_eq!(~[A,B,C], elems)
 263     }
 264
 265     #[test]
 266     fn test_each() {
 267         let mut e1: EnumSet<Foo> = EnumSet::empty();
 268
 269         assert_eq!(~[], collect(e1))
 270
 271         e1.add(A);
 272         assert_eq!(~[A], collect(e1))
 273
 274         e1.add(C);
 275         assert_eq!(~[A,C], collect(e1))
 276
 277         e1.add(C);
 278         assert_eq!(~[A,C], collect(e1))
 279
 280         e1.add(B);
 281         assert_eq!(~[A,B,C], collect(e1))
 282     }
 283
 284     fn collect(e: EnumSet<Foo>) -> ~[Foo] {
 285         let mut elems = ~[];
 286         e.each(|elem| {
 287            elems.push(elem);
 288            true
 289         });
 290         elems
 291     }
 292
 293     ///////////////////////////////////////////////////////////////////////////
 294     // operators
 295
 296     #[test]
 297     fn test_operators() {
 298         let mut e1: EnumSet<Foo> = EnumSet::empty();
 299         e1.add(A);
 300         e1.add(C);
 301
 302         let mut e2: EnumSet<Foo> = EnumSet::empty();
 303         e2.add(B);
 304         e2.add(C);
 305
 306         let e_union = e1 | e2;
 307         let elems: ~[Foo] = e_union.iter().collect();
 308         assert_eq!(~[A,B,C], elems)
 309
 310         let e_intersection = e1 & e2;
 311         let elems: ~[Foo] = e_intersection.iter().collect();
 312         assert_eq!(~[C], elems)
 313
 314         let e_subtract = e1 - e2;
 315         let elems: ~[Foo] = e_subtract.iter().collect();
 316         assert_eq!(~[A], elems)
 317     }
 318 }