Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Iterator over all the enabled options #278

Merged
merged 10 commits into from Apr 28, 2022
89 changes: 89 additions & 0 deletions src/lib.rs
Expand Up @@ -730,6 +730,48 @@ macro_rules! __impl_bitflags {
Self::from_bits_truncate(!self.bits)
}

/// Returns an iterator over all the flags in this set.
pub fn iter(mut self) -> impl Iterator<Item = Self> {
const NUM_FLAGS: usize = {
#[allow(unused_mut)]
let mut num_flags = 0;

$(
#[allow(unused_doc_comments, unused_attributes)]
$(#[$attr $($args)*])*
{
num_flags += 1;
}
)*

num_flags
};
const OPTIONS: [$BitFlags; NUM_FLAGS] = [
$(
#[allow(unused_doc_comments, unused_attributes)]
$(#[$attr $($args)*])*
$BitFlags::$Flag,
)*
];
let mut start = 0;

$crate::_core::iter::from_fn(move || {
if self.is_empty() || NUM_FLAGS == 0 {
None
}else{
for flag in OPTIONS[start..NUM_FLAGS].iter().copied() {
start += 1;
if self.contains(flag) {
self.remove(flag);
return Some(flag)
}
}

None
}
})
}

}

impl $crate::_core::ops::BitOr for $BitFlags {
Expand Down Expand Up @@ -1797,4 +1839,51 @@ mod tests {
const D = 8;
}
}

#[test]
fn test_iter() {
bitflags! {
struct Flags: u32 {
const ONE = 0b001;
const TWO = 0b010;
const THREE = 0b100;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be possible to add a cfg'd-out flag to this as well just to make sure we handle those? Something like:

#[cfg(windows)]
const FOUR_WIN = 0b1000;
#[cfg(unix)]
const FOUR_UNIX = 0b10000;

and then use those matching cfgs to check that we iterate over those flags when they're available, but don't when they're not?

}
}

let flags = Flags::all();
assert_eq!(flags.iter().count(), 3);
let mut iter = flags.iter();
assert_eq!(iter.next().unwrap(), Flags::ONE);
assert_eq!(iter.next().unwrap(), Flags::TWO);
assert_eq!(iter.next().unwrap(), Flags::THREE);
assert_eq!(iter.next(), None);

let flags = Flags::empty();
assert_eq!(flags.iter().count(), 0);

let flags = Flags::ONE | Flags::THREE;
assert_eq!(flags.iter().count(), 2);
let mut iter = flags.iter();
assert_eq!(iter.next().unwrap(), Flags::ONE);
assert_eq!(iter.next().unwrap(), Flags::THREE);
assert_eq!(iter.next(), None);
}

#[test]
fn test_iter_edge_cases() {
bitflags! {
struct Flags: u8 {
const A = 0b00000001;
const BC = 0b00000110;
}
}


let flags = Flags::all();
assert_eq!(flags.iter().count(), 2);
let mut iter = flags.iter();
assert_eq!(iter.next().unwrap(), Flags::A);
assert_eq!(iter.next().unwrap(), Flags::BC);
assert_eq!(iter.next(), None);
}
}