enum_set/
lib.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
16use std::fmt;
17use std::hash;
18use std::marker::PhantomData;
19use std::iter;
20use std::ops;
21
22#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
23/// A specialized set implementation to use enum types.
24pub struct EnumSet<E> {
25    // We must maintain the invariant that no bits are set
26    // for which no variant exists
27    bits: u32,
28    phantom: PhantomData<E>,
29}
30
31impl<E: CLike + fmt::Debug> fmt::Debug for EnumSet<E> {
32    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
33        fmt.debug_set().entries(self).finish()
34    }
35}
36
37impl<E: CLike> hash::Hash for EnumSet<E> {
38    fn hash<H: hash::Hasher>(&self, state: &mut H) {
39        self.bits.hash(state);
40    }
41}
42
43/// An interface for casting C-like enum to `u32` and back.
44///
45/// The returned value must be no more than 31: `EnumSet` does not support more cases than this.
46///
47/// A typical implementation can be seen below:
48///
49/// ```
50/// use enum_set::CLike;
51/// use std::mem;
52///
53/// #[derive(Clone, Copy)]
54/// #[repr(u32)]
55/// enum Foo {
56///     A, B, C
57/// }
58///
59/// impl CLike for Foo {
60///     fn to_u32(&self) -> u32 {
61///         *self as u32
62///     }
63///
64///     unsafe fn from_u32(v: u32) -> Foo {
65///         mem::transmute(v)
66///     }
67/// }
68/// ```
69pub trait CLike {
70    /// Converts a C-like enum to a `u32`. The value must be `<= 31`.
71    fn to_u32(&self) -> u32;
72
73    /// Converts a `u32` to a C-like enum. This method only needs to be safe
74    /// for possible return values of `to_u32` of this trait.
75    unsafe fn from_u32(u32) -> Self;
76}
77
78fn bit<E: CLike>(e: &E) -> u32 {
79    let value = e.to_u32();
80    assert!(value < 32, "EnumSet only supports up to {} variants.", 31);
81    1 << value
82}
83
84impl<E: CLike> EnumSet<E> {
85    /// Returns an empty `EnumSet`.
86    pub fn new() -> Self {
87        Self::new_with_bits(0)
88    }
89
90    fn new_with_bits(bits: u32) -> Self {
91        EnumSet { bits: bits, phantom: PhantomData }
92    }
93
94    /// Returns the number of elements in the set.
95    pub fn len(&self) -> usize {
96        self.bits.count_ones() as usize
97    }
98
99    /// Checks if the set is empty.
100    pub fn is_empty(&self) -> bool {
101        self.bits == 0
102    }
103
104    /// Removes all elements from the set.
105    pub fn clear(&mut self) {
106        self.bits = 0;
107    }
108
109    /// Returns `true` if the set has no elements in common with `other`.
110    ///
111    /// This is equivalent to checking for an empty intersection.
112    pub fn is_disjoint(&self, other: &Self) -> bool {
113        (self.bits & other.bits) == 0
114    }
115
116    /// Returns `true` if the set is a superset of `other`.
117    pub fn is_superset(&self, other: &Self) -> bool {
118        (self.bits & other.bits) == other.bits
119    }
120
121    /// Returns `true` if the set is a subset of `other`.
122    pub fn is_subset(&self, other: &Self) -> bool {
123        other.is_superset(self)
124    }
125
126    /// Returns the union of the set and `other`.
127    pub fn union(&self, other: Self) -> Self {
128        Self::new_with_bits(self.bits | other.bits)
129    }
130
131    /// Returns the intersection of the set and `other`.
132    pub fn intersection(&self, other: Self) -> Self {
133        Self::new_with_bits(self.bits & other.bits)
134    }
135
136    /// Returns the difference between the set and `other`.
137    pub fn difference(&self, other: Self) -> Self {
138        Self::new_with_bits(self.bits & !other.bits)
139    }
140
141    /// Returns the symmetric difference between the set and `other`.
142    pub fn symmetric_difference(&self, other: Self) -> Self {
143        Self::new_with_bits(self.bits ^ other.bits)
144    }
145
146    /// Adds the given value to the set.
147    ///
148    /// Returns `true` if the value was not already present in the set.
149    pub fn insert(&mut self, value: E) -> bool {
150        let result = !self.contains(&value);
151        self.bits |= bit(&value);
152        result
153    }
154
155    /// Removes a value from the set.
156    ///
157    /// Returns `true` if the value was present in the set.
158    pub fn remove(&mut self, value: &E) -> bool {
159        let result = self.contains(value);
160        self.bits &= !bit(value);
161        result
162    }
163
164    /// Returns `true` if the set contains the given value.
165    pub fn contains(&self, value: &E) -> bool {
166        (self.bits & bit(value)) != 0
167    }
168
169    /// Returns an iterator over the set's elements.
170    pub fn iter(&self) -> Iter<E> {
171        Iter { index: 0, bits: self.bits, phantom: PhantomData }
172    }
173}
174
175impl<E: CLike> ops::Sub for EnumSet<E> {
176    type Output = Self;
177
178    fn sub(self, other: Self) -> Self {
179        self.difference(other)
180    }
181}
182
183impl<E: CLike> ops::BitOr for EnumSet<E> {
184    type Output = Self;
185
186    fn bitor(self, other: Self) -> Self {
187        self.union(other)
188    }
189}
190
191impl<E: CLike> ops::BitAnd for EnumSet<E> {
192    type Output = Self;
193
194    fn bitand(self, other: Self) -> Self {
195        self.intersection(other)
196    }
197}
198
199impl<E: CLike> ops::BitXor for EnumSet<E> {
200    type Output = Self;
201
202    fn bitxor(self, other: Self) -> Self {
203        self.symmetric_difference(other)
204    }
205}
206
207impl<E: CLike> ops::SubAssign for EnumSet<E> {
208    fn sub_assign(&mut self, other: Self) {
209        self.bits &= !other.bits;
210    }
211}
212
213impl<E: CLike> ops::BitOrAssign for EnumSet<E> {
214    fn bitor_assign(&mut self, other: Self) {
215        self.bits |= other.bits;
216    }
217}
218
219impl<E: CLike> ops::BitAndAssign for EnumSet<E> {
220    fn bitand_assign(&mut self, other: Self) {
221        self.bits &= other.bits;
222    }
223}
224
225impl<E: CLike> ops::BitXorAssign for EnumSet<E> {
226    fn bitxor_assign(&mut self, other: Self) {
227        self.bits ^= other.bits;
228    }
229}
230
231#[derive(Clone)]
232/// An iterator over an `EnumSet`.
233pub struct Iter<E> {
234    index: u32,
235    bits: u32,
236    phantom: PhantomData<*mut E>,
237}
238
239impl<E: CLike> Iterator for Iter<E> {
240    type Item = E;
241
242    fn next(&mut self) -> Option<E> {
243        if self.bits == 0 {
244            return None;
245        }
246
247        while (self.bits & 1) == 0 {
248            self.index += 1;
249            self.bits >>= 1;
250        }
251
252        // Safe because of the invariant that only valid bits are set (see
253        // comment on the `bit` member of this struct).
254        let elem = unsafe { CLike::from_u32(self.index) };
255        self.index += 1;
256        self.bits >>= 1;
257        Some(elem)
258    }
259
260    fn size_hint(&self) -> (usize, Option<usize>) {
261        let exact = self.bits.count_ones() as usize;
262        (exact, Some(exact))
263    }
264}
265
266impl<E: CLike> ExactSizeIterator for Iter<E> {}
267
268impl<E: CLike> Default for EnumSet<E> {
269    fn default() -> Self {
270        Self::new()
271    }
272}
273
274impl<E: CLike> iter::FromIterator<E> for EnumSet<E> {
275    fn from_iter<I: IntoIterator<Item = E>>(iterator: I) -> Self {
276        let mut ret = Self::new();
277        ret.extend(iterator);
278        ret
279    }
280}
281
282impl<E: CLike> Extend<E> for EnumSet<E> {
283    fn extend<I: IntoIterator<Item = E>>(&mut self, iter: I) {
284        for element in iter {
285            self.insert(element);
286        }
287    }
288}
289
290impl<'a, E: CLike> IntoIterator for &'a EnumSet<E> {
291    type Item = E;
292    type IntoIter = Iter<E>;
293    fn into_iter(self) -> Iter<E> { self.iter() }
294}
295
296#[cfg(test)]
297mod tests {
298    use self::Foo::*;
299    use std::mem;
300
301    use super::{EnumSet, CLike};
302
303    #[derive(Copy, Clone, PartialEq, Debug)]
304    #[repr(u32)]
305    enum Foo {
306        A, B, C
307    }
308
309    impl CLike for Foo {
310        fn to_u32(&self) -> u32 {
311            *self as u32
312        }
313
314        unsafe fn from_u32(v: u32) -> Foo {
315            mem::transmute(v)
316        }
317    }
318
319    #[test]
320    fn test_new() {
321        let e: EnumSet<Foo> = EnumSet::new();
322        assert!(e.is_empty());
323    }
324
325    #[test]
326    fn test_debug() {
327        let mut e = EnumSet::new();
328        assert_eq!("{}", format!("{:?}", e));
329        e.insert(A);
330        assert_eq!("{A}", format!("{:?}", e));
331        e.insert(C);
332        assert_eq!("{A, C}", format!("{:?}", e));
333    }
334
335    #[test]
336    fn test_len() {
337        let mut e = EnumSet::new();
338        assert_eq!(e.len(), 0);
339        e.insert(A);
340        e.insert(B);
341        e.insert(C);
342        assert_eq!(e.len(), 3);
343        e.remove(&A);
344        assert_eq!(e.len(), 2);
345        e.clear();
346        assert_eq!(e.len(), 0);
347    }
348
349    ///////////////////////////////////////////////////////////////////////////
350    // intersect
351
352    #[test]
353    fn test_two_empties_do_not_intersect() {
354        let e1: EnumSet<Foo> = EnumSet::new();
355        let e2: EnumSet<Foo> = EnumSet::new();
356        assert!(e1.is_disjoint(&e2));
357    }
358
359    #[test]
360    fn test_empty_does_not_intersect_with_full() {
361        let e1: EnumSet<Foo> = EnumSet::new();
362
363        let mut e2: EnumSet<Foo> = EnumSet::new();
364        e2.insert(A);
365        e2.insert(B);
366        e2.insert(C);
367
368        assert!(e1.is_disjoint(&e2));
369    }
370
371    #[test]
372    fn test_disjoint_intersects() {
373        let mut e1: EnumSet<Foo> = EnumSet::new();
374        e1.insert(A);
375
376        let mut e2: EnumSet<Foo> = EnumSet::new();
377        e2.insert(B);
378
379        assert!(e1.is_disjoint(&e2));
380    }
381
382    #[test]
383    fn test_overlapping_intersects() {
384        let mut e1: EnumSet<Foo> = EnumSet::new();
385        e1.insert(A);
386
387        let mut e2: EnumSet<Foo> = EnumSet::new();
388        e2.insert(A);
389        e2.insert(B);
390
391        assert!(!e1.is_disjoint(&e2));
392    }
393
394    ///////////////////////////////////////////////////////////////////////////
395    // contains and contains_elem
396
397    #[test]
398    fn test_superset() {
399        let mut e1: EnumSet<Foo> = EnumSet::new();
400        e1.insert(A);
401
402        let mut e2: EnumSet<Foo> = EnumSet::new();
403        e2.insert(A);
404        e2.insert(B);
405
406        let mut e3: EnumSet<Foo> = EnumSet::new();
407        e3.insert(C);
408
409        assert!(e1.is_subset(&e2));
410        assert!(e2.is_superset(&e1));
411        assert!(!e3.is_superset(&e2));
412        assert!(!e2.is_superset(&e3));
413    }
414
415    #[test]
416    fn test_contains() {
417        let mut e1: EnumSet<Foo> = EnumSet::new();
418        e1.insert(A);
419        assert!(e1.contains(&A));
420        assert!(!e1.contains(&B));
421        assert!(!e1.contains(&C));
422
423        e1.insert(A);
424        e1.insert(B);
425        assert!(e1.contains(&A));
426        assert!(e1.contains(&B));
427        assert!(!e1.contains(&C));
428    }
429
430    ///////////////////////////////////////////////////////////////////////////
431    // iter
432
433    #[test]
434    fn test_iterator() {
435        let mut e1: EnumSet<Foo> = EnumSet::new();
436
437        let elems: Vec<Foo> = e1.iter().collect();
438        assert!(elems.is_empty());
439
440        e1.insert(A);
441        let elems: Vec<_> = e1.iter().collect();
442        assert_eq!(vec![A], elems);
443
444        e1.insert(C);
445        let elems: Vec<_> = e1.iter().collect();
446        assert_eq!(vec![A,C], elems);
447
448        e1.insert(C);
449        let elems: Vec<_> = e1.iter().collect();
450        assert_eq!(vec![A,C], elems);
451
452        e1.insert(B);
453        let elems: Vec<_> = e1.iter().collect();
454        assert_eq!(vec![A,B,C], elems);
455    }
456
457    #[test]
458    fn test_clone_iterator() {
459        let mut e: EnumSet<Foo> = EnumSet::new();
460        e.insert(A);
461        e.insert(B);
462        e.insert(C);
463
464        let mut iter1 = e.iter();
465        let first_elem = iter1.next();
466        assert_eq!(Some(A), first_elem);
467
468        let iter2 = iter1.clone();
469        let elems1: Vec<_> = iter1.collect();
470        assert_eq!(vec![B, C], elems1);
471
472        let elems2: Vec<_> = iter2.collect();
473        assert_eq!(vec![B, C], elems2);
474    }
475
476    ///////////////////////////////////////////////////////////////////////////
477    // operators
478
479    #[test]
480    fn test_operators() {
481        let mut e1: EnumSet<Foo> = EnumSet::new();
482        e1.insert(A);
483        e1.insert(C);
484
485        let mut e2: EnumSet<Foo> = EnumSet::new();
486        e2.insert(B);
487        e2.insert(C);
488
489        let e_union = e1 | e2;
490        let elems: Vec<_> = e_union.iter().collect();
491        assert_eq!(vec![A,B,C], elems);
492
493        let e_intersection = e1 & e2;
494        let elems: Vec<_> = e_intersection.iter().collect();
495        assert_eq!(vec![C], elems);
496
497        // Another way to express intersection
498        let e_intersection = e1 - (e1 - e2);
499        let elems: Vec<_> = e_intersection.iter().collect();
500        assert_eq!(vec![C], elems);
501
502        let e_subtract = e1 - e2;
503        let elems: Vec<_> = e_subtract.iter().collect();
504        assert_eq!(vec![A], elems);
505
506        // Bitwise XOR of two sets, aka symmetric difference
507        let e_symmetric_diff = e1 ^ e2;
508        let elems: Vec<_> = e_symmetric_diff.iter().collect();
509        assert_eq!(vec![A,B], elems);
510
511        // Another way to express symmetric difference
512        let e_symmetric_diff = (e1 - e2) | (e2 - e1);
513        let elems: Vec<_> = e_symmetric_diff.iter().collect();
514        assert_eq!(vec![A,B], elems);
515
516        // Yet another way to express symmetric difference
517        let e_symmetric_diff = (e1 | e2) - (e1 & e2);
518        let elems: Vec<_> = e_symmetric_diff.iter().collect();
519        assert_eq!(vec![A,B], elems);
520    }
521
522    #[test]
523    fn test_assign_operators() {
524        let mut e1: EnumSet<Foo> = EnumSet::new();
525        e1.insert(A);
526        e1.insert(C);
527
528        let mut e2: EnumSet<Foo> = EnumSet::new();
529        e2.insert(B);
530        e2.insert(C);
531
532        let mut e_union = e1.clone();
533        e_union |= e2;
534        let elems: Vec<_> = e_union.iter().collect();
535        assert_eq!(vec![A,B,C], elems);
536
537        let mut e_intersection = e1.clone();
538        e_intersection &= e2;
539        let elems: Vec<_> = e_intersection.iter().collect();
540        assert_eq!(vec![C], elems);
541
542        let mut e_subtract = e1.clone();
543        e_subtract -= e2;
544        let elems: Vec<_> = e_subtract.iter().collect();
545        assert_eq!(vec![A], elems);
546
547        // Bitwise XOR of two sets, aka symmetric difference
548        let mut e_symmetric_diff = e1.clone();
549        e_symmetric_diff ^= e2;
550        let elems: Vec<_> = e_symmetric_diff.iter().collect();
551        assert_eq!(vec![A,B], elems);
552    }
553
554    #[test]
555    #[should_panic]
556    fn test_overflow() {
557        #[allow(dead_code)]
558        #[repr(u32)]
559        #[derive(Clone, Copy)]
560        enum Bar {
561            V00, V01, V02, V03, V04, V05, V06, V07, V08, V09,
562            V10, V11, V12, V13, V14, V15, V16, V17, V18, V19,
563            V20, V21, V22, V23, V24, V25, V26, V27, V28, V29,
564            V30, V31, V32, V33, V34, V35, V36, V37, V38, V39,
565        }
566
567        impl CLike for Bar {
568            fn to_u32(&self) -> u32 {
569                *self as u32
570            }
571
572            unsafe fn from_u32(v: u32) -> Bar {
573                mem::transmute(v)
574            }
575        }
576
577        let mut set = EnumSet::new();
578        set.insert(Bar::V32);
579    }
580}