Merge pull request #278 from arturoc/iterators

Iterator over all the enabled options
bitflags · Apr 28, 2022 · 0141a07 · 0141a07
2 parents 186a7bb + 6656792
commit 0141a07
Showing 1 changed file with 145 additions and 0 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -730,6 +730,99 @@ 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
+                    }
+                })
+            }
+
+            /// Returns an iterator over all the flags names as &'static str in this set.
+            pub fn iter_names(mut self) -> impl Iterator<Item = &'static str> {
+                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,
+                    )*
+                ];
+                #[allow(unused_doc_comments, unused_attributes)]
+                const OPTIONS_NAMES: [&'static str; NUM_FLAGS] = [
+                    $(
+                        $(#[$attr $($args)*])*
+                        $crate::_core::stringify!($Flag),
+                    )*
+                ];
+                let mut start = 0;
+
+                $crate::_core::iter::from_fn(move || {
+                    if self.is_empty() || NUM_FLAGS == 0 {
+                        None
+                    }else{
+                        for (flag, flag_name) in OPTIONS[start..NUM_FLAGS].iter().copied()
+                            .zip(OPTIONS_NAMES[start..NUM_FLAGS].iter().copied())
+                        {
+                            start += 1;
+                            if self.contains(flag) {
+                                self.remove(flag);
+                                return Some(flag_name)
+                            }
+                        }
+
+                        None
+                    }
+                })
+            }
+
         }
 
         impl $crate::_core::ops::BitOr for $BitFlags {
@@ -1797,4 +1890,56 @@ mod tests {
             const D = 8;
         }
     }
+
+    #[test]
+    fn test_iter() {
+        bitflags! {
+            struct Flags: u32 {
+                const ONE  = 0b001;
+                const TWO  = 0b010;
+                const THREE = 0b100;
+                #[cfg(windows)]
+                const FOUR_WIN = 0b1000;
+                #[cfg(unix)]
+                const FOUR_UNIX = 0b10000;
+            }
+        }
+
+        let flags = Flags::all();
+        assert_eq!(flags.iter().count(), 4);
+        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().unwrap(), Flags::FOUR_UNIX);
+        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);
+    }
 }