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