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