bunch of write stuff

bitbuffer_derive/src/discriminant.rs

@@ -0,0 +1,36 @@
+use syn::{Expr, Lit, Variant};
+use syn_util::get_attribute_value;
+
+pub enum Discriminant {
+    Int(usize),
+    Default,
+    Wildcard,
+}
+
+impl From<Lit> for Discriminant {
+    fn from(lit: Lit) -> Self {
+        match lit {
+            Lit::Int(lit) => Discriminant::Int(lit.base10_parse::<usize>().unwrap()),
+            Lit::Str(lit) => match lit.value().as_str() {
+                "_" => Discriminant::Wildcard,
+                _ => panic!("discriminant is required to be an integer literal or \"_\""),
+            },
+            _ => panic!("discriminant is required to be an integer literal or \"_\""),
+        }
+    }
+}
+
+impl From<&Variant> for Discriminant {
+    fn from(variant: &Variant) -> Self {
+        variant
+            .discriminant
+            .as_ref()
+            .map(|(_, expr)| match expr {
+                Expr::Lit(expr_lit) => expr_lit.lit.clone(),
+                _ => panic!("discriminant is required to be an integer literal"),
+            })
+            .or_else(|| get_attribute_value(&variant.attrs, &["discriminant"]))
+            .map(Discriminant::from)
+            .unwrap_or(Discriminant::Default)
+    }
+}

bitbuffer_derive/src/lib.rs

@@ -131,14 +131,20 @@
 //!     stream: BitReadStream<'a, BigEndian>,
 //! }
 //! ```
+//!
+mod discriminant;
+mod write;
+
 extern crate proc_macro;
 
+use crate::write::derive_bitwrite_trait;
+use discriminant::Discriminant;
 use proc_macro2::{Span, TokenStream};
 use quote::{quote, quote_spanned};
 use syn::spanned::Spanned;
 use syn::{
     parse_macro_input, parse_quote, parse_str, Attribute, Data, DataStruct, DeriveInput, Expr,
-    Fields, GenericParam, Ident, Lit, LitStr, Path, Variant,
+    Fields, GenericParam, Ident, Lit, LitStr, Path,
 };
 use syn_util::get_attribute_value;
 
@@ -162,6 +168,31 @@ pub fn derive_bitread_sized(input: proc_macro::TokenStream) -> proc_macro::Token
     derive_bitread_trait(input, "BitReadSized".to_owned(), Some(extra_param))
 }
 
+/// See the [crate documentation](index.html) for details
+#[proc_macro_derive(
+    BitWrite,
+    attributes(size, size_bits, discriminant_bits, discriminant, endianness)
+)]
+pub fn derive_bitwrite(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    derive_bitwrite_trait(input, "BitWrite".into(), "write".into(), None)
+}
+
+//
+/// See the [crate documentation](index.html) for details
+#[proc_macro_derive(
+    BitWriteSized,
+    attributes(size, size_bits, discriminant_bits, discriminant, endianness)
+)]
+pub fn derive_bitwrite_sized(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let extra_param = parse_str::<TokenStream>(", input_size: usize").unwrap();
+    derive_bitwrite_trait(
+        input,
+        "BitWriteSized".into(),
+        "write_sized".into(),
+        Some(extra_param),
+    )
+}
+
 fn derive_bitread_trait(
     input: proc_macro::TokenStream,
     trait_name: String,
@@ -514,37 +545,3 @@ fn get_field_size(attrs: &[Attribute], span: Span) -> Option<TokenStream> {
             })
         })
 }
