Merge e744eb0 into e92700a

sharksforarms · May 4, 2024 · 2a61000 · 2a61000
2 parents e92700a + e744eb0
commit 2a61000
Show file tree

Hide file tree

Showing 30 changed files with 471 additions and 149 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -19,6 +19,7 @@
      VariantC((u8, u8)),
  }
 ```
+- Add attributes `seek_from_start`, `seek_from_current`, `seek_from_end`, and `seek_rewind` to control the position of the reader before reading a field ([#360](https://github.com/sharksforarms/deku/pull/360))
 
 ### Updated Reader API
 - Changed API of reading to use `io::Read`, bringing massive performance and usability improvements ([#352](https://github.com/sharksforarms/deku/pull/352))

diff --git a/benches/deku.rs b/benches/deku.rs
@@ -1,4 +1,4 @@
-use std::io::{Cursor, Read};
+use std::io::{Cursor, Read, Seek};
 
 use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion};
 use deku::prelude::*;
@@ -38,7 +38,7 @@ fn deku_write<T: DekuContainerWrite>(input: &T) {
     let _v = input.to_bytes().unwrap();
 }
 
-fn deku_read<T: for<'a> DekuContainerRead<'a>>(mut reader: impl Read) {
+fn deku_read<T: for<'a> DekuContainerRead<'a>>(mut reader: impl Read + Seek) {
     let mut reader = Reader::new(&mut reader);
     let _v = T::from_reader_with_ctx(&mut reader, ()).unwrap();
 }

diff --git a/deku-derive/src/lib.rs b/deku-derive/src/lib.rs
@@ -140,6 +140,18 @@ struct DekuData {
 
     /// enum only: byte size of the enum `id`
     bytes: Option<Num>,
+
+    /// struct only: seek from current position
+    seek_rewind: bool,
+
+    /// struct only: seek from current position
+    seek_from_current: Option<Num>,
+
+    /// struct only: seek from end position
+    seek_from_end: Option<Num>,
+
+    /// struct only: seek from start position
+    seek_from_start: Option<Num>,
 }
 
 impl DekuData {
@@ -188,6 +200,10 @@ impl DekuData {
             id_type: receiver.id_type?,
             bits: receiver.bits,
             bytes: receiver.bytes,
+            seek_rewind: receiver.seek_rewind,
+            seek_from_current: receiver.seek_from_current,
+            seek_from_end: receiver.seek_from_end,
+            seek_from_start: receiver.seek_from_start,
         };
 
         DekuData::validate(&data)?;
@@ -269,12 +285,14 @@ impl DekuData {
 
     /// Emit a reader. On error, a compiler error is emitted
     fn emit_reader(&self) -> TokenStream {
-        self.emit_reader_checked().unwrap_or_else(|e| e.to_compile_error())
+        self.emit_reader_checked()
+            .unwrap_or_else(|e| e.to_compile_error())
     }
 
     /// Emit a writer. On error, a compiler error is emitted
     fn emit_writer(&self) -> TokenStream {
-        self.emit_writer_checked().unwrap_or_else(|e| e.to_compile_error())
+        self.emit_writer_checked()
+            .unwrap_or_else(|e| e.to_compile_error())
     }
 
     /// Same as `emit_reader`, but won't auto convert error to compile error
@@ -442,6 +460,18 @@ struct FieldData {
 
     // assert value of field
     assert_eq: Option<TokenStream>,
+
+    /// seek from current position
+    seek_rewind: bool,
+
+    /// seek from current position
+    seek_from_current: Option<TokenStream>,
+
+    /// seek from end position
+    seek_from_end: Option<TokenStream>,
+
+    /// seek from start position
+    seek_from_start: Option<TokenStream>,
 }
 
 impl FieldData {
@@ -479,6 +509,10 @@ impl FieldData {
             cond: receiver.cond?,
             assert: receiver.assert?,
             assert_eq: receiver.assert_eq?,
+            seek_rewind: receiver.seek_rewind,
+            seek_from_current: receiver.seek_from_current?,
+            seek_from_end: receiver.seek_from_end?,
+            seek_from_start: receiver.seek_from_start?,
         };
 
         FieldData::validate(&data)?;
@@ -542,6 +576,19 @@ impl FieldData {
             ));
         }
 
+        // Validate usage of seek_*
+        if (data.seek_from_current.is_some() as u8
+            + data.seek_from_end.is_some() as u8
+            + data.seek_from_start.is_some() as u8
+            + data.seek_rewind as u8)
+            > 1
+        {
+            return Err(cerror(
+                data.bits.span(),
+                "conflicting: only one `seek` attribute can be used at one time",
+            ));
+        }
+
         Ok(())
     }
 
@@ -666,6 +713,22 @@ struct DekuReceiver {
     /// enum only: byte size of the enum `id`
     #[darling(default)]
     bytes: Option<Num>,
+
+    /// struct only: seek from current position
+    #[darling(default)]
+    seek_rewind: bool,
+
+    /// struct only: seek from current position
+    #[darling(default)]
+    seek_from_current: Option<Num>,
+
+    /// struct only: seek from end position
+    #[darling(default)]
+    seek_from_end: Option<Num>,
+
+    /// struct only: seek from start position
+    #[darling(default)]
+    seek_from_start: Option<Num>,
 }
 
 type ReplacementError = TokenStream;
@@ -846,6 +909,22 @@ struct DekuFieldReceiver {
     // assert value of field
     #[darling(default = "default_res_opt", map = "map_litstr_as_tokenstream")]
     assert_eq: Result<Option<TokenStream>, ReplacementError>,
+
+    /// seek from current position
+    #[darling(default)]
+    seek_rewind: bool,
+
+    /// seek from current position
+    #[darling(default = "default_res_opt", map = "map_litstr_as_tokenstream")]
+    seek_from_current: Result<Option<TokenStream>, ReplacementError>,
+
+    /// seek from end position
+    #[darling(default = "default_res_opt", map = "map_litstr_as_tokenstream")]
+    seek_from_end: Result<Option<TokenStream>, ReplacementError>,
+
+    /// seek from start position
+    #[darling(default = "default_res_opt", map = "map_litstr_as_tokenstream")]
+    seek_from_start: Result<Option<TokenStream>, ReplacementError>,
 }
 
 /// Receiver for the variant-level attributes inside a enum

diff --git a/deku-derive/src/macros/deku_read.rs b/deku-derive/src/macros/deku_read.rs
@@ -118,7 +118,7 @@ fn emit_struct(input: &DekuData) -> Result<TokenStream, syn::Error> {
     tokens.extend(quote! {
         impl #imp ::#crate_::DekuReader<#lifetime, #ctx_types> for #ident #wher {
             #[inline]
-            fn from_reader_with_ctx<R: ::#crate_::no_std_io::Read>(__deku_reader: &mut ::#crate_::reader::Reader<R>, #ctx_arg) -> core::result::Result<Self, ::#crate_::DekuError> {
+            fn from_reader_with_ctx<R: ::#crate_::no_std_io::Read + ::#crate_::no_std_io::Seek>(__deku_reader: &mut ::#crate_::reader::Reader<R>, #ctx_arg) -> core::result::Result<Self, ::#crate_::DekuError> {
                 #read_body
             }
         }
@@ -130,7 +130,7 @@ fn emit_struct(input: &DekuData) -> Result<TokenStream, syn::Error> {
         tokens.extend(quote! {
             impl #imp ::#crate_::DekuReader<#lifetime> for #ident #wher {
                 #[inline]
-                fn from_reader_with_ctx<R: ::#crate_::no_std_io::Read>(__deku_reader: &mut ::#crate_::reader::Reader<R>, _: ()) -> core::result::Result<Self, ::#crate_::DekuError> {
+                fn from_reader_with_ctx<R: ::#crate_::no_std_io::Read + ::#crate_::no_std_io::Seek>(__deku_reader: &mut ::#crate_::reader::Reader<R>, _: ()) -> core::result::Result<Self, ::#crate_::DekuError> {
                     #read_body
                 }
             }
@@ -381,7 +381,7 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
         #[allow(non_snake_case)]
         impl #imp ::#crate_::DekuReader<#lifetime, #ctx_types> for #ident #wher {
             #[inline]
-            fn from_reader_with_ctx<R: ::#crate_::no_std_io::Read>(__deku_reader: &mut ::#crate_::reader::Reader<R>, #ctx_arg) -> core::result::Result<Self, ::#crate_::DekuError> {
+            fn from_reader_with_ctx<R: ::#crate_::no_std_io::Read + ::#crate_::no_std_io::Seek>(__deku_reader: &mut ::#crate_::reader::Reader<R>, #ctx_arg) -> core::result::Result<Self, ::#crate_::DekuError> {
                 #read_body
             }
         }
@@ -394,7 +394,7 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
             #[allow(non_snake_case)]
             impl #imp ::#crate_::DekuReader<#lifetime> for #ident #wher {
                 #[inline]
-                fn from_reader_with_ctx<R: ::#crate_::no_std_io::Read>(__deku_reader: &mut ::#crate_::reader::Reader<R>, _: ()) -> core::result::Result<Self, ::#crate_::DekuError> {
+                fn from_reader_with_ctx<R: ::#crate_::no_std_io::Read + ::#crate_::no_std_io::Seek>(__deku_reader: &mut ::#crate_::reader::Reader<R>, _: ()) -> core::result::Result<Self, ::#crate_::DekuError> {
                     #read_body
                 }
             }
@@ -568,6 +568,49 @@ fn emit_field_read(
         &f.assert_eq,
     ];
 
+    let seek = if let Some(num) = &f.seek_from_current {
+        quote! {
+            {
+                use ::#crate_::no_std_io::Seek;
+                use ::#crate_::no_std_io::SeekFrom;
+                if let Err(e) = __deku_reader.seek(SeekFrom::Current(i64::try_from(#num).unwrap())) {
+                    return Err(DekuError::Io(e.kind()));
+                }
+            }
+        }
+    } else if let Some(num) = &f.seek_from_end {
+        quote! {
+            {
+                use ::#crate_::no_std_io::Seek;
+                use ::#crate_::no_std_io::SeekFrom;
+                if let Err(e) = __deku_reader.seek(SeekFrom::End(i64::try_from(#num).unwrap())) {
+                    return Err(DekuError::Io(e.kind()));
+                }
+            }
+        }
+    } else if let Some(num) = &f.seek_from_start {
+        quote! {
+            {
+                use ::#crate_::no_std_io::Seek;
+                use ::#crate_::no_std_io::SeekFrom;
+                if let Err(e) = __deku_reader.seek(SeekFrom::Start(u64::try_from(#num).unwrap())) {
+                    return Err(DekuError::Io(e.kind()));
+                }
+            }
+        }
+    } else if f.seek_rewind {
+        quote! {
+            {
+                use ::#crate_::no_std_io::Seek;
+                if let Err(e) = __deku_reader.rewind() {
+                    return Err(DekuError::Io(e.kind()));
+                }
+            }
+        }
+    } else {
+        quote! {}
+    };
+
     let (bit_offset, byte_offset) = emit_bit_byte_offsets(&field_check_vars);
 
     let field_map = f
@@ -761,6 +804,7 @@ fn emit_field_read(
     };
 
     let field_read = quote! {
+        #seek
         #pad_bits_before
 
         #bit_offset
@@ -795,7 +839,7 @@ pub fn emit_container_read(
         impl #imp ::#crate_::DekuContainerRead<#lifetime> for #ident #wher {
             #[allow(non_snake_case)]
             #[inline]
-            fn from_reader<'a, R: ::#crate_::no_std_io::Read>(__deku_input: (&'a mut R, usize)) -> core::result::Result<(usize, Self), ::#crate_::DekuError> {
+            fn from_reader<'a, R: ::#crate_::no_std_io::Read + ::#crate_::no_std_io::Seek>(__deku_input: (&'a mut R, usize)) -> core::result::Result<(usize, Self), ::#crate_::DekuError> {
                 #from_reader_body
             }
 

diff --git a/examples/custom_reader_and_writer.rs b/examples/custom_reader_and_writer.rs
@@ -5,7 +5,7 @@ use deku::writer::Writer;
 use deku::{prelude::*, DekuWriter};
 use no_std_io::io::Write;
 
-fn bit_flipper_read<R: std::io::Read>(
+fn bit_flipper_read<R: std::io::Read + std::io::Seek>(
     field_a: u8,
     reader: &mut Reader<R>,
     bit_size: BitSize,
@@ -59,8 +59,9 @@ struct DekuTest {
 
 fn main() {
     let test_data = [0x01, 0b1001_0110];
+    let mut cursor = std::io::Cursor::new(test_data);
 
-    let (_read_amt, ret_read) = DekuTest::from_reader((&mut test_data.as_slice(), 0)).unwrap();
+    let (_read_amt, ret_read) = DekuTest::from_reader((&mut cursor, 0)).unwrap();
 
     assert_eq!(
         ret_read,

diff --git a/examples/deku_input.rs b/examples/deku_input.rs
@@ -1,15 +1,15 @@
 //! Example of a close replacement for deku::input
 use deku::prelude::*;
-use std::io::{self, Cursor, Read};
+use std::io::{self, Cursor, Read, Seek, SeekFrom};
 
 /// Every read to this struct will be saved into an internal cache. This is to keep the cache
 /// around for the crc without reading from the buffer twice
-struct ReaderCrc<R: Read> {
+struct ReaderCrc<R: Read + Seek> {
     reader: R,
     pub cache: Vec<u8>,
 }
 
-impl<R: Read> ReaderCrc<R> {
+impl<R: Read + Seek> ReaderCrc<R> {
     pub fn new(reader: R) -> Self {
         Self {
             reader,
@@ -18,14 +18,20 @@ impl<R: Read> ReaderCrc<R> {
     }
 }
 
-impl<R: Read> Read for ReaderCrc<R> {
+impl<R: Read + Seek> Read for ReaderCrc<R> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
         let n = self.reader.read(buf);
         self.cache.extend_from_slice(buf);
         n
     }
 }
 
+impl<R: Read + Seek> Seek for ReaderCrc<R> {
+    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
+        self.reader.seek(pos)
+    }
+}
+
 #[derive(Debug, DekuRead)]
 pub struct DekuStruct {
     pub a: u8,