Small write impl

jam1garner · Mar 6, 2024 · 7cfe67d · 7cfe67d
1 parent ed02e5f
commit 7cfe67d
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 1 deletion.
diff --git a/binrw/tests/derive/write/padding.rs b/binrw/tests/derive/write/padding.rs
@@ -73,3 +73,37 @@ fn padding_one_way() {
 
     assert_eq!(x.into_inner(), data);
 }
+
+#[test]
+fn padding_struct() {
+    #[derive(BinRead, BinWrite)]
+    #[brw(pad_size_to = 20)]
+    struct Test {
+        #[brw(pad_before = 0x2_u32, align_after = 0x8)]
+        x: u8,
+
+        #[brw(align_before = 0x4_u32, pad_after = 0x3_u32)]
+        y: u8,
+
+        #[brw(pad_size_to = 0x6_u32)]
+        z: u32,
+    }
+
+    let data = &[
+        /* pad_before: */ 0, 0, /* x */ 1, /* align: */ 0, 0, 0, 0, 0,
+        /* align_before: (none)*/ /* y */ 2, /* pad_after: */ 0, 0, 0, /* z */ 0xef,
+        0xcd, 0xab, 0, /* pad_size_to */ 0, 0, /* struct padding */ 0, 0,
+    ];
+
+    let mut x = Cursor::new(Vec::new());
+
+    Test {
+        x: 1,
+        y: 2,
+        z: 0xabcdef,
+    }
+    .write_options(&mut x, Endian::Little, ())
+    .unwrap();
+
+    assert_eq!(x.into_inner(), data);
+}
diff --git a/binrw_derive/src/binrw/codegen/sanitization.rs b/binrw_derive/src/binrw/codegen/sanitization.rs
@@ -74,6 +74,7 @@ ident_str! {
     pub(crate) READ_FUNCTION = "__binrw_generated_read_function";
     pub(crate) WRITE_FUNCTION = "__binrw_generated_write_function";
     pub(crate) BEFORE_POS = "__binrw_generated_before_pos";
+    pub(crate) STRUCT_POS = "__binrw_generated_struct_pos";
     pub(crate) DBG_EPRINTLN = from_crate!(__private::eprintln);
 }
 

diff --git a/binrw_derive/src/binrw/codegen/write_options/struct.rs b/binrw_derive/src/binrw/codegen/write_options/struct.rs
@@ -1,6 +1,6 @@
 use super::{prelude::PreludeGenerator, struct_field::write_field};
 use crate::binrw::{
-    codegen::sanitization::{THIS, WRITER},
+    codegen::sanitization::{SEEK_TRAIT, STRUCT_POS, THIS, WRITER, WRITE_ZEROES},
     parser::{Input, Struct},
 };
 use proc_macro2::TokenStream;
@@ -10,6 +10,7 @@ use syn::Ident;
 pub(super) fn generate_struct(input: &Input, name: Option<&Ident>, st: &Struct) -> TokenStream {
     StructGenerator::new(input, st, name, &input.stream_ident_or(WRITER))
         .write_fields()
+        .wrap_pad_size()
         .prefix_prelude()
         .prefix_borrow_fields()
         .prefix_imports()
@@ -59,6 +60,28 @@ impl<'input> StructGenerator<'input> {
         self
     }
 
+    fn wrap_pad_size(mut self) -> Self {
+        if let Some(size) = &self.st.pad_size_to {
+            let writer_var = self.writer_var;
+            let out = self.out;
+            self.out = quote! {
+                let #STRUCT_POS = #SEEK_TRAIT::stream_position(#writer_var)?;
+                #out
+                {
+                    let pad_to_size = (#size) as u64;
+                    let after_pos = #SEEK_TRAIT::stream_position(#writer_var)?;
+                    if let Some(size) = after_pos.checked_sub(#STRUCT_POS) {
+                        if let Some(padding) = pad_to_size.checked_sub(size) {
+                            #WRITE_ZEROES(#writer_var, padding)?;
+                        }
+                    }
+                }
+            };
+        }
+
+        self
+    }
+
     pub(super) fn write_fields(mut self) -> Self {
         let write_fields = self
             .st