include any bits not corresponding to a flag in into_iter

bitflags · Mar 13, 2023 · 777d640 · 777d640
1 parent d1bd5e2
commit 777d640
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 3 deletions.
diff --git a/src/internal.rs b/src/internal.rs
@@ -20,7 +20,10 @@ macro_rules! __declare_internal_bitflags {
             bits: $T,
         }
 
-        $iter_vis struct $Iter($IterNames);
+        $iter_vis struct $Iter {
+            inner: $IterNames,
+            done: bool,
+        }
 
         $iter_names_vis struct $IterNames {
             idx: usize,
@@ -245,7 +248,10 @@ macro_rules! __impl_internal_bitflags {
 
             #[inline]
             pub const fn iter(&self) -> $Iter {
-                $Iter(self.iter_names())
+                $Iter {
+                    inner: self.iter_names(),
+                    done: false,
+                }
             }
 
             #[inline]
@@ -348,7 +354,22 @@ macro_rules! __impl_internal_bitflags {
             type Item = $BitFlags;
 
             fn next(&mut self) -> $crate::__private::core::option::Option<Self::Item> {
-                self.0.next().map(|(_, value)| value)
+                match self.inner.next().map(|(_, value)| value) {
+                    $crate::__private::core::option::Option::Some(value) => $crate::__private::core::option::Option::Some(value),
+                    $crate::__private::core::option::Option::None if !self.done => {
+                        self.done = true;
+
+                        // After iterating through valid names, if there are any bits left over
+                        // then return one final value that includes them. This makes `into_iter`
+                        // and `from_iter` roundtrip
+                        if self.inner.state != $InternalBitFlags::empty() {
+                            $crate::__private::core::option::Option::Some($BitFlags::from_bits_retain(self.inner.state.bits()))
+                        } else {
+                            $crate::__private::core::option::Option::None
+                        }
+                    },
+                    _ => $crate::__private::core::option::Option::None,
+                }
             }
         }
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -1498,6 +1498,17 @@ mod tests {
         assert_eq!(iter.next().unwrap(), ("ONE", Flags::ONE));
         assert_eq!(iter.next().unwrap(), ("THREE", Flags::THREE));
         assert_eq!(iter.next(), None);
+
+        let flags = Flags::from_bits_retain(0b1000_0000);
+        assert_eq!(flags.into_iter().count(), 1);
+        assert_eq!(flags.iter_names().count(), 0);
+    }
+
+    #[test]
+    fn into_iter_from_iter_roundtrip() {
+        let flags = Flags::ABC | Flags::from_bits_retain(0b1000_0000);
+
+        assert_eq!(flags, flags.into_iter().collect::<Flags>());
     }
 
     #[test]