Move ConcatSlice to util

libjxl · Sep 19, 2024 · ce8c89f · ce8c89f
1 parent d85166c
commit ce8c89f
Show file tree

Hide file tree

Showing 3 changed files with 116 additions and 75 deletions.
diff --git a/jxl/src/container/box_header.rs b/jxl/src/container/box_header.rs
@@ -5,7 +5,7 @@
 //
 // Originally written for jxl-oxide.
 
-use crate::error::Error;
+use crate::{error::Error, util::ConcatSlice};
 
 /// Box header used in JPEG XL containers.
 #[derive(Debug, Clone)]
@@ -24,7 +24,7 @@ pub enum HeaderParseResult {
 }
 
 impl ContainerBoxHeader {
-    pub(super) fn parse(reader: &super::ConcatSlice<'_, '_>) -> Result<HeaderParseResult, Error> {
+    pub(super) fn parse(reader: &ConcatSlice<'_, '_>) -> Result<HeaderParseResult, Error> {
         let mut buf = [0u8; 16];
         let buf = reader.peek(&mut buf);
         let (tbox, box_size, header_size) = match *buf {

diff --git a/jxl/src/container/mod.rs b/jxl/src/container/mod.rs
@@ -9,7 +9,7 @@ pub mod box_header;
 
 use box_header::*;
 
-use crate::error::Error;
+use crate::{error::Error, util::ConcatSlice};
 
 /// Container format parser.
 #[derive(Default)]
@@ -70,78 +70,6 @@ enum JxlpIndexState {
     JxlpFinished,
 }
 
-struct ConcatSlice<'first, 'second> {
-    slices: (&'first [u8], &'second [u8]),
-    ptr: usize,
-}
-
-impl<'first, 'second> ConcatSlice<'first, 'second> {
-    fn new(slice0: &'first [u8], slice1: &'second [u8]) -> Self {
-        Self {
-            slices: (slice0, slice1),
-            ptr: 0,
-        }
-    }
-
-    fn len(&self) -> usize {
-        self.slices.0.len() + self.slices.1.len()
-    }
-
-    fn remaining_slices(&self) -> (&'first [u8], &'second [u8]) {
-        let (slice0, slice1) = self.slices;
-        let total_len = self.len();
-        let ptr = self.ptr;
-        if ptr >= total_len {
-            (&[], &[])
-        } else if let Some(second_slice_ptr) = ptr.checked_sub(slice0.len()) {
-            (&[], &slice1[second_slice_ptr..])
-        } else {
-            (&slice0[ptr..], slice1)
-        }
-    }
-
-    fn advance(&mut self, bytes: usize) {
-        self.ptr += bytes;
-    }
-
-    fn peek<'out>(&self, out_buf: &'out mut [u8]) -> &'out mut [u8] {
-        let (slice0, slice1) = self.remaining_slices();
-        let total_len = slice0.len() + slice1.len();
-
-        let out_bytes = out_buf.len().min(total_len);
-        let out_buf = &mut out_buf[..out_bytes];
-
-        if out_bytes <= slice0.len() {
-            out_buf.copy_from_slice(&slice0[..out_bytes]);
-        } else {
-            let (out_first, out_second) = out_buf.split_at_mut(slice0.len());
-            out_first.copy_from_slice(slice0);
-            out_second.copy_from_slice(&slice1[..out_second.len()]);
-        }
-
-        out_buf
-    }
-
-    fn fill_vec(&mut self, max_bytes: Option<usize>, v: &mut Vec<u8>) -> Result<usize, Error> {
-        let (slice0, slice1) = self.remaining_slices();
-        let total_len = slice0.len() + slice1.len();
-
-        let out_bytes = max_bytes.unwrap_or(usize::MAX).min(total_len);
-        v.try_reserve(out_bytes)?;
-
-        if out_bytes <= slice0.len() {
-            v.extend_from_slice(&slice0[..out_bytes]);
-        } else {
-            let second_slice_len = out_bytes - slice0.len();
-            v.extend_from_slice(slice0);
-            v.extend_from_slice(&slice1[..second_slice_len]);
-        }
-
-        self.advance(out_bytes);
-        Ok(out_bytes)
-    }
-}
-
 impl ContainerParser {
     const CODESTREAM_SIG: [u8; 2] = [0xff, 0x0a];
     const CONTAINER_SIG: [u8; 12] = [0, 0, 0, 0xc, b'J', b'X', b'L', b' ', 0xd, 0xa, 0x87, 0xa];

diff --git a/jxl/src/util.rs b/jxl/src/util.rs
@@ -3,6 +3,8 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+use crate::error::Error;
+
 pub trait FloorLog2 {
     fn floor_log2(&self) -> Self;
 }
@@ -44,6 +46,78 @@ where
     }
 }
 
+pub struct ConcatSlice<'first, 'second> {
+    slices: (&'first [u8], &'second [u8]),
+    ptr: usize,
+}
+
+impl<'first, 'second> ConcatSlice<'first, 'second> {
+    pub fn new(slice0: &'first [u8], slice1: &'second [u8]) -> Self {
+        Self {
+            slices: (slice0, slice1),
+            ptr: 0,
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        self.slices.0.len() + self.slices.1.len()
+    }
+
+    pub fn remaining_slices(&self) -> (&'first [u8], &'second [u8]) {
+        let (slice0, slice1) = self.slices;
+        let total_len = self.len();
+        let ptr = self.ptr;
+        if ptr >= total_len {
+            (&[], &[])
+        } else if let Some(second_slice_ptr) = ptr.checked_sub(slice0.len()) {
+            (&[], &slice1[second_slice_ptr..])
+        } else {
+            (&slice0[ptr..], slice1)
+        }
+    }
+
+    pub fn advance(&mut self, bytes: usize) {
+        self.ptr += bytes;
+    }
+
+    pub fn peek<'out>(&self, out_buf: &'out mut [u8]) -> &'out mut [u8] {
+        let (slice0, slice1) = self.remaining_slices();
+        let total_len = slice0.len() + slice1.len();
+
+        let out_bytes = out_buf.len().min(total_len);
+        let out_buf = &mut out_buf[..out_bytes];
+
+        if out_bytes <= slice0.len() {
+            out_buf.copy_from_slice(&slice0[..out_bytes]);
+        } else {
+            let (out_first, out_second) = out_buf.split_at_mut(slice0.len());
+            out_first.copy_from_slice(slice0);
+            out_second.copy_from_slice(&slice1[..out_second.len()]);
+        }
+
+        out_buf
+    }
+
+    pub fn fill_vec(&mut self, max_bytes: Option<usize>, v: &mut Vec<u8>) -> Result<usize, Error> {
+        let (slice0, slice1) = self.remaining_slices();
+        let total_len = slice0.len() + slice1.len();
+
+        let out_bytes = max_bytes.unwrap_or(usize::MAX).min(total_len);
+        v.try_reserve(out_bytes)?;
+
+        if out_bytes <= slice0.len() {
+            v.extend_from_slice(&slice0[..out_bytes]);
+        } else {
+            let second_slice_len = out_bytes - slice0.len();
+            v.extend_from_slice(slice0);
+            v.extend_from_slice(&slice1[..second_slice_len]);
+        }
+
+        self.advance(out_bytes);
+        Ok(out_bytes)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -75,4 +149,43 @@ mod test {
         assert_eq!(2, 3usize.ceil_log2());
         assert_eq!(2, 4usize.ceil_log2());
     }
+
+    #[test]
+    fn test_concat_slice_peek_advance() {
+        let mut reader = ConcatSlice::new(&[0, 1, 2, 3], &[4, 5, 6, 7]);
+        let mut buf = [0u8; 8];
+
+        let actual = reader.peek(&mut buf[..1]);
+        assert_eq!(actual, &[0]);
+        reader.advance(actual.len());
+
+        let actual = reader.peek(&mut buf[..2]);
+        assert_eq!(actual, &[1, 2]);
+        reader.advance(actual.len());
+
+        let actual = reader.peek(&mut buf[..3]);
+        assert_eq!(actual, &[3, 4, 5]);
+        reader.advance(actual.len());
+
+        let actual = reader.peek(&mut buf);
+        assert_eq!(actual, &[6, 7]);
+        reader.advance(actual.len());
+
+        let actual = reader.peek(&mut buf);
+        assert!(actual.is_empty());
+    }
+
+    #[test]
+    fn test_concat_slice_fill_vec() {
+        let mut reader = ConcatSlice::new(&[0, 1, 2, 3], &[4, 5, 6, 7]);
+        let mut v = Vec::new();
+
+        let count = reader.fill_vec(Some(3), &mut v).unwrap();
+        assert_eq!(count, 3);
+        assert_eq!(&v, &[0, 1, 2]);
+
+        let count = reader.fill_vec(None, &mut v).unwrap();
+        assert_eq!(count, 5);
+        assert_eq!(&v, &[0, 1, 2, 3, 4, 5, 6, 7]);
+    }
 }