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: usize
  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 usize and back.
  51 /// A typically implementation is as below.
  52 ///
  53 /// ```{rust,ignore}
  54 /// #[repr(usize)]
  55 /// enum Foo {
  56 ///     A, B, C
  57 /// }
  58 ///
  59 /// impl CLike for Foo {
  60 ///     fn to_usize(&self) -> usize {
  61 ///         *self as usize
  62 ///     }
  63 ///
  64 ///     fn from_usize(v: usize) -> Foo {
  65 ///         unsafe { mem::transmute(v) }
  66 ///     }
  67 /// }
  68 /// ```
  69 pub trait CLike {
  70     /// Converts a C-like enum to a `usize`.
  71     fn to_usize(&self) -> usize;
  72     /// Converts a `usize` to a C-like enum.
  73     fn from_usize(usize) -> Self;
  74 }
  75
  76 fn bit<E:CLike>(e: &E) -> usize {
  77     use core::usize;
  78     let value = e.to_usize();
  79     assert!(value < usize::BITS,
  80             "EnumSet only supports up to {} variants.", usize::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) -> usize {
  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: usize,
 209     bits: usize,
 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: usize) -> 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_usize(self.index);
 241         self.index += 1;
 242         self.bits >>= 1;
 243         Some(elem)
 244     }
 245
 246     fn size_hint(&self) -> (usize, Option<usize>) {
 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 #[stable(feature = "rust1", since = "1.0.0")]
 261 impl<'a, E> IntoIterator for &'a EnumSet<E> where E: CLike {
 262     type Item = E;
 263     type IntoIter = Iter<E>;
 264
 265     fn into_iter(self) -> Iter<E> {
 266         self.iter()
 267     }
 268 }
 269
 270 impl<E:CLike> Extend<E> for EnumSet<E> {
 271     fn extend<I: Iterator<Item=E>>(&mut self, iterator: I) {
 272         for element in iterator {
 273             self.insert(element);
 274         }
 275     }
 276 }
 277
 278 #[cfg(test)]
 279 mod test {
 280     use self::Foo::*;
 281     use prelude::*;
 282     use core::mem;
 283
 284     use super::{EnumSet, CLike};
 285
 286     #[derive(Copy, PartialEq, Debug)]
 287     #[repr(usize)]
 288     enum Foo {
 289         A, B, C
 290     }
 291
 292     impl CLike for Foo {
 293         fn to_usize(&self) -> usize {
 294             *self as usize
 295         }
 296
 297         fn from_usize(v: usize) -> Foo {
 298             unsafe { mem::transmute(v) }
 299         }
 300     }
 301
 302     #[test]
 303     fn test_new() {
 304         let e: EnumSet<Foo> = EnumSet::new();
 305         assert!(e.is_empty());
 306     }
 307
 308     #[test]
 309     fn test_show() {
 310         let mut e = EnumSet::new();
 311         assert!(format!("{:?}", e) == "EnumSet {}");
 312         e.insert(A);
 313         assert!(format!("{:?}", e) == "EnumSet {A}");
 314         e.insert(C);
 315         assert!(format!("{:?}", e) == "EnumSet {A, C}");
 316     }
 317
 318     #[test]
 319     fn test_len() {
 320         let mut e = EnumSet::new();
 321         assert_eq!(e.len(), 0);
 322         e.insert(A);
 323         e.insert(B);
 324         e.insert(C);
 325         assert_eq!(e.len(), 3);
 326         e.remove(&A);
 327         assert_eq!(e.len(), 2);
 328         e.clear();
 329         assert_eq!(e.len(), 0);
 330     }
 331
 332     ///////////////////////////////////////////////////////////////////////////
 333     // intersect
 334
 335     #[test]
 336     fn test_two_empties_do_not_intersect() {
 337         let e1: EnumSet<Foo> = EnumSet::new();
 338         let e2: EnumSet<Foo> = EnumSet::new();
 339         assert!(e1.is_disjoint(&e2));
 340     }
 341
 342     #[test]
 343     fn test_empty_does_not_intersect_with_full() {
 344         let e1: EnumSet<Foo> = EnumSet::new();
 345
 346         let mut e2: EnumSet<Foo> = EnumSet::new();
 347         e2.insert(A);
 348         e2.insert(B);
 349         e2.insert(C);
 350
 351         assert!(e1.is_disjoint(&e2));
 352     }
 353
 354     #[test]
 355     fn test_disjoint_intersects() {
 356         let mut e1: EnumSet<Foo> = EnumSet::new();
 357         e1.insert(A);
 358
 359         let mut e2: EnumSet<Foo> = EnumSet::new();
 360         e2.insert(B);
 361
 362         assert!(e1.is_disjoint(&e2));
 363     }
 364
 365     #[test]
 366     fn test_overlapping_intersects() {
 367         let mut e1: EnumSet<Foo> = EnumSet::new();
 368         e1.insert(A);
 369
 370         let mut e2: EnumSet<Foo> = EnumSet::new();
 371         e2.insert(A);
 372         e2.insert(B);
 373
 374         assert!(!e1.is_disjoint(&e2));
 375     }
 376
 377     ///////////////////////////////////////////////////////////////////////////
 378     // contains and contains_elem
 379
 380     #[test]
 381     fn test_superset() {
 382         let mut e1: EnumSet<Foo> = EnumSet::new();
 383         e1.insert(A);
 384
 385         let mut e2: EnumSet<Foo> = EnumSet::new();
 386         e2.insert(A);
 387         e2.insert(B);
 388
 389         let mut e3: EnumSet<Foo> = EnumSet::new();
 390         e3.insert(C);
 391
 392         assert!(e1.is_subset(&e2));
 393         assert!(e2.is_superset(&e1));
 394         assert!(!e3.is_superset(&e2));
 395         assert!(!e2.is_superset(&e3))
 396     }
 397
 398     #[test]
 399     fn test_contains() {
 400         let mut e1: EnumSet<Foo> = EnumSet::new();
 401         e1.insert(A);
 402         assert!(e1.contains(&A));
 403         assert!(!e1.contains(&B));
 404         assert!(!e1.contains(&C));
 405
 406         e1.insert(A);
 407         e1.insert(B);
 408         assert!(e1.contains(&A));
 409         assert!(e1.contains(&B));
 410         assert!(!e1.contains(&C));
 411     }
 412
 413     ///////////////////////////////////////////////////////////////////////////
 414     // iter
 415
 416     #[test]
 417     fn test_iterator() {
 418         let mut e1: EnumSet<Foo> = EnumSet::new();
 419
 420         let elems: ::vec::Vec<Foo> = e1.iter().collect();
 421         assert!(elems.is_empty());
 422
 423         e1.insert(A);
 424         let elems: ::vec::Vec<_> = e1.iter().collect();
 425         assert_eq!(vec![A], elems);
 426
 427         e1.insert(C);
 428         let elems: ::vec::Vec<_> = e1.iter().collect();
 429         assert_eq!(vec![A,C], elems);
 430
 431         e1.insert(C);
 432         let elems: ::vec::Vec<_> = e1.iter().collect();
 433         assert_eq!(vec![A,C], elems);
 434
 435         e1.insert(B);
 436         let elems: ::vec::Vec<_> = e1.iter().collect();
 437         assert_eq!(vec![A,B,C], elems);
 438     }
 439
 440     ///////////////////////////////////////////////////////////////////////////
 441     // operators
 442
 443     #[test]
 444     fn test_operators() {
 445         let mut e1: EnumSet<Foo> = EnumSet::new();
 446         e1.insert(A);
 447         e1.insert(C);
 448
 449         let mut e2: EnumSet<Foo> = EnumSet::new();
 450         e2.insert(B);
 451         e2.insert(C);
 452
 453         let e_union = e1 | e2;
 454         let elems: ::vec::Vec<_> = e_union.iter().collect();
 455         assert_eq!(vec![A,B,C], elems);
 456
 457         let e_intersection = e1 & e2;
 458         let elems: ::vec::Vec<_> = e_intersection.iter().collect();
 459         assert_eq!(vec![C], elems);
 460
 461         // Another way to express intersection
 462         let e_intersection = e1 - (e1 - e2);
 463         let elems: ::vec::Vec<_> = e_intersection.iter().collect();
 464         assert_eq!(vec![C], elems);
 465
 466         let e_subtract = e1 - e2;
 467         let elems: ::vec::Vec<_> = e_subtract.iter().collect();
 468         assert_eq!(vec![A], elems);
 469
 470         // Bitwise XOR of two sets, aka symmetric difference
 471         let e_symmetric_diff = e1 ^ e2;
 472         let elems: ::vec::Vec<_> = e_symmetric_diff.iter().collect();
 473         assert_eq!(vec![A,B], elems);
 474
 475         // Another way to express symmetric difference
 476         let e_symmetric_diff = (e1 - e2) | (e2 - e1);
 477         let elems: ::vec::Vec<_> = e_symmetric_diff.iter().collect();
 478         assert_eq!(vec![A,B], elems);
 479
 480         // Yet another way to express symmetric difference
 481         let e_symmetric_diff = (e1 | e2) - (e1 & e2);
 482         let elems: ::vec::Vec<_> = e_symmetric_diff.iter().collect();
 483         assert_eq!(vec![A,B], elems);
 484     }
 485
 486     #[test]
 487     #[should_fail]
 488     fn test_overflow() {
 489         #[allow(dead_code)]
 490         #[derive(Copy)]
 491         #[repr(usize)]
 492         enum Bar {
 493             V00, V01, V02, V03, V04, V05, V06, V07, V08, V09,
 494             V10, V11, V12, V13, V14, V15, V16, V17, V18, V19,
 495             V20, V21, V22, V23, V24, V25, V26, V27, V28, V29,
 496             V30, V31, V32, V33, V34, V35, V36, V37, V38, V39,
 497             V40, V41, V42, V43, V44, V45, V46, V47, V48, V49,
 498             V50, V51, V52, V53, V54, V55, V56, V57, V58, V59,
 499             V60, V61, V62, V63, V64, V65, V66, V67, V68, V69,
 500         }
 501
 502         impl CLike for Bar {
 503             fn to_usize(&self) -> usize {
 504                 *self as usize
 505             }
 506
 507             fn from_usize(v: usize) -> Bar {
 508                 unsafe { mem::transmute(v) }
 509             }
 510         }
 511         let mut set = EnumSet::new();
 512         set.insert(Bar::V64);
 513     }
 514 }