Draft: Improve DekuUpdate

- Allow DekuUpdate::update to have ctx
- Add call_update, to call `update()` on self.
- Add update_custom to call custom function

Author: wcampbell0x2a

deku-derive/src/lib.rs

@@ -133,6 +133,8 @@ struct DekuData {
     /// default context passed to the field
     ctx_default: Option<Punctuated<syn::Expr, syn::token::Comma>>,
 
+    update_ctx: Option<syn::punctuated::Punctuated<syn::FnArg, syn::token::Comma>>,
+
     /// A magic value that must appear at the start of this struct/enum's data
     magic: Option<syn::LitByteStr>,
 
@@ -209,6 +211,7 @@ impl DekuData {
             endian: receiver.endian,
             ctx: receiver.ctx,
             ctx_default: receiver.ctx_default,
+            update_ctx: receiver.update_ctx,
             magic: receiver.magic,
             id: receiver.id,
             id_type: receiver.id_type?,
@@ -455,6 +458,13 @@ struct FieldData {
     /// map field when updating struct
     update: Option<TokenStream>,
 
+    /// map field when updating struct
+    update_custom: Option<TokenStream>,
+
+    update_ctx: Option<Punctuated<syn::Expr, syn::token::Comma>>,
+
+    call_update: bool,
+
     /// custom field reader code
     reader: Option<TokenStream>,
 
@@ -523,6 +533,12 @@ impl FieldData {
             .transpose()
             .map_err(|e| e.to_compile_error())?;
 
+        let update_ctx = receiver
+            .update_ctx?
+            .map(|s| s.parse_with(Punctuated::parse_terminated))
+            .transpose()
+            .map_err(|e| e.to_compile_error())?;
+
         let data = Self {
             ident: receiver.ident,
             ty: receiver.ty,
@@ -539,6 +555,9 @@ impl FieldData {
             map: receiver.map?,
             ctx,
             update: receiver.update?,
+            update_custom: receiver.update_custom?,
+            update_ctx,
+            call_update: receiver.call_update,
             reader: receiver.reader?,
             writer: receiver.writer?,
             skip: receiver.skip,
@@ -753,6 +772,9 @@ struct DekuReceiver {
     #[darling(default)]
     ctx_default: Option<syn::punctuated::Punctuated<syn::Expr, syn::token::Comma>>,
 
+    #[darling(default)]
+    update_ctx: Option<syn::punctuated::Punctuated<syn::FnArg, syn::token::Comma>>,
+
     /// A magic value that must appear at the start of this struct/enum's data
     #[darling(default)]
     magic: Option<syn::LitByteStr>,
@@ -928,6 +950,19 @@ struct DekuFieldReceiver {
     #[darling(default = "default_res_opt", map = "map_litstr_as_tokenstream")]
     update: Result<Option<TokenStream>, ReplacementError>,
 
+    /// map field when updating struct
+    #[darling(default = "default_res_opt", map = "map_litstr_as_tokenstream")]
+    update_custom: Result<Option<TokenStream>, ReplacementError>,
+
+    /// skip field reading/writing
+    #[darling(default)]
+    call_update: bool,
+
+    // TODO: The type of it should be `Punctuated<Expr, Comma>`
+    //       https://github.com/TedDriggs/darling/pull/98
+    #[darling(default = "default_res_opt", map = "map_option_litstr")]
+    update_ctx: Result<Option<syn::LitStr>, ReplacementError>,
+
     /// custom field reader code
     #[darling(default = "default_res_opt", map = "map_litstr_as_tokenstream")]
     reader: Result<Option<TokenStream>, ReplacementError>,

deku-derive/src/macros/deku_write.rs

@@ -126,6 +126,7 @@ fn emit_struct(input: &DekuData) -> Result<TokenStream, syn::Error> {
     }
 
     let (ctx_types, ctx_arg) = gen_ctx_types_and_arg(input.ctx.as_ref())?;
+    let (update_ctx_types, update_ctx_arg) = gen_ctx_types_and_arg(input.update_ctx.as_ref())?;
 
     let write_body = quote! {
         match *self {
@@ -144,9 +145,8 @@ fn emit_struct(input: &DekuData) -> Result<TokenStream, syn::Error> {
 
     tokens.extend(quote! {
         #[automatically_derived]
-        impl #imp ::#crate_::DekuUpdate for #ident #wher {
-            #[inline]
-            fn update(&mut self) -> core::result::Result<(), ::#crate_::DekuError> {
+        impl #imp ::#crate_::DekuUpdate<#update_ctx_types> for #ident #wher {
+            fn update(&mut self, #update_ctx_arg) -> core::result::Result<(), ::#crate_::DekuError> {
                 #update_use
                 #(#field_updates)*
 
@@ -339,6 +339,7 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
     }
 
     let (ctx_types, ctx_arg) = gen_ctx_types_and_arg(input.ctx.as_ref())?;
+    let (update_ctx_types, update_ctx_arg) = gen_ctx_types_and_arg(input.update_ctx.as_ref())?;
 
     let write_body = quote! {
         #magic_write
@@ -355,9 +356,8 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
 
     tokens.extend(quote! {
         #[automatically_derived]
-        impl #imp ::#crate_::DekuUpdate for #ident #wher {
-            #[inline]
-            fn update(&mut self) -> core::result::Result<(), ::#crate_::DekuError> {
+        impl #imp ::#crate_::DekuUpdate<#update_ctx_types> for #ident #wher {
+            fn update(&mut self, #update_ctx_arg) -> core::result::Result<(), ::#crate_::DekuError> {
                 #update_use
 
                 match self {
@@ -441,17 +441,34 @@ fn emit_field_update(
         return None;
     }
     let field_ident = f.get_ident(i, object_prefix.is_none());
-    let deref = if object_prefix.is_none() {
-        Some(quote! { * })
-    } else {
-        None
+    let deref = match object_prefix.is_none() {
+        true => Some(quote! { * }),
+        false => None,
     };
 
-    f.update.as_ref().map(|field_update| {
-        quote! {
-            #deref #object_prefix #field_ident = (#field_update).try_into()?;
+    if f.call_update {
+        let ctx = match f.update_ctx.as_ref() {
+            Some(ctx) => quote! {(#ctx)},
+            None => quote! {()},
+        };
+        let a = quote! {
+            #deref #object_prefix #field_ident.update(#ctx)?;
+        };
+
+        Some(a)
+    } else {
+        if let Some(custom) = &f.update_custom {
+            Some(quote! {
+                #custom(&mut #object_prefix #field_ident)?;
+            })
+        } else {
+            f.update.as_ref().map(|field_update| {
+                quote! {
+                    #deref #object_prefix #field_ident = (#field_update).try_into()?;
+                }
+            })
         }
-    })
+    }
 }
 
 fn emit_bit_byte_offsets(

examples/update.rs

@@ -0,0 +1,68 @@
+use deku::prelude::*;
+
+#[derive(Debug, DekuRead, DekuWrite, PartialEq)]
+pub struct Test {
+    #[deku(call_update, update_ctx = "self.val.len() as u16, 0")]
+    hdr: Hdr,
+
+    #[deku(count = "hdr.length")]
+    val: Vec<u8>,
+
+    #[deku(call_update)]
+    no_update_ctx: NoUpdateCtx,
+
+    #[deku(update_custom = "Self::custom")]
+    num: u8,
+
+    #[deku(update_custom = "Self::other_custom")]
+    other_num: (u8, u32),
+}
+
+impl Test {
+    fn custom(num: &mut u8) -> Result<(), DekuError> {
+        *num = 1;
+
+        Ok(())
+    }
+
+    fn other_custom(num: &mut (u8, u32)) -> Result<(), DekuError> {
+        *num = (0xf0, 0x0f);
+
+        Ok(())
+    }
+}
+
+#[derive(Debug, DekuRead, DekuWrite, PartialEq)]
+#[deku(update_ctx = "val_len: u16, _na: u8")]
+struct Hdr {
+    #[deku(update = "val_len")]
+    length: u8,
+}
+
+#[derive(Debug, DekuRead, DekuWrite, PartialEq)]
+struct NoUpdateCtx {
+    #[deku(update = "0xff")]
+    val: u8,
+}
+
+fn main() {
+    let mut test = Test {
+        hdr: Hdr { length: 2 },
+        val: vec![1, 2],
+        no_update_ctx: NoUpdateCtx { val: 0 },
+        num: 0,
+        other_num: (0, 0),
+    };
+
+    test.val = vec![1, 2, 3];
+    test.update(()).unwrap();
+
+    let expected = Test {
+        hdr: Hdr { length: 3 },
+        val: test.val.clone(),
+        no_update_ctx: NoUpdateCtx { val: 0xff },
+        num: 1,
+        other_num: (0xf0, 0x0f),
+    };
+    assert_eq!(expected, test);
+}

src/attributes.rs

@@ -908,7 +908,7 @@ assert_eq!(
 value.items.push(0xFF);
 
 // update it, this will update the `count` field
-value.update().unwrap();
+value.update(()).unwrap();
 
 assert_eq!(
     DekuTest { count: 0x03, items: vec![0xAB, 0xCD, 0xFF] },

src/lib.rs

@@ -138,7 +138,7 @@ let data_out = val.to_bytes().unwrap();
 assert_eq!(vec![0x02, 0xBE, 0xEF, 0xAA], data_out);
 
 // Use `update` to update `count`
-val.update().unwrap();
+val.update(()).unwrap();
 
 assert_eq!(DekuTest {
     count: 0x03,
@@ -591,9 +591,9 @@ pub trait DekuContainerWrite: DekuWriter<()> {
 }
 
 /// "Updater" trait: apply mutations to a type
-pub trait DekuUpdate {
+pub trait DekuUpdate<Ctx = ()> {
     /// Apply updates
-    fn update(&mut self) -> Result<(), DekuError>;
+    fn update(&mut self, ctx: Ctx) -> Result<(), DekuError>;
 }
 
 /// "Extended Enum" trait: obtain additional enum information

tests/test_attributes/test_update.rs

@@ -18,7 +18,7 @@ fn test_update() {
     assert_eq!(TestStruct { field_a: 0x01 }, ret_read);
 
     // `field_a` field should now be increased
-    ret_read.update().unwrap();
+    ret_read.update(()).unwrap();
     assert_eq!(0x05, ret_read.field_a);
 
     let ret_write: Vec<u8> = ret_read.try_into().unwrap();
@@ -53,7 +53,7 @@ fn test_update_from_field() {
     ret_read.data.push(0xff);
 
     // `count` field should now be increased
-    ret_read.update().unwrap();
+    ret_read.update(()).unwrap();
     assert_eq!(3, ret_read.count);
 
     // Write
@@ -75,5 +75,5 @@ fn test_update_error() {
 
     let mut val = TestStruct { count: 0x01 };
 
-    val.update().unwrap();
+    val.update(()).unwrap();
 }