-
-enum Discriminant {
-    Int(usize),
-    Default,
-    Wildcard,
-}
-
-impl From<Lit> for Discriminant {
-    fn from(lit: Lit) -> Self {
-        match lit {
-            Lit::Int(lit) => Discriminant::Int(lit.base10_parse::<usize>().unwrap()),
-            Lit::Str(lit) => match lit.value().as_str() {
-                "_" => Discriminant::Wildcard,
-                _ => panic!("discriminant is required to be an integer literal or \"_\""),
-            },
-            _ => panic!("discriminant is required to be an integer literal or \"_\""),
-        }
-    }
-}
-
-impl From<&Variant> for Discriminant {
-    fn from(variant: &Variant) -> Self {
-        variant
-            .discriminant
-            .as_ref()
-            .map(|(_, expr)| match expr {
-                Expr::Lit(expr_lit) => expr_lit.lit.clone(),
-                _ => panic!("discriminant is required to be an integer literal"),
-            })
-            .or_else(|| get_attribute_value(&variant.attrs, &["discriminant"]))
-            .map(Discriminant::from)
-            .unwrap_or(Discriminant::Default)
-    }
-}

bitbuffer_derive/src/write.rs

@@ -0,0 +1,275 @@
+use crate::discriminant::Discriminant;
+use crate::size;
+use proc_macro2::{Span, TokenStream};
+use quote::{quote, quote_spanned};
+use syn::spanned::Spanned;
+use syn::{
+    parse_macro_input, parse_quote, parse_str, Attribute, Data, DataStruct, DeriveInput, Expr,
+    Fields, GenericParam, Ident, Index, Lit, Member, Path,
+};
+use syn_util::get_attribute_value;
+
+pub fn derive_bitwrite_trait(
+    input: proc_macro::TokenStream,
+    trait_name: String,
+    write_method_name: String,
+    extra_param: Option<TokenStream>,
+) -> proc_macro::TokenStream {
+    let input: DeriveInput = parse_macro_input!(input as DeriveInput);
+
+    let name = &input.ident;
+
+    let endianness = get_attribute_value(&input.attrs, &["endianness"]);
+    let mut trait_generics = input.generics.clone();
+    // we need these separate generics to only add out Endianness param to the 'impl'
+    let (_, ty_generics, where_clause) = input.generics.split_for_impl();
+    let lifetime: Option<&GenericParam> = trait_generics
+        .params
+        .iter()
+        .find(|param| matches!(param, GenericParam::Lifetime(_)));
+    let _lifetime = match lifetime {
+        Some(GenericParam::Lifetime(lifetime)) => lifetime.lifetime.clone(),
+        _ => {
+            // trait_generics.params.push(parse_quote!('a));
+            parse_quote!('a)
+        }
+    };
+    if endianness.is_none() {
+        trait_generics
+            .params
+            .push(parse_quote!(_E: ::bitbuffer::Endianness));
+    }
+    let (impl_generics, _, _) = trait_generics.split_for_impl();
+    let span = input.span();
+
+    let _size = size(
+        input.data.clone(),
+        &name,
+        &input.attrs,
+        extra_param.is_some(),
+    );
+    let parsed = write(input.data.clone(), &name, &input.attrs);
+    let _parsed_unchecked = write(input.data.clone(), &name, &input.attrs);
+
+    let endianness_placeholder = endianness.unwrap_or_else(|| "_E".to_owned());
+    let trait_def_str = format!("::bitbuffer::{}<{}>", trait_name, &endianness_placeholder);
+    let trait_def = parse_str::<Path>(&trait_def_str).expect("trait");
+
+    let endianness_ident = Ident::new(&endianness_placeholder, span);
+
+    let _size_extra_param = if extra_param.is_some() {
+        Some(quote!(input_size: usize))
+    } else {
+        None
+    };
+
+    let _extra_param_call = if extra_param.is_some() {
+        Some(quote!(input_size,))
+    } else {
+        None
+    };
+
+    let write_method = Ident::new(&write_method_name, span);
+
+    let expanded = quote! {
+        impl #impl_generics #trait_def for #name #ty_generics #where_clause {
+            fn #write_method(&self, __target__stream: &mut ::bitbuffer::BitWriteStream<#endianness_ident>#extra_param) -> ::bitbuffer::Result<()> {
+                #parsed
+            }
+        }
+    };
+
+    // panic!("{}", TokenStream::to_string(&expanded));
+
+    proc_macro::TokenStream::from(expanded)
+}
+
+fn write(data: Data, struct_name: &Ident, attrs: &[Attribute]) -> TokenStream {
+    let span = struct_name.span();
+
+    match data {
+        Data::Struct(DataStruct { fields, .. }) => {
+            let expand = fields.iter().enumerate().map(|(i, field)| {
+                let name = field
+                    .ident
+                    .clone()
+                    .unwrap_or_else(|| Ident::new(&format!("__{}", i), span));
+                let member = field.ident.clone().map(Member::Named).unwrap_or_else(|| {
+                    Member::Unnamed(Index {
+                        index: i as u32,
+                        span: field.span(),
+                    })
+                });
+                quote_spanned! { field.span() =>
+                    let #name = &self.#member;
+                }
+            });
+
+            let writes = fields.iter().enumerate().map(|(i, f)| {
+                // Get attributes `#[..]` on each field
+                let size = get_field_size(&f.attrs, f.span());
+                let span = f.span();
+                let name = f
+                    .ident
+                    .clone()
+                    .unwrap_or_else(|| Ident::new(&format!("__{}", i), span));
+                match size {
+                    Some(size) => {
+                        quote_spanned! { span =>
+                            {
+                                let _size: usize = #size;
+                                __target__stream.write_sized(#name, _size)?;
+                            }
+                        }
+                    }
+                    None => {
+                        quote_spanned! { span => {
+                            __target__stream.write(#name)?;
+                        }}
+                    }
+                }
+            });
+
+            quote_spanned! {span=>
+                #(#expand)*
+                #(#writes)*
+                Ok(())
+            }
+        }
+        Data::Enum(data) => {
+            let discriminant_bits: u64 = get_attribute_value(attrs, &["discriminant_bits"]).expect(
+                "'discriminant_bits' attribute is required when deriving `BinRead` for enums",
+            );
+
+            let mut last_discriminant = -1;
+
+            let max_discriminant = data
+                .variants
+                .iter()
+                .map(|variant| match Discriminant::from(variant) {
+                    Discriminant::Int(discriminant) => {
+                        last_discriminant = discriminant as isize;
+                        discriminant
+                    }
+                    Discriminant::Wildcard => 0,
+                    Discriminant::Default => {
+                        let new_discriminant = (last_discriminant + 1) as usize;
+                        last_discriminant += 1;
+                        new_discriminant
+                    }
+                })
+                .max()
+                .unwrap_or(0);
+
+            let mut last_discriminant = -1;
+
+            let discriminant_value = data.variants.iter().map(|variant| {
+                let span = variant.span();
+                let variant_name = &variant.ident;
+
+                let discriminant_token: TokenStream = match Discriminant::from(variant) {
+                    Discriminant::Int(discriminant) => {
+                        last_discriminant = discriminant as isize;
+                        quote_spanned! { span => #discriminant }
+                    }
+                    Discriminant::Wildcard => {
+                        let free_discriminant = max_discriminant + 1;
+                        quote_spanned! { span => #free_discriminant }
+                    }
+                    Discriminant::Default => {
+                        let new_discriminant = (last_discriminant + 1) as usize;
+                        last_discriminant += 1;
+                        quote_spanned! { span => #new_discriminant }
+                    }
+                };
+
+                match &variant.fields {
+                    Fields::Unit => quote_spanned! {span =>
+                        #struct_name::#variant_name => #discriminant_token
+                    },
+                    Fields::Unnamed(_f) => {
+                        quote_spanned! { span =>
+                            #struct_name::#variant_name(_) => #discriminant_token
+                        }
+                    }
+                    _ => unimplemented!(),
+                }
+            });
+
+            let write_inner = data.variants.iter().map(|variant| {
+                let span = variant.span();
+                let variant_name = &variant.ident;
+
+                match &variant.fields {
+                    Fields::Unit => quote_spanned! {span =>
+                        #struct_name::#variant_name => {},
+                    },
+                    Fields::Unnamed(f) => {
+                        let size = get_field_size(&variant.attrs, f.span());
+                        match size {
+                            Some(size) => {
+                                quote_spanned! { span =>
+                                    #struct_name::#variant_name(inner) => {
+                                        let size:usize = #size;
+                                        __target__stream.write_sized(inner, size)?;
+                                    }
+                                }
+                            }
+                            None => {
+                                quote_spanned! { span =>
+                                    #struct_name::#variant_name(inner) => { __target__stream.write(inner)?; }
+                                }
+                            }
+                        }
+                    }
+                    _ => unimplemented!(),
+                }
+            });
+
+            let span = data.enum_token.span();
+
+            quote_spanned! {span=>
+                let discriminant = match &self {
+                    #(#discriminant_value),*
+                };
+                __target__stream.write_int(discriminant as usize, #discriminant_bits as usize)?;
+                match &self {
+                    #(#write_inner)*
+                }
+                Ok(())
+            }
+        }
+        _ => unimplemented!(),
+    }
+}
+
+fn get_field_size(attrs: &[Attribute], span: Span) -> Option<TokenStream> {
+    get_attribute_value(attrs, &["size"])
+        .map(|size_lit| match size_lit {
+            Lit::Int(size) => {
+                quote_spanned! {span =>
+                    #size
+                }
+            }
+            Lit::Str(size_field) => {
+                let size = parse_str::<Expr>(&size_field.value()).expect("size");
+                if size_field.value() == "input_size" {
+                    quote_spanned! {span =>
+                        (#size) as usize
+                    }
+                } else {
+                    quote_spanned! {span =>
+                        (*#size) as usize
+                    }
+                }
+            }
+            _ => panic!("Unsupported value for size attribute"),
+        })
+        .or_else(|| {
+            get_attribute_value::<Lit>(attrs, &["size_bits"]).map(|_| {
+                quote_spanned! {span =>
+                    compile_error!("#[size_bits] is not supported when deriving BitWrite or BitWriteSized")
+                }
+            })
+        })
+}

bitbuffer_derive/tests/expand.rs

@@ -1,10 +1,15 @@
 #![allow(dead_code)]
 #![allow(unreachable_patterns)]
+#![allow(unused_imports)]
 
-use bitbuffer_derive::BitRead;
+use bitbuffer::{BitReadStream, Endianness};
+use bitbuffer_derive::{BitWrite, BitWriteSized};
 
-#[derive(BitRead)]
+#[derive(BitWrite)]
 struct TestStruct {
     foo: u8,
     str: String,
 }
+
+#[derive(BitWrite)]
+struct UnnamedSize(u8, #[size = 5] String, bool);

bitbuffer_derive/tests/write.rs

@@ -0,0 +1,262 @@
+#![allow(dead_code)]
+#![allow(unreachable_patterns)]
+
+use bitbuffer::{
+    BigEndian, BitReadBuffer, BitReadSized, BitReadStream, BitWriteStream, Endianness, LittleEndian,
+};
+use bitbuffer_derive::{BitRead, BitWrite, BitWriteSized};
+
+#[derive(BitWrite, PartialEq, Debug)]
+struct TestStruct {
+    foo: u8,
+    str: String,
+    #[size = 2]
+    truncated: String,
+    bar: u16,
+    float: f32,
+    #[size = 3]
+    asd: u8,
+    #[size = "asd"]
+    previous_field: u8,
+}
+
+#[test]
+fn test_read_struct() {
+    let float: [u8; 4] = 12.5f32.to_bits().to_le_bytes();
+    let bytes = vec![
+        12,
+        'h' as u8,
+        'e' as u8,
+        'l' as u8,
+        'l' as u8,
+        'o' as u8,
+        0,
+        'f' as u8,
+        'o' as u8,
+        'o' as u8,
+        0,
+        float[0],
+        float[1],
+        float[2],
+        float[3],
+        0b1010_0101,
+    ];
+    let val = TestStruct {
+        foo: 12,
+        str: "hello".to_owned(),
+        truncated: "fo".to_owned(),
+        bar: 'o' as u16,
+        float: 12.5,
+        asd: 0b101,
+        previous_field: 0b1010_0,
+    };
+    let mut stream = BitWriteStream::new(LittleEndian);
+    stream.write(&val).unwrap();
+    assert_eq!(bytes, stream.finish());
+}
+
+#[derive(BitWrite, PartialEq, Debug)]
+#[discriminant_bits = 2]
+enum TestBareEnum {
+    Foo,
+    Bar,
+    Asd = 3,
+}
+
+#[test]
+fn test_read_bare_enum() {
+    let bytes = vec![0b1100_0100];
+    let mut stream = BitWriteStream::new(BigEndian);
+    stream.write(&TestBareEnum::Asd).unwrap();
+    stream.write(&TestBareEnum::Foo).unwrap();
+    stream.write(&TestBareEnum::Bar).unwrap();
+
+    assert_eq!(bytes, stream.finish());
+}
+
+#[derive(BitWrite, BitRead, PartialEq, Debug)]
+#[discriminant_bits = 2]
+enum TestUnnamedFieldEnum {
+    #[size = 5]
+    Foo(i8),
+    Bar(bool),
+    #[discriminant = 3]
+    Asd(u8),
+}
+
+#[test]
+fn test_read_unnamed_field_enum() {
+    let bytes = vec![0b1100_0110, 0b1000_0110, 0b1011_0000];
+    let mut stream = BitWriteStream::new(BigEndian);
+    stream
+        .write(&TestUnnamedFieldEnum::Asd(0b_00_0110_10))
+        .unwrap();
+    assert_eq!(10, stream.bit_len());
+
+    stream.write(&TestUnnamedFieldEnum::Foo(0b0110_1)).unwrap();
+    assert_eq!(17, stream.bit_len());
+
+    stream.write(&TestUnnamedFieldEnum::Bar(true)).unwrap();
+
+    let result = stream.finish();
+    let mut read = BitReadStream::<BigEndian>::from(result.as_slice());
+
+    assert_eq!(
+        TestUnnamedFieldEnum::Asd(0b_00_0110_10),
+        read.read().unwrap()
+    );
+    assert_eq!(TestUnnamedFieldEnum::Foo(0b11_0_1), read.read().unwrap());
+    assert_eq!(TestUnnamedFieldEnum::Bar(true), read.read().unwrap());
+
+    assert_eq!(bytes, result);
+}
+
+#[derive(BitWriteSized, BitReadSized, PartialEq, Debug)]
+struct TestStructSized {
+    foo: u8,
+    #[size = "input_size"]
+    string: String,
+    #[size = "input_size"]
+    int: u8,
+}
+
+#[test]
+fn test_read_struct_sized() {
+    let bytes = vec![12, 'h' as u8, 'e' as u8, 'l' as u8, 0b1000_0000];
+    let mut stream = BitWriteStream::new(BigEndian);
+    let val = TestStructSized {
+        foo: 12,
+        string: "hel".to_owned(),
+        int: 4,
+    };
+    stream.write_sized(&val, 3).unwrap();
+    let result = stream.finish();
+    let mut read = BitReadStream::<BigEndian>::from(result.as_slice());
+
+    assert_eq!(val, read.read_sized(3).unwrap());
+
+    assert_eq!(bytes, result);
+}
+
+#[derive(BitWriteSized, PartialEq, Debug)]
+#[discriminant_bits = 2]
+enum TestUnnamedFieldEnumSized {
+    #[size = 5]
+    Foo(i8),
+    Bar(bool),
+    #[discriminant = 3]
+    #[size = "input_size"]
+    Asd(u8),
+}
+
+#[test]
+fn test_read_unnamed_field_enum_sized() {
+    let bytes = vec![0b1100_0110];
+    let mut stream = BitWriteStream::new(BigEndian);
+    stream
+        .write_sized(&TestUnnamedFieldEnumSized::Asd(0b_00_0110), 6)
+        .unwrap();
+    assert_eq!(bytes, stream.finish());
+}
+
+#[derive(BitWrite, PartialEq, Debug)]
+struct TestStruct2 {
+    size: u8,
+    #[size = "size * 2"]
+    str: String,
+}
+
+#[test]
+fn test_read_struct2() {
+    let bytes = vec![
+        0b0000_0101,
+        'h' as u8,
+        'e' as u8,
+        'l' as u8,
+        'l' as u8,
+        'o' as u8,
+        ' ' as u8,
+        'w' as u8,
+        'o' as u8,
+        'r' as u8,
+        'l' as u8,
+    ];
+    let mut stream = BitWriteStream::new(BigEndian);
+    stream
+        .write(&TestStruct2 {
+            size: 5,
+            str: "hello worl".to_owned(),
+        })
+        .unwrap();
+    assert_eq!(bytes, stream.finish());
+}
+
+#[derive(BitWrite)]
+#[endianness = "E"]
+struct TestStruct3<'a, E: Endianness> {
+    size: u8,
+    #[size = "size"]
+    stream: BitReadStream<'a, E>,
+}
+
+#[test]
+fn test_read_struct3() {
+    let bytes = vec![0b0000_0101, 0b1010_1000];
+    let mut stream = BitWriteStream::new(BigEndian);
+    let mut inner = BitReadStream::from(BitReadBuffer::new(&[0b1010_1010], BigEndian));
+
+    let inner = inner.read_bits(5).unwrap();
+
+    let val: TestStruct3<BigEndian> = TestStruct3 {
+        size: 5,
+        stream: inner,
+    };
+    stream.write(&val).unwrap();
+    assert_eq!(bytes, stream.finish());
+}
+
+#[derive(BitWrite, PartialEq, Debug)]
+#[discriminant_bits = 2]
+enum TestEnumRest {
+    Foo,
+    Bar,
+    #[discriminant = "_"]
+    Asd,
+}
+
+#[test]
+fn test_read_rest_enum() {
+    let bytes = vec![0b1000_0110];
+    let mut stream = BitWriteStream::new(BigEndian);
+
+    stream.write(&TestEnumRest::Asd).unwrap();
+    stream.write(&TestEnumRest::Foo).unwrap();
+    stream.write(&TestEnumRest::Bar).unwrap();
+    stream.write(&TestEnumRest::Asd).unwrap();
+
+    assert_eq!(bytes, stream.finish());
+}
+
+#[derive(BitWrite, PartialEq, Debug)]
+struct UnnamedSize(u8, #[size = 5] String, bool);
+
+fn test_unnamed_struct() {
+    let bytes = vec![
+        12, 'h' as u8, 'e' as u8, 'l' as u8, 'l' as u8, 'o' as u8, 0, 0, 0, 0, 0, 0,
+    ];
+    let mut stream = BitWriteStream::new(LittleEndian);
+    stream
+        .write(&UnnamedSize(12, "hello".to_string(), false))
+        .unwrap();
+
+    assert_eq!(bytes, stream.finish());
+}
+
+#[derive(BitWrite, PartialEq, Debug)]
+struct EmptyStruct;
+
+fn test_empty_struct() {
+    let mut stream = BitWriteStream::new(LittleEndian);
+    stream.write(&EmptyStruct).unwrap();
+    assert_eq!(Vec::<u8>::new(), stream.finish());
+}