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 core::prelude::*;
  17 use core::fmt;
  18 use core::num::Int;
  19 use core::iter::FromIterator;
  20 use core::ops::{Sub, BitOr, BitAnd, BitXor};
  21
  22 // FIXME(contentions): implement union family of methods? (general design may be wrong here)
  23
  24 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
  25 /// A specialized set implementation to use enum types.
  26 pub struct EnumSet<E> {
  27     // We must maintain the invariant that no bits are set
  28     // for which no variant exists
  29     bits: uint
  30 }
  31
  32 impl<E> Copy for EnumSet<E> {}
  33
  34 impl<E:CLike+fmt::Show> fmt::Show for EnumSet<E> {
  35     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
  36         try!(write!(fmt, "{{"));
  37         let mut first = true;
  38         for e in self.iter() {
  39             if !first {
  40                 try!(write!(fmt, ", "));
  41             }
  42             try!(write!(fmt, "{}", e));
  43             first = false;
  44         }
  45         write!(fmt, "}}")
  46     }
  47 }
  48
  49 /// An interface for casting C-like enum to uint and back.
  50 /// A typically implementation is as below.
  51 ///
  52 /// ```{rust,ignore}
  53 /// #[repr(uint)]
  54 /// enum Foo {
  55 ///     A, B, C
  56 /// }
  57 ///
  58 /// impl CLike for Foo {
  59 ///     fn to_uint(&self) -> uint {
  60 ///         *self as uint
  61 ///     }
  62 ///
  63 ///     fn from_uint(v: uint) -> Foo {
  64 ///         unsafe { mem::transmute(v) }
  65 ///     }
  66 /// }
  67 /// ```
  68 pub trait CLike {
  69     /// Converts a C-like enum to a `uint`.
  70     fn to_uint(&self) -> uint;
  71     /// Converts a `uint` to a C-like enum.
  72     fn from_uint(uint) -> Self;
  73 }
  74
  75 fn bit<E:CLike>(e: &E) -> uint {
  76     use core::uint;
  77     let value = e.to_uint();
  78     assert!(value < uint::BITS,
  79             "EnumSet only supports up to {} variants.", uint::BITS - 1);
  80     1 << value
  81 }
  82
  83 impl<E:CLike> EnumSet<E> {
  84     /// Returns an empty `EnumSet`.
  85     #[unstable = "matches collection reform specification, waiting for dust to settle"]
  86     pub fn new() -> EnumSet<E> {
  87         EnumSet {bits: 0}
  88     }
  89
  90     /// Returns the number of elements in the given `EnumSet`.
  91     #[unstable = "matches collection reform specification, waiting for dust to settle"]
  92     pub fn len(&self) -> uint {
  93         self.bits.count_ones()
  94     }
  95
  96     /// Returns true if the `EnumSet` is empty.
  97     #[unstable = "matches collection reform specification, waiting for dust to settle"]
  98     pub fn is_empty(&self) -> bool {
  99         self.bits == 0
 100     }
 101
 102     pub fn clear(&mut self) {
 103         self.bits = 0;
 104     }
 105
 106     /// Returns `false` if the `EnumSet` contains any enum of the given `EnumSet`.
 107     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 108     pub fn is_disjoint(&self, other: &EnumSet<E>) -> bool {
 109         (self.bits & other.bits) == 0
 110     }
 111
 112     /// Returns `true` if a given `EnumSet` is included in this `EnumSet`.
 113     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 114     pub fn is_superset(&self, other: &EnumSet<E>) -> bool {
 115         (self.bits & other.bits) == other.bits
 116     }
 117
 118     /// Returns `true` if this `EnumSet` is included in the given `EnumSet`.
 119     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 120     pub fn is_subset(&self, other: &EnumSet<E>) -> bool {
 121         other.is_superset(self)
 122     }
 123
 124     /// Returns the union of both `EnumSets`.
 125     pub fn union(&self, e: EnumSet<E>) -> EnumSet<E> {
 126         EnumSet {bits: self.bits | e.bits}
 127     }
 128
 129     /// Returns the intersection of both `EnumSets`.
 130     pub fn intersection(&self, e: EnumSet<E>) -> EnumSet<E> {
 131         EnumSet {bits: self.bits & e.bits}
 132     }
 133
 134     /// Adds an enum to the `EnumSet`, and returns `true` if it wasn't there before
 135     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 136     pub fn insert(&mut self, e: E) -> bool {
 137         let result = !self.contains(&e);
 138         self.bits |= bit(&e);
 139         result
 140     }
 141
 142     /// Removes an enum from the EnumSet
 143     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 144     pub fn remove(&mut self, e: &E) -> bool {
 145         let result = self.contains(e);
 146         self.bits &= !bit(e);
 147         result
 148     }
 149
 150     /// Returns `true` if an `EnumSet` contains a given enum.
 151     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 152     pub fn contains(&self, e: &E) -> bool {
 153         (self.bits & bit(e)) != 0
 154     }
 155
 156     /// Returns an iterator over an `EnumSet`.
 157     #[unstable = "matches collection reform specification, waiting for dust to settle"]
 158     pub fn iter(&self) -> Iter<E> {
 159         Iter::new(self.bits)
 160     }
 161 }
 162
 163 impl<E:CLike> Sub for EnumSet<E> {
 164     type Output = EnumSet<E>;
 165
 166     fn sub(self, e: EnumSet<E>) -> EnumSet<E> {
 167         EnumSet {bits: self.bits & !e.bits}
 168     }
 169 }
 170
 171 impl<E:CLike> BitOr for EnumSet<E> {
 172     type Output = EnumSet<E>;
 173
 174     fn bitor(self, e: EnumSet<E>) -> EnumSet<E> {
 175         EnumSet {bits: self.bits | e.bits}
 176     }
 177 }
 178
 179 impl<E:CLike> BitAnd for EnumSet<E> {
 180     type Output = EnumSet<E>;
 181
 182     fn bitand(self, e: EnumSet<E>) -> EnumSet<E> {
 183         EnumSet {bits: self.bits & e.bits}
 184     }
 185 }
 186
 187 impl<E:CLike> BitXor for EnumSet<E> {
 188     type Output = EnumSet<E>;
 189
 190     fn bitxor(self, e: EnumSet<E>) -> EnumSet<E> {
 191         EnumSet {bits: self.bits ^ e.bits}
 192     }
 193 }
 194
 195 /// An iterator over an EnumSet
 196 pub struct Iter<E> {
 197     index: uint,
 198     bits: uint,
 199 }
 200
 201 // FIXME(#19839) Remove in favor of `#[derive(Clone)]`
 202 impl<E> Clone for Iter<E> {
 203     fn clone(&self) -> Iter<E> {
 204         Iter {
 205             index: self.index,
 206             bits: self.bits,
 207         }
 208     }
 209 }
 210
 211 impl<E:CLike> Iter<E> {
 212     fn new(bits: uint) -> Iter<E> {
 213         Iter { index: 0, bits: bits }
 214     }
 215 }
 216
 217 impl<E:CLike> Iterator for Iter<E> {
 218     type Item = E;
 219
 220     fn next(&mut self) -> Option<E> {
 221         if self.bits == 0 {
 222             return None;
 223         }
 224
 225         while (self.bits & 1) == 0 {
 226             self.index += 1;
 227             self.bits >>= 1;
 228         }
 229         let elem = CLike::from_uint(self.index);
 230         self.index += 1;
 231         self.bits >>= 1;
 232         Some(elem)
 233     }
 234
 235     fn size_hint(&self) -> (uint, Option<uint>) {
 236         let exact = self.bits.count_ones();
 237         (exact, Some(exact))
 238     }
 239 }
 240
 241 impl<E:CLike> FromIterator<E> for EnumSet<E> {
 242     fn from_iter<I:Iterator<Item=E>>(iterator: I) -> EnumSet<E> {
 243         let mut ret = EnumSet::new();
 244         ret.extend(iterator);
 245         ret
 246     }
 247 }
 248
 249 impl<E:CLike> Extend<E> for EnumSet<E> {
 250     fn extend<I: Iterator<Item=E>>(&mut self, mut iterator: I) {
 251         for element in iterator {
 252             self.insert(element);
 253         }
 254     }
 255 }
 256
 257 #[cfg(test)]
 258 mod test {
 259     use self::Foo::*;
 260     use prelude::*;
 261     use core::mem;
 262
 263     use super::{EnumSet, CLike};
 264
 265     #[derive(Copy, PartialEq, Show)]
 266     #[repr(uint)]
 267     enum Foo {
 268         A, B, C
 269     }
 270
 271     impl CLike for Foo {
 272         fn to_uint(&self) -> uint {
 273             *self as uint
 274         }
 275
 276         fn from_uint(v: uint) -> Foo {
 277             unsafe { mem::transmute(v) }
 278         }
 279     }
 280
 281     #[test]
 282     fn test_new() {
 283         let e: EnumSet<Foo> = EnumSet::new();
 284         assert!(e.is_empty());
 285     }
 286
 287     #[test]
 288     fn test_show() {
 289         let mut e = EnumSet::new();
 290         assert_eq!("{}", e.to_string());
 291         e.insert(A);
 292         assert_eq!("{A}", e.to_string());
 293         e.insert(C);
 294         assert_eq!("{A, C}", e.to_string());
 295     }
 296
 297     #[test]
 298     fn test_len() {
 299         let mut e = EnumSet::new();
 300         assert_eq!(e.len(), 0);
 301         e.insert(A);
 302         e.insert(B);
 303         e.insert(C);
 304         assert_eq!(e.len(), 3);
 305         e.remove(&A);
 306         assert_eq!(e.len(), 2);
 307         e.clear();
 308         assert_eq!(e.len(), 0);
 309     }
 310
 311     ///////////////////////////////////////////////////////////////////////////
 312     // intersect
 313
 314     #[test]
 315     fn test_two_empties_do_not_intersect() {
 316         let e1: EnumSet<Foo> = EnumSet::new();
 317         let e2: EnumSet<Foo> = EnumSet::new();
 318         assert!(e1.is_disjoint(&e2));
 319     }
 320
 321     #[test]
 322     fn test_empty_does_not_intersect_with_full() {
 323         let e1: EnumSet<Foo> = EnumSet::new();
 324
 325         let mut e2: EnumSet<Foo> = EnumSet::new();
 326         e2.insert(A);
 327         e2.insert(B);
 328         e2.insert(C);
 329
 330         assert!(e1.is_disjoint(&e2));
 331     }
 332
 333     #[test]
 334     fn test_disjoint_intersects() {
 335         let mut e1: EnumSet<Foo> = EnumSet::new();
 336         e1.insert(A);
 337
 338         let mut e2: EnumSet<Foo> = EnumSet::new();
 339         e2.insert(B);
 340
 341         assert!(e1.is_disjoint(&e2));
 342     }
 343
 344     #[test]
 345     fn test_overlapping_intersects() {
 346         let mut e1: EnumSet<Foo> = EnumSet::new();
 347         e1.insert(A);
 348
 349         let mut e2: EnumSet<Foo> = EnumSet::new();
 350         e2.insert(A);
 351         e2.insert(B);
 352
 353         assert!(!e1.is_disjoint(&e2));
 354     }
 355
 356     ///////////////////////////////////////////////////////////////////////////
 357     // contains and contains_elem
 358
 359     #[test]
 360     fn test_superset() {
 361         let mut e1: EnumSet<Foo> = EnumSet::new();
 362         e1.insert(A);
 363
 364         let mut e2: EnumSet<Foo> = EnumSet::new();
 365         e2.insert(A);
 366         e2.insert(B);
 367
 368         let mut e3: EnumSet<Foo> = EnumSet::new();
 369         e3.insert(C);
 370
 371         assert!(e1.is_subset(&e2));
 372         assert!(e2.is_superset(&e1));
 373         assert!(!e3.is_superset(&e2));
 374         assert!(!e2.is_superset(&e3))
 375     }
 376
 377     #[test]
 378     fn test_contains() {
 379         let mut e1: EnumSet<Foo> = EnumSet::new();
 380         e1.insert(A);
 381         assert!(e1.contains(&A));
 382         assert!(!e1.contains(&B));
 383         assert!(!e1.contains(&C));
 384
 385         e1.insert(A);
 386         e1.insert(B);
 387         assert!(e1.contains(&A));
 388         assert!(e1.contains(&B));
 389         assert!(!e1.contains(&C));
 390     }
 391
 392     ///////////////////////////////////////////////////////////////////////////
 393     // iter
 394
 395     #[test]
 396     fn test_iterator() {
 397         let mut e1: EnumSet<Foo> = EnumSet::new();
 398
 399         let elems: ::vec::Vec<Foo> = e1.iter().collect();
 400         assert!(elems.is_empty());
 401
 402         e1.insert(A);
 403         let elems: ::vec::Vec<_> = e1.iter().collect();
 404         assert_eq!(vec![A], elems);
 405
 406         e1.insert(C);
 407         let elems: ::vec::Vec<_> = e1.iter().collect();
 408         assert_eq!(vec![A,C], elems);
 409
 410         e1.insert(C);
 411         let elems: ::vec::Vec<_> = e1.iter().collect();
 412         assert_eq!(vec![A,C], elems);
 413
 414         e1.insert(B);
 415         let elems: ::vec::Vec<_> = e1.iter().collect();
 416         assert_eq!(vec![A,B,C], elems);
 417     }
 418
 419     ///////////////////////////////////////////////////////////////////////////
 420     // operators
 421
 422     #[test]
 423     fn test_operators() {
 424         let mut e1: EnumSet<Foo> = EnumSet::new();
 425         e1.insert(A);
 426         e1.insert(C);
 427
 428         let mut e2: EnumSet<Foo> = EnumSet::new();
 429         e2.insert(B);
 430         e2.insert(C);
 431
 432         let e_union = e1 | e2;
 433         let elems: ::vec::Vec<_> = e_union.iter().collect();
 434         assert_eq!(vec![A,B,C], elems);
 435
 436         let e_intersection = e1 & e2;
 437         let elems: ::vec::Vec<_> = e_intersection.iter().collect();
 438         assert_eq!(vec![C], elems);
 439
 440         // Another way to express intersection
 441         let e_intersection = e1 - (e1 - e2);
 442         let elems: ::vec::Vec<_> = e_intersection.iter().collect();
 443         assert_eq!(vec![C], elems);
 444
 445         let e_subtract = e1 - e2;
 446         let elems: ::vec::Vec<_> = e_subtract.iter().collect();
 447         assert_eq!(vec![A], elems);
 448
 449         // Bitwise XOR of two sets, aka symmetric difference
 450         let e_symmetric_diff = e1 ^ e2;
 451         let elems: ::vec::Vec<_> = e_symmetric_diff.iter().collect();
 452         assert_eq!(vec![A,B], elems);
 453
 454         // Another way to express symmetric difference
 455         let e_symmetric_diff = (e1 - e2) | (e2 - e1);
 456         let elems: ::vec::Vec<_> = e_symmetric_diff.iter().collect();
 457         assert_eq!(vec![A,B], elems);
 458
 459         // Yet another way to express symmetric difference
 460         let e_symmetric_diff = (e1 | e2) - (e1 & e2);
 461         let elems: ::vec::Vec<_> = e_symmetric_diff.iter().collect();
 462         assert_eq!(vec![A,B], elems);
 463     }
 464
 465     #[test]
 466     #[should_fail]
 467     fn test_overflow() {
 468         #[allow(dead_code)]
 469         #[derive(Copy)]
 470         #[repr(uint)]
 471         enum Bar {
 472             V00, V01, V02, V03, V04, V05, V06, V07, V08, V09,
 473             V10, V11, V12, V13, V14, V15, V16, V17, V18, V19,
 474             V20, V21, V22, V23, V24, V25, V26, V27, V28, V29,
 475             V30, V31, V32, V33, V34, V35, V36, V37, V38, V39,
 476             V40, V41, V42, V43, V44, V45, V46, V47, V48, V49,
 477             V50, V51, V52, V53, V54, V55, V56, V57, V58, V59,
 478             V60, V61, V62, V63, V64, V65, V66, V67, V68, V69,
 479         }
 480
 481         impl CLike for Bar {
 482             fn to_uint(&self) -> uint {
 483                 *self as uint
 484             }
 485
 486             fn from_uint(v: uint) -> Bar {
 487                 unsafe { mem::transmute(v) }
 488             }
 489         }
 490         let mut set = EnumSet::new();
 491         set.insert(Bar::V64);
 492     }
 493 }