Rewrite bitstream reader, enable partial decoding

crates/jxl-bitstream/src/container.rs

@@ -1,5 +1,3 @@
-use std::io::prelude::*;
-
 use crate::Error;
 
 #[derive(Debug, Clone)]
@@ -9,38 +7,53 @@ pub struct ContainerBoxHeader {
     is_last: bool,
 }
 
-impl ContainerBoxHeader {
-    pub fn parse<R: Read>(mut reader: R) -> std::io::Result<Self> {
-        let mut sbox = [0u8; 4];
-        reader.read_exact(&mut sbox)?;
-        let sbox = u32::from_be_bytes(sbox);
+pub enum HeaderParseResult {
+    Done {
+        header: ContainerBoxHeader,
+        size: usize,
+    },
+    NeedMoreData,
+}
 
-        let mut tbox = [0u8; 4];
-        reader.read_exact(&mut tbox)?;
-        let tbox = ContainerBoxType(tbox);
+impl ContainerBoxHeader {
+    pub fn parse(buf: &[u8]) -> std::io::Result<HeaderParseResult> {
+        if buf.len() < 8 {
+            return Ok(HeaderParseResult::NeedMoreData);
+        }
 
-        let size = if sbox == 1 {
-            let mut xlbox = [0u8; 8];
-            reader.read_exact(&mut xlbox)?;
-            let xlbox = u64::from_be_bytes(xlbox).checked_sub(16).ok_or(
-                std::io::Error::new(std::io::ErrorKind::InvalidData, Error::InvalidBoxSize)
-            )?;
-            Some(xlbox)
-        } else if sbox == 0 {
-            None
-        } else {
-            let sbox = sbox.checked_sub(8).ok_or(
-                std::io::Error::new(std::io::ErrorKind::InvalidData, Error::InvalidBoxSize)
-            )?;
-            Some(sbox as u64)
+        let (tbox, size, header_size) = match *buf {
+            [0, 0, 0, 1, t0, t1, t2, t3, s0, s1, s2, s3, s4, s5, s6, s7, ..] => {
+                let xlbox = u64::from_be_bytes([s0, s1, s2, s3, s4, s5, s6, s7]);
+                let tbox = ContainerBoxType([t0, t1, t2, t3]);
+                let xlbox = xlbox.checked_sub(16).ok_or(
+                    std::io::Error::new(std::io::ErrorKind::InvalidData, Error::InvalidBoxSize)
+                )?;
+                (tbox, Some(xlbox), 16)
+            },
+            [s0, s1, s2, s3, t0, t1, t2, t3, ..] => {
+                let sbox = u32::from_be_bytes([s0, s1, s2, s3]);
+                let tbox = ContainerBoxType([t0, t1, t2, t3]);
+                let sbox = if sbox == 0 {
+                    None
+                } else if let Some(sbox) = sbox.checked_sub(8) {
+                    Some(sbox as u64)
+                } else {
+                    return Err(
+                        std::io::Error::new(std::io::ErrorKind::InvalidData, Error::InvalidBoxSize)
+                    );
+                };
+                (tbox, sbox, 8)
+            },
+            _ => return Ok(HeaderParseResult::NeedMoreData),
         };
         let is_last = size.is_none();
 
-        Ok(Self {
+        let header = Self {
             ty: tbox,
             size,
             is_last,
-        })
+        };
+        Ok(HeaderParseResult::Done { header, size: header_size })
     }
 }
 

crates/jxl-bitstream/src/error.rs

@@ -76,4 +76,13 @@ impl From<std::io::Error> for Error {
     }
 }
 
+impl Error {
+    pub fn unexpected_eof(&self) -> bool {
+        if let Error::Io(e) = self {
+            return e.kind() == std::io::ErrorKind::UnexpectedEof;
+        }
+        false
+    }
+}
+
 pub type Result<T> = std::result::Result<T, Error>;