add config debug_as_vec

Lukas3674 · Nov 3, 2023 · cf530f8 · cf530f8
1 parent 79c43f2
commit cf530f8
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 5 deletions.
diff --git a/src/parser.rs b/src/parser.rs
@@ -28,13 +28,16 @@ pub fn parse(attr: TokenStream, mut item: ItemEnum) -> Result<TokenStream> {
     }
 
     let mut all_flags = Vec::with_capacity(flags_amount);
+    let mut all_flags_names = Vec::with_capacity(flags_amount);
 
     let mut i: usize = 0;
-    let flags = item.variants.iter().map(|v| {
+    let mut flags = Vec::with_capacity(flags_amount);
+    for v in item.variants.iter() {
         let v_attrs = &v.attrs;
         let v_ident = &v.ident;
 
         all_flags.push(quote::quote!(Self::#v_ident));
+        all_flags_names.push(quote::quote!(stringify!(#v_ident)));
 
         let expr = if let Some((_, expr)) = v.discriminant.as_ref() {
             quote::quote!(#expr)
@@ -48,25 +51,52 @@ pub fn parse(attr: TokenStream, mut item: ItemEnum) -> Result<TokenStream> {
             let i_ident = Ident::new(&format!("Inverted{}", v_ident), v_ident.span());
 
             all_flags.push(quote::quote!(Self::#i_ident));
+            all_flags_names.push(quote::quote!(stringify!(#v_ident)));
 
             quote::quote!(
                 #(#v_attrs)*
                 #vis const #i_ident: #ident = Self { bits: (#expr) ^ !0 };
             )
         }).into_iter();
 
-        quote::quote!(
+        flags.push(quote::quote!(
             #(#v_attrs)*
             #vis const #v_ident: #ident = Self { bits: #expr };
 
             #(#i_flag)*
-        )
-    });
+        ))
+    }
+
+    let debug_impl = if config.debug_as_vec {
+        quote::quote! {
+            impl core::fmt::Debug for #ident {
+                fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
+                    let mut matches = Vec::new();
+
+                    #(if self.contains(#all_flags) {
+                        matches.push(#all_flags_names);
+                    })*
+
+                    write!(f, "[{}]", matches.join(", "))
+                }
+            }
+        }
+    } else { 
+        quote::quote! {
+            impl core::fmt::Debug for #ident {
+                fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
+                    f.debug_struct(stringify!(#ident))
+                        .field("bits", &self.bits)
+                        .finish()
+                }
+            }
+        }
+    };
 
     Ok(TokenStream::from(quote::quote! {
         #(#attrs)*
         #[repr(transparent)]
-        #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+        #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
         #vis struct #ident {
             bits: #typ,
         }
@@ -248,6 +278,8 @@ pub fn parse(attr: TokenStream, mut item: ItemEnum) -> Result<TokenStream> {
             }
         }
 
+        #debug_impl
+
         impl core::fmt::Binary for #ident {
             #[inline]
             fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
@@ -294,12 +326,14 @@ fn parse_typ(attr: TokenStream) -> Result<Ident> {
 
 struct Config {
     inverted_flags: bool,
+    debug_as_vec: bool
 }
 
 impl Config {
     fn new() -> Self {
         Self {
             inverted_flags: false,
+            debug_as_vec: false
         }
     }
 }
@@ -311,6 +345,7 @@ impl Parse for Config {
         for arg in args {
             match arg.to_string().as_str() {
                 "inverted_flags" => config.inverted_flags = true,
+                "debug_as_vec" => config.debug_as_vec = true,
                 _ => return Err(Error::new_spanned(arg, "unknown config option")),
             }
         }

diff --git a/tests/tests.rs b/tests/tests.rs
@@ -335,4 +335,23 @@ mod tests {
         assert_eq!(BitmaskCustomTyped::InvertedFlag123, !0b111);
         assert_eq!(BitmaskCustomTyped::InvertedFlag4, !0b1000);
     }
+
+    #[test]
+    fn test_debug_as_vec() {
+        #[bitmask]
+        #[bitmask_config(debug_as_vec)]
+        pub enum BitmaskDebugAsVec {
+            Flag1,
+            Flag2,
+            Flag12 = Self::Flag1.or(Self::Flag2).bits,
+            Flag3,
+        }
+
+        assert_eq!(format!("{:?}", BitmaskDebugAsVec::none()), "[]");
+        assert_eq!(format!("{:?}", BitmaskDebugAsVec::Flag1), "[Flag1]");
+        assert_eq!(format!("{:?}", BitmaskDebugAsVec::Flag2), "[Flag2]");
+        assert_eq!(format!("{:?}", BitmaskDebugAsVec::Flag12), "[Flag1, Flag2, Flag12]");
+        assert_eq!(format!("{:?}", BitmaskDebugAsVec::Flag3), "[Flag3]");
+        assert_eq!(format!("{:?}", BitmaskDebugAsVec::full()), "[Flag1, Flag2, Flag12, Flag3]");
+    }
 }