Implement Iterator for bitflags

src/bitflags_trait.rs

@@ -6,6 +6,8 @@ pub trait ImplementedByBitFlagsMacro {}
 /// It should not be implemented manually.
 pub trait BitFlags: ImplementedByBitFlagsMacro {
     type Bits;
+    const NUM_FLAGS: usize;
+
     /// Returns an empty set of flags.
     fn empty() -> Self;
     /// Returns the set containing all flags.

src/lib.rs

@@ -559,18 +559,41 @@ macro_rules! __impl_bitflags {
             /// representation contains bits that do not correspond to a flag.
             #[inline]
             pub const fn from_bits(bits: $T) -> $crate::_core::option::Option<Self> {
-                if (bits & !Self::all().bits()) == 0 {
-                    $crate::_core::option::Option::Some(Self { bits })
-                } else {
-                    $crate::_core::option::Option::None
+                if bits == 0 {
+                    return Some(Self{ bits })
                 }
+
+                $(
+                    #[allow(unused_doc_comments, unused_attributes)]
+                    $(#[$attr $($args)*])*
+                    if bits & Self::$Flag.bits == Self::$Flag.bits {
+                        return Some(Self{ bits })
+                    }
+                )*
+
+                None
             }
 
             /// Convert from underlying bit representation, dropping any bits
             /// that do not correspond to flags.
             #[inline]
             pub const fn from_bits_truncate(bits: $T) -> Self {
-                Self { bits: bits & Self::all().bits }
+                if bits == 0 {
+                    return Self{ bits }
+                }
+
+                #[allow(unused_mut)]
+                let mut truncated = 0;
+
+                $(
+                    #[allow(unused_doc_comments, unused_attributes)]
+                    $(#[$attr $($args)*])*
+                    if bits & Self::$Flag.bits == Self::$Flag.bits {
+                        truncated |= Self::$Flag.bits
+                    }
+                )*
+
+                Self { bits: truncated }
             }
 
             /// Convert from underlying bit representation, preserving all
@@ -730,6 +753,17 @@ macro_rules! __impl_bitflags {
                 Self::from_bits_truncate(!self.bits)
             }
 
+
+            /// Returns an iterator over all the flags in this set.
+            pub const fn iter(self) -> $crate::Iter<Self, { <Self as $crate::BitFlags>::NUM_FLAGS }> {
+                $crate::Iter::new(self, [
+                    $(
+                        #[allow(unused_doc_comments, unused_attributes)]
+                        $(#[$attr $($args)*])*
+                        Self::$Flag,
+                    )*
+                ])
+            }
         }
 
         impl $crate::_core::ops::BitOr for $BitFlags {
@@ -832,6 +866,20 @@ macro_rules! __impl_bitflags {
 
         impl $crate::BitFlags for $BitFlags {
             type Bits = $T;
+            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
+            };
 
             fn empty() -> Self {
                 $BitFlags::empty()
@@ -1002,6 +1050,39 @@ macro_rules! __impl_bitflags {
     };
 }
 
+pub struct Iter<T, const LEN: usize> {
+    bitflags: T,
+    possible: [T;LEN],
+    len: usize,
+}
+
+impl<T, const LEN: usize> Iter<T, LEN> {
+    pub const fn new(bitflags: T, possible: [T;LEN]) -> Iter<T,LEN> {
+        Iter{ bitflags, possible, len: LEN }
+    }
+}
+
+impl<T: BitFlags + Copy, const LEN: usize> Iterator for Iter<T, LEN>  {
+    type Item = T;
+
+    fn next(&mut self) -> Option<T> {
+        if self.bitflags.is_empty() || LEN == 0 || self.len == 0 {
+            None
+        }else{
+            for (pos, flag) in self.possible.iter().copied().take(self.len).enumerate() {
+                if self.bitflags.contains(flag) {
+                    self.bitflags.remove(flag);
+                    self.len -= 1;
+                    self.possible.swap(pos, self.len);
+                    return Some(flag)
+                }
+            }
+
+            None
+        }
+    }
+}
+
 #[cfg(feature = "example_generated")]
 pub mod example_generated;
 
@@ -1789,12 +1870,87 @@ mod tests {
     }
 
     bitflags! {
-        #[derive(serde::Serialize, serde::Deserialize)]
+        #[derive(serde_derive::Serialize, serde_derive::Deserialize)]
         struct SerdeFlags: u32 {
             const A = 1;
             const B = 2;
             const C = 4;
             const D = 8;
         }
     }
-}
+
+    #[test]
+    fn test_iter() {
+        bitflags! {
+            struct Flags: u32 {
+                const ONE  = 0b001;
+                const TWO  = 0b010;
+                const THREE = 0b100;
+            }
+        }
+
+        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::THREE);
+        assert_eq!(iter.next().unwrap(), Flags::TWO);
+        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);
+    }
+
+    #[test]
+    fn test_from_bits_edge_cases() {
+        bitflags! {
+            struct Flags: u8 {
+                const A = 0b00000001;
+                const BC = 0b00000110;
+            }
+        }
+
+
+        let flags = Flags::from_bits(0b00000100);
+        assert!(flags.is_none());
+    }
+
+    #[test]
+    fn test_from_bits_truncate_edge_cases() {
+        bitflags! {
+            struct Flags: u8 {
+                const A = 0b00000001;
+                const BC = 0b00000110;
+            }
+        }
+
+
+        let flags = Flags::from_bits_truncate(0b00000100);
+        assert!(flags.is_empty());
+    }
+}