Infer stream type when type checking parse_with/write_with args

Fixes jam1garner#240.
csnover · Nov 28, 2023 · d6f97e1 · d6f97e1
1 parent 24f5f4f
commit d6f97e1
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 2 deletions.
diff --git a/binrw/tests/derive/struct.rs b/binrw/tests/derive/struct.rs
@@ -549,6 +549,31 @@ fn map_stream_field() {
     );
 }
 
+#[test]
+fn map_stream_parse_with_args() {
+    use binrw::{helpers::until_eof, io::TakeSeekExt};
+    #[derive(BinRead, Debug, PartialEq)]
+    #[br(import(extra: u8))]
+    struct Inner(#[br(map = |v: u8| v + extra)] u8);
+
+    #[derive(BinRead, Debug, PartialEq)]
+    struct Test {
+        a: u8,
+        #[br(map_stream = |s| s.take_seek(4), parse_with = until_eof, args(a))]
+        b: Vec<Inner>,
+        c: u8,
+    }
+
+    assert_eq!(
+        Test::read_le(&mut Cursor::new(b"\x0a\x00\x01\x02\x03\x04")).unwrap(),
+        Test {
+            a: 10,
+            b: vec![Inner(10), Inner(11), Inner(12), Inner(13)],
+            c: 4
+        }
+    );
+}
+
 #[test]
 fn named_args_trailing_commas() {
     #[rustfmt::skip]

diff --git a/binrw/tests/derive/write/map_stream.rs b/binrw/tests/derive/write/map_stream.rs
@@ -74,3 +74,30 @@ fn map_stream_field() {
         &[0, 1, 3, 0, b'm', b'a', b'g', b'i', b'c', 132, 129, 135, 128],
     );
 }
+
+#[test]
+fn map_stream_write_with_args() {
+    #[binrw::writer(writer)]
+    fn adder(value: &Vec<u8>, amt: u8) -> binrw::BinResult<()> {
+        for v in value {
+            writer.write_all(&[*v + amt])?;
+        }
+        Ok(())
+    }
+
+    #[derive(BinWrite, Debug, PartialEq)]
+    #[bw(big)]
+    struct Test {
+        #[bw(map_stream = |inner| BadCrypt { inner, key: 0x80 }, write_with = adder, args(10))]
+        a: Vec<u8>,
+    }
+
+    let mut out = Cursor::new(vec![]);
+    Test {
+        a: vec![0, 1, 2, 3],
+    }
+    .write(&mut out)
+    .unwrap();
+
+    assert_eq!(out.into_inner(), &[138, 129, 141, 128],);
+}
diff --git a/binrw_derive/src/binrw/codegen/read_options/struct.rs b/binrw_derive/src/binrw/codegen/read_options/struct.rs
@@ -387,7 +387,7 @@ impl<'field> FieldGenerator<'field> {
 
             if let FieldMode::Function(_) = &self.field.field_mode {
                 quote_spanned! {ty.span()=>
-                    let #args_var = #ARGS_TYPE_HINT::<R, #ty, _, _>(&#READ_FUNCTION, #args);
+                    let #args_var = #ARGS_TYPE_HINT::<_, #ty, _, _>(&#READ_FUNCTION, #args);
                 }
             } else {
                 match &self.field.map {

diff --git a/binrw_derive/src/binrw/codegen/write_options/struct_field.rs b/binrw_derive/src/binrw/codegen/write_options/struct_field.rs
@@ -295,7 +295,7 @@ impl<'a> StructFieldGenerator<'a> {
             FieldMode::Function(_) => {
                 let ty = &self.field.ty;
                 quote! {
-                    let #args = #WRITE_ARGS_TYPE_HINT::<#ty, W, _, _>(
+                    let #args = #WRITE_ARGS_TYPE_HINT::<#ty, _, _, _>(
                         #WRITE_FUNCTION, #args_val
                     );
                     #out