Add a test of streaming bitstream parser

jxl/Cargo.toml

@@ -16,6 +16,9 @@ half = "1.7.1"
 tracing = "0.1.40"
 jxl_headers_derive = { path = "../jxl_headers_derive" }
 
+[dev-dependencies]
+arbtest = "0.3.1"
+
 [features]
 debug_tools = []
 

jxl/src/container/mod.rs

@@ -113,3 +113,74 @@ impl ContainerParser {
         Ok(codestream)
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[rustfmt::skip]
+    const HEADER: &[u8] = &[
+        0x00, 0x00, 0x00, 0x0c, b'J', b'X', b'L', b' ', 0x0d, 0x0a, 0x87, 0x0a, 0x00, 0x00, 0x00, 0x14,
+        b'f', b't', b'y', b'p', b'j', b'x', b'l', b' ', 0x00, 0x00, 0x00, 0x00, b'j', b'x', b'l', b' ',
+    ];
+
+    #[test]
+    fn parse_partial() {
+        arbtest::arbtest(|u| {
+            // Prepare arbitrary container format data with two jxlp boxes.
+            let total_len = u.arbitrary_len::<u8>()?;
+            let mut codestream0 = vec![0u8; total_len / 2];
+            u.fill_buffer(&mut codestream0)?;
+            let mut codestream1 = vec![0u8; total_len - codestream0.len()];
+            u.fill_buffer(&mut codestream1)?;
+
+            let mut container = HEADER.to_vec();
+            container.extend_from_slice(&(12 + codestream0.len() as u32).to_be_bytes());
+            container.extend_from_slice(b"jxlp\x00\x00\x00\x00");
+            container.extend_from_slice(&codestream0);
+
+            container.extend_from_slice(&(12 + codestream1.len() as u32).to_be_bytes());
+            container.extend_from_slice(b"jxlp\x80\x00\x00\x01");
+            container.extend_from_slice(&codestream1);
+
+            let mut expected = codestream0;
+            expected.extend(codestream1);
+
+            // Create a list of arbitrary splits.
+            let mut tests = Vec::new();
+            u.arbitrary_loop(Some(1), Some(10), |u| {
+                let split_at_idx = u.choose_index(container.len())?;
+                tests.push(container.split_at(split_at_idx));
+                Ok(std::ops::ControlFlow::Continue(()))
+            })?;
+
+            // Test if split index doesn't affect final codestream.
+            for (first, second) in tests {
+                let mut codestream = Vec::new();
+                let mut parser = ContainerParser::new();
+
+                for event in parser.process_bytes(first) {
+                    let event = event.unwrap();
+                    if let ParseEvent::Codestream(data) = event {
+                        codestream.extend_from_slice(data);
+                    }
+                }
+
+                let consumed = parser.previous_consumed_bytes();
+                let mut second_chunk = first[consumed..].to_vec();
+                second_chunk.extend_from_slice(second);
+
+                for event in parser.process_bytes(&second_chunk) {
+                    let event = event.unwrap();
+                    if let ParseEvent::Codestream(data) = event {
+                        codestream.extend_from_slice(data);
+                    }
+                }
+
+                assert_eq!(codestream, expected);
+            }
+
+            Ok(())
+        });
+    }
+}