From e83f76a4f65a15a6f4d484be3963fe9b714a816e Mon Sep 17 00:00:00 2001
From: Wonwoo Choi <chwo9843@gmail.com>
Date: Sun, 22 Sep 2024 04:42:53 +0900
Subject: [PATCH] Add a test of streaming bitstream parser

---
 Cargo.lock               | 16 +++++++++
 jxl/Cargo.toml           |  3 ++
 jxl/src/container/mod.rs | 71 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 90 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index fb79721..e0bec4d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2,6 +2,21 @@
 # It is not intended for manual editing.
 version = 3
 
+[[package]]
+name = "arbitrary"
+version = "1.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7d5a26814d8dcb93b0e5a0ff3c6d80a8843bafb21b39e8e18a6f05471870e110"
+
+[[package]]
+name = "arbtest"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23909d5fb517fac2a8a4c887e847dbe41dd22ec46914586f5727980d0a193fdc"
+dependencies = [
+ "arbitrary",
+]
+
 [[package]]
 name = "array-init"
 version = "2.1.0"
@@ -30,6 +45,7 @@ checksum = "1b43ede17f21864e81be2fa654110bf1e793774238d86ef8555c37e6519c0403"
 name = "jxl"
 version = "0.1.0"
 dependencies = [
+ "arbtest",
  "array-init",
  "byteorder",
  "half",
diff --git a/jxl/Cargo.toml b/jxl/Cargo.toml
index ca26ffa..fa65d41 100644
--- a/jxl/Cargo.toml
+++ b/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 = []
 
diff --git a/jxl/src/container/mod.rs b/jxl/src/container/mod.rs
index d6c6489..af96f17 100644
--- a/jxl/src/container/mod.rs
+++ b/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(())
+        });
+    }
+}