Add DekuSize impl

* Add new derive DekuSize, which emits SIZE_BYTES and SIZE_BITS.
  This can be use to get the max bit/byte size of the parsed enum/struct
  so that your stack/heap can be properly allocated before usage.
* Generated for all primitive and std types we have, others such as Vec
  are impossible.

Closees #599, #307

Author: wcampbell0x2a

deku-derive/src/lib.rs

@@ -462,6 +462,17 @@ impl DekuData {
     fn emit_writer_checked(&self) -> Result<TokenStream, syn::Error> {
         emit_deku_write(self)
     }
+
+    /// Emit a size implementation
+    fn emit_size(&self) -> TokenStream {
+        self.emit_size_checked()
+            .unwrap_or_else(|e| e.to_compile_error())
+    }
+
+    /// Emit a size implementation, no compile_error
+    fn emit_size_checked(&self) -> Result<TokenStream, syn::Error> {
+        macros::deku_size::emit_deku_size(self)
+    }
 }
 
 /// Common variables from `DekuData` for `emit_enum` read/write functions
@@ -1241,6 +1252,15 @@ pub fn proc_deku_write(input: proc_macro::TokenStream) -> proc_macro::TokenStrea
     }
 }
 
+/// Entry function for `DekuSize` proc-macro
+#[proc_macro_derive(DekuSize, attributes(deku))]
+pub fn proc_deku_size(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    match DekuData::from_input(input.into()) {
+        Ok(data) => data.emit_size().into(),
+        Err(err) => err.into(),
+    }
+}
+
 fn is_not_deku(attr: &syn::Attribute) -> bool {
     attr.path()
         .get_ident()

deku-derive/src/macros/deku_size.rs

@@ -0,0 +1,175 @@
+use std::convert::TryFrom;
+
+use darling::ast::Data;
+use proc_macro2::TokenStream;
+use quote::quote;
+
+use crate::{DekuData, DekuDataEnum, DekuDataStruct, FieldData};
+
+pub(crate) fn emit_deku_size(input: &DekuData) -> Result<TokenStream, syn::Error> {
+    // NOTE: This is overkill, but keeps the same validating requirements of DekuRead. This ensures
+    // the required values are set so we calculate the correct size.
+    // TODO: We only care about validating and not the codegen. Split the two.
+    let _ = super::deku_read::emit_deku_read(input)?;
+
+    match &input.data {
+        Data::Enum(_) => emit_enum(input),
+        Data::Struct(_) => emit_struct(input),
+    }
+}
+
+/// Calculate the size of a collection of fields
+fn calculate_fields_size<'a>(
+    fields: impl IntoIterator<Item = &'a FieldData>,
+    crate_: &syn::Ident,
+) -> TokenStream {
+    let field_sizes = fields.into_iter().filter_map(|f| {
+        if !f.temp {
+            let field_type = &f.ty;
+
+            #[cfg(feature = "bits")]
+            if let Some(bits) = &f.bits {
+                return Some(quote! { (#bits) });
+            }
+
+            if let Some(bytes) = &f.bytes {
+                return Some(quote! { (#bytes) * 8 });
+            }
+
+            Some(quote! { <#field_type as ::#crate_::DekuSize>::SIZE_BITS })
+        } else {
+            None
+        }
+    });
+
+    quote! { 0 #(+ #field_sizes)* }
+}
+
+/// Add DekuSize trait bounds to where clause for fields that need them
+fn add_field_bounds<'a>(
+    where_clause: &mut Option<syn::WhereClause>,
+    fields: impl IntoIterator<Item = &'a FieldData>,
+    crate_: &syn::Ident,
+) {
+    for field in fields {
+        if !field.temp {
+            let field_type = &field.ty;
+            #[cfg(feature = "bits")]
+            let needs_bound = field.bits.is_none() && field.bytes.is_none();
+            #[cfg(not(feature = "bits"))]
+            let needs_bound = field.bytes.is_none();
+
+            if needs_bound {
+                let where_clause = where_clause.get_or_insert_with(|| syn::parse_quote! { where });
+                where_clause.predicates.push(syn::parse_quote! {
+                    #field_type: ::#crate_::DekuSize
+                });
+            }
+        }
+    }
+}
+
+/// Calculate the discriminant size for an enum
+fn calculate_discriminant_size(
+    input: &DekuData,
+    id: Option<&crate::Id>,
+    id_type: Option<&TokenStream>,
+    crate_: &syn::Ident,
+) -> TokenStream {
+    #[cfg(feature = "bits")]
+    if let Some(bits) = &input.bits {
+        return quote! { (#bits) };
+    }
+
+    if let Some(bytes) = &input.bytes {
+        return quote! { (#bytes) * 8 };
+    }
+
+    if id.is_some() {
+        return quote! { 0 };
+    }
+
+    if let Some(id_type) = id_type {
+        return quote! { <#id_type as ::#crate_::DekuSize>::SIZE_BITS };
+    }
+
+    // This unwrap is ensured by us calling DekuRead validation
+    let repr = &input.repr.unwrap();
+    let repr_type = TokenStream::from(*repr);
+    quote! { <#repr_type as ::#crate_::DekuSize>::SIZE_BITS }
+}
+
+fn emit_struct(input: &DekuData) -> Result<TokenStream, syn::Error> {
+    let crate_ = super::get_crate_name();
+
+    let DekuDataStruct {
+        imp: _,
+        wher: _,
+        ident,
+        fields,
+    } = DekuDataStruct::try_from(input)?;
+
+    let size_calculation = calculate_fields_size(fields.iter().copied(), &crate_);
+
+    let (imp_generics, ty_generics, where_clause) = input.generics.split_for_impl();
+
+    let mut where_clause = where_clause.cloned();
+    add_field_bounds(&mut where_clause, fields.iter().copied(), &crate_);
+
+    let tokens = quote! {
+        impl #imp_generics ::#crate_::DekuSize for #ident #ty_generics #where_clause {
+            const SIZE_BITS: usize = #size_calculation;
+        }
+    };
+
+    Ok(tokens)
+}
+
+fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
+    let crate_ = super::get_crate_name();
+
+    let DekuDataEnum {
+        imp: _,
+        wher: _,
+        variants,
+        ident,
+        id,
+        id_type,
+        id_args: _,
+    } = DekuDataEnum::try_from(input)?;
+
+    let discriminant_size = calculate_discriminant_size(input, id, id_type, &crate_);
+
+    let variant_sizes = variants
+        .iter()
+        .map(|variant| calculate_fields_size(variant.fields.iter(), &crate_));
+
+    let max_variant_size = quote! {
+        {
+            const fn const_max(a: usize, b: usize) -> usize {
+                if a > b { a } else { b }
+            }
+
+            let mut max = 0;
+            #(
+                max = const_max(max, #variant_sizes);
+            )*
+            max
+        }
+    };
+
+    let (imp_generics, ty_generics, where_clause) = input.generics.split_for_impl();
+
+    let mut where_clause = where_clause.cloned();
+    for variant in variants.iter() {
+        add_field_bounds(&mut where_clause, variant.fields.iter(), &crate_);
+    }
+
+    let tokens = quote! {
+        impl #imp_generics ::#crate_::DekuSize for #ident #ty_generics #where_clause {
+            const SIZE_BITS: usize = #discriminant_size + #max_variant_size;
+        }
+    };
+
+    Ok(tokens)
+}

deku-derive/src/macros/mod.rs

@@ -11,6 +11,7 @@ use syn::LitStr;
 use crate::Num;
 
 pub(crate) mod deku_read;
+pub(crate) mod deku_size;
 pub(crate) mod deku_write;
 
 #[cfg(feature = "proc-macro-crate")]

ensure_no_std/src/no_alloc_imports.rs

@@ -1,6 +1,6 @@
 use deku::prelude::*;
 
-#[derive(Debug, PartialEq, DekuRead, DekuWrite)]
+#[derive(Debug, PartialEq, DekuRead, DekuWrite, DekuSize)]
 struct DekuTest {
     field_a: u8,
     field_b: u8,
@@ -23,6 +23,7 @@ pub fn rw() {
     );
 
     // Test writing
-    let mut buf = [0; 20];
+    const BUF_SIZE: usize = DekuTest::SIZE_BYTES.unwrap();
+    let mut buf = [0; BUF_SIZE];
     let _val = val.to_slice(&mut buf).unwrap();
 }

src/impls/bool.rs

@@ -47,6 +47,10 @@ where
     }
 }
 
+impl crate::DekuSize for bool {
+    const SIZE_BITS: usize = 8;
+}
+
 #[cfg(test)]
 mod tests {
     use hexlit::hex;

src/impls/primitive.rs

@@ -1259,6 +1259,33 @@ ImplDekuTraitsBytesUnsigned!(f32, u32);
 ImplDekuTraitsUnsigned!(f64, u64);
 ImplDekuTraitsBytesUnsigned!(f64, u64);
 
+use crate::DekuSize;
+
+macro_rules! ImplDekuSize {
+    ($typ:ty) => {
+        impl DekuSize for $typ {
+            const SIZE_BITS: usize = core::mem::size_of::<$typ>() * 8;
+        }
+    };
+}
+
+ImplDekuSize!(u8);
+ImplDekuSize!(u16);
+ImplDekuSize!(u32);
+ImplDekuSize!(u64);
+ImplDekuSize!(u128);
+ImplDekuSize!(usize);
+
+ImplDekuSize!(i8);
+ImplDekuSize!(i16);
+ImplDekuSize!(i32);
+ImplDekuSize!(i64);
+ImplDekuSize!(i128);
+ImplDekuSize!(isize);
+
+ImplDekuSize!(f32);
+ImplDekuSize!(f64);
+
 #[cfg(feature = "std")]
 #[cfg(test)]
 mod tests {

src/impls/slice.rs

@@ -94,6 +94,10 @@ where
     }
 }
 
+impl<T: crate::DekuSize, const N: usize> crate::DekuSize for [T; N] {
+    const SIZE_BITS: usize = T::SIZE_BITS * N;
+}
+
 #[cfg(feature = "std")]
 #[cfg(test)]
 mod tests {

src/impls/tuple.rs

@@ -81,6 +81,26 @@ ImplDekuTupleTraits! { A, B, C, D, E, F, G, H, I, }
 ImplDekuTupleTraits! { A, B, C, D, E, F, G, H, I, J, }
 ImplDekuTupleTraits! { A, B, C, D, E, F, G, H, I, J, K, }
 
+macro_rules! ImplDekuSizeTuple {
+    ( $($T:ident,)+ ) => {
+        impl<$($T: crate::DekuSize),+> crate::DekuSize for ($($T,)+) {
+            const SIZE_BITS: usize = 0 $(+ $T::SIZE_BITS)+;
+        }
+    };
+}
+
+ImplDekuSizeTuple! { A, }
+ImplDekuSizeTuple! { A, B, }
+ImplDekuSizeTuple! { A, B, C, }
+ImplDekuSizeTuple! { A, B, C, D, }
+ImplDekuSizeTuple! { A, B, C, D, E, }
+ImplDekuSizeTuple! { A, B, C, D, E, F, }
+ImplDekuSizeTuple! { A, B, C, D, E, F, G, }
+ImplDekuSizeTuple! { A, B, C, D, E, F, G, H, }
+ImplDekuSizeTuple! { A, B, C, D, E, F, G, H, I, }
+ImplDekuSizeTuple! { A, B, C, D, E, F, G, H, I, J, }
+ImplDekuSizeTuple! { A, B, C, D, E, F, G, H, I, J, K, }
+
 #[cfg(all(feature = "alloc", feature = "std"))]
 #[cfg(test)]
 mod tests {