Extract a separate unfiltering_buffer module.

anforowicz · anforowicz · commit 70a217d74cf1 · 2024-11-01T17:48:16.000Z
diff --git a/src/decoder/mod.rs b/src/decoder/mod.rs
@@ -2,11 +2,13 @@ mod interlace_info;
 mod read_decoder;
 pub(crate) mod stream;
 pub(crate) mod transform;
+mod unfiltering_buffer;
 mod zlib;
 
 use self::read_decoder::{ImageDataCompletionStatus, ReadDecoder};
 use self::stream::{DecodeOptions, DecodingError, FormatErrorInner, CHUNK_BUFFER_SIZE};
 use self::transform::{create_transform_fn, TransformFn};
+use self::unfiltering_buffer::UnfilteringBuffer;
 
 use std::io::Read;
 use std::mem;
@@ -15,7 +17,6 @@ use crate::adam7::{self, Adam7Info};
 use crate::common::{
     BitDepth, BytesPerPixel, ColorType, Info, ParameterErrorKind, Transformations,
 };
-use crate::filter::{unfilter, FilterType};
 use crate::FrameControl;
 
 pub use interlace_info::InterlaceInfo;
@@ -192,9 +193,7 @@ impl<R: Read> Decoder<R> {
             bpp: BytesPerPixel::One,
             subframe: SubframeInfo::not_yet_init(),
             remaining_frames: 0, // Temporary value - fixed below after reading `acTL` and `fcTL`.
-            data_stream: Vec::new(),
-            prev_start: 0,
-            current_start: 0,
+            unfiltering_buffer: UnfilteringBuffer::new(),
             transform: self.transform,
             transform_fn: None,
             scratch_buffer: Vec::new(),
@@ -288,12 +287,8 @@ pub struct Reader<R: Read> {
     subframe: SubframeInfo,
     /// How many frames remain to be decoded.  Decremented after each `IDAT` or `fdAT` sequence.
     remaining_frames: usize,
-    /// Vec containing the uncompressed image data currently being processed.
-    data_stream: Vec<u8>,
-    /// Index in `data_stream` where the previous row starts.
-    prev_start: usize,
-    /// Index in `data_stream` where the current row starts.
-    current_start: usize,
+    /// Buffer with not-yet-`unfilter`-ed image rows
+    unfiltering_buffer: UnfilteringBuffer,
     /// Output transformations
     transform: Transformations,
     /// Function that can transform decompressed, unfiltered rows into final output.
@@ -362,16 +357,12 @@ impl<R: Read> Reader<R> {
 
         self.subframe = SubframeInfo::new(self.info());
         self.bpp = self.info().bpp_in_prediction();
-        self.data_stream.clear();
-        self.current_start = 0;
-        self.prev_start = 0;
+        self.unfiltering_buffer = UnfilteringBuffer::new();
 
         // Allocate output buffer.
         let buflen = self.output_line_size(self.subframe.width);
         self.decoder.reserve_bytes(buflen)?;
 
-        self.prev_start = self.current_start;
-
         Ok(())
     }
 
@@ -499,7 +490,7 @@ impl<R: Read> Reader<R> {
             Some(interlace) => *interlace,
         };
         if interlace.line_number() == 0 {
-            self.prev_start = self.current_start;
+            self.unfiltering_buffer.reset_prev_row();
         }
         let rowlen = match interlace {
             InterlaceInfo::Null(_) => self.subframe.rowlen,
@@ -537,9 +528,7 @@ impl<R: Read> Reader<R> {
         }
 
         self.remaining_frames = 0;
-        self.data_stream.clear();
-        self.current_start = 0;
-        self.prev_start = 0;
+        self.unfiltering_buffer = UnfilteringBuffer::new();
         self.decoder.read_until_end_of_input()?;
 
         self.finished = true;
@@ -553,8 +542,8 @@ impl<R: Read> Reader<R> {
         output_buffer: &mut [u8],
     ) -> Result<(), DecodingError> {
         self.next_raw_interlaced_row(rowlen)?;
-        assert_eq!(self.current_start - self.prev_start, rowlen - 1);
-        let row = &self.data_stream[self.prev_start..self.current_start];
+        let row = self.unfiltering_buffer.prev_row();
+        assert_eq!(row.len(), rowlen - 1);
 
         // Apply transformations and write resulting data to buffer.
         let transform_fn = {
@@ -619,51 +608,26 @@ impl<R: Read> Reader<R> {
         color.raw_row_length_from_width(depth, width) - 1
     }
 
-    /// Write the next raw interlaced row into `self.prev`.
-    ///
-    /// The scanline is filtered against the previous scanline according to the specification.
+    /// Unfilter the next raw interlaced row into `self.unfiltering_buffer`.
     fn next_raw_interlaced_row(&mut self, rowlen: usize) -> Result<(), DecodingError> {
         // Read image data until we have at least one full row (but possibly more than one).
-        while self.data_stream.len() - self.current_start < rowlen {
+        while self.unfiltering_buffer.curr_row_len() < rowlen {
             if self.subframe.consumed_and_flushed {
                 return Err(DecodingError::Format(
                     FormatErrorInner::NoMoreImageData.into(),
                 ));
             }
 
-            // Clear the current buffer before appending more data.
-            if self.prev_start > 0 {
-                self.data_stream.copy_within(self.prev_start.., 0);
-                self.data_stream
-                    .truncate(self.data_stream.len() - self.prev_start);
-                self.current_start -= self.prev_start;
-                self.prev_start = 0;
-            }
-
-            match self.decoder.decode_image_data(&mut self.data_stream)? {
+            match self
+                .decoder
+                .decode_image_data(self.unfiltering_buffer.as_mut_vec())?
+            {
                 ImageDataCompletionStatus::ExpectingMoreData => (),
                 ImageDataCompletionStatus::Done => self.mark_subframe_as_consumed_and_flushed(),
             }
         }
 
-        // Get a reference to the current row and point scan_start to the next one.
-        let (prev, row) = self.data_stream.split_at_mut(self.current_start);
-
-        // Unfilter the row.
-        let filter = FilterType::from_u8(row[0]).ok_or(DecodingError::Format(
-            FormatErrorInner::UnknownFilterMethod(row[0]).into(),
-        ))?;
-        unfilter(
-            filter,
-            self.bpp,
-            &prev[self.prev_start..],
-            &mut row[1..rowlen],
-        );
-
-        self.prev_start = self.current_start + 1;
-        self.current_start += rowlen;
-
-        Ok(())
+        self.unfiltering_buffer.unfilter_curr_row(rowlen, self.bpp)
     }
 }
 
diff --git a/src/decoder/stream.rs b/src/decoder/stream.rs
@@ -2381,17 +2381,16 @@ mod tests {
     /// Creates a ready-to-test [`Reader`] which decodes an animated PNG that contains:
     /// IHDR, acTL, fcTL, IDAT, fcTL, fdAT, IEND.  (i.e. IDAT is part of the animation)
     fn create_reader_of_ihdr_actl_fctl_idat_fctl_fdat() -> Reader<VecDeque<u8>> {
-        let width = 16;
-        let frame_data = generate_rgba8_with_width_and_height(width, width);
+        let idat_width = 16;
         let mut fctl = crate::FrameControl {
-            width,
-            height: width,
+            width: idat_width,
+            height: idat_width, // same height and width
             ..Default::default()
         };
 
         let mut png = VecDeque::new();
         write_png_sig(&mut png);
-        write_rgba8_ihdr_with_width(&mut png, width);
+        write_rgba8_ihdr_with_width(&mut png, idat_width);
         write_actl(
             &mut png,
             &crate::AnimationControl {
@@ -2401,10 +2400,21 @@ mod tests {
         );
         fctl.sequence_number = 0;
         write_fctl(&mut png, &fctl);
-        write_chunk(&mut png, b"IDAT", &frame_data);
+        // Using `fctl.height + 1` means that the `IDAT` will have "left-over" data after
+        // processing.  This helps to verify that `Reader.read_until_image_data` discards the
+        // left-over data when resetting `UnfilteredRowsBuffer`.
+        let idat_data = generate_rgba8_with_width_and_height(fctl.width, fctl.height + 1);
+        write_chunk(&mut png, b"IDAT", &idat_data);
+
+        let fdat_width = 10;
         fctl.sequence_number = 1;
+        // Using different width in `IDAT` and `fDAT` frames helps to catch problems that
+        // may arise when `Reader.read_until_image_data` doesn't properly reset
+        // `UnfilteredRowsBuffer`.
+        fctl.width = fdat_width;
         write_fctl(&mut png, &fctl);
-        write_fdat(&mut png, 2, &frame_data);
+        let fdat_data = generate_rgba8_with_width_and_height(fctl.width, fctl.height);
+        write_fdat(&mut png, 2, &fdat_data);
         write_iend(&mut png);
 
         Decoder::new(png).read_info().unwrap()
diff --git a/src/decoder/unfiltering_buffer.rs b/src/decoder/unfiltering_buffer.rs
@@ -0,0 +1,112 @@
+use super::stream::{DecodingError, FormatErrorInner};
+use crate::common::BytesPerPixel;
+use crate::filter::{unfilter, FilterType};
+
+// Buffer for temporarily holding decompressed, not-yet-`unfilter`-ed rows.
+pub(crate) struct UnfilteringBuffer {
+    /// Vec containing the uncompressed image data currently being processed.
+    data_stream: Vec<u8>,
+    /// Index in `data_stream` where the previous row starts.
+    /// This excludes the filter type byte - it points at the first byte of actual pixel data.
+    /// The pixel data is already-`unfilter`-ed.
+    /// If `prev_start == current_start` then it means that there is no previous row.
+    prev_start: usize,
+    /// Index in `data_stream` where the current row starts.
+    /// This points at the filter type byte of the current row (i.e. the actual pixel data starts at `current_start + 1`)
+    /// The pixel data is not-yet-`unfilter`-ed.
+    current_start: usize,
+}
+
+impl UnfilteringBuffer {
+    /// Asserts in debug builds that all the invariants hold.  No-op in release
+    /// builds.  Intended to be called after creating or mutating `self` to
+    /// ensure that the final state preserves the invariants.
+    fn debug_assert_invariants(&self) {
+        debug_assert!(self.prev_start <= self.current_start);
+        debug_assert!(self.prev_start <= self.data_stream.len());
+        debug_assert!(self.current_start <= self.data_stream.len());
+    }
+
+    pub fn new() -> Self {
+        let result = Self {
+            data_stream: Vec::new(),
+            prev_start: 0,
+            current_start: 0,
+        };
+        result.debug_assert_invariants();
+        result
+    }
+
+    /// Called to indicate that there is no previous row (e.g. when the current
+    /// row is the first scanline of a given Adam7 pass).
+    pub fn reset_prev_row(&mut self) {
+        self.prev_start = self.current_start;
+        self.debug_assert_invariants();
+    }
+
+    /// Returns the previous (already `unfilter`-ed) row.
+    pub fn prev_row(&self) -> &[u8] {
+        // No point calling this if there is no previous row.
+        debug_assert!(self.prev_start < self.current_start);
+
+        &self.data_stream[self.prev_start..self.current_start]
+    }
+
+    /// Returns how many bytes of the current row are present in the buffer.
+    pub fn curr_row_len(&self) -> usize {
+        self.data_stream.len() - self.current_start
+    }
+
+    /// Returns a `&mut Vec<u8>` suitable for passing to
+    /// `ReadDecoder.decode_image_data` or `StreamingDecoder.update`.
+    ///
+    /// Invariants of `self` depend on the assumption that the caller will only
+    /// append new bytes to the returned vector (which is indeed the behavior of
+    /// `ReadDecoder` and `StreamingDecoder`).  TODO: Consider protecting the
+    /// invariants by returning an append-only view of the vector
+    /// (`FnMut(&[u8])`??? or maybe `std::io::Write`???).
+    pub fn as_mut_vec(&mut self) -> &mut Vec<u8> {
+        // Opportunistically compact the current buffer by discarding bytes
+        // before `prev_start`.
+        if self.prev_start > 0 {
+            self.data_stream.copy_within(self.prev_start.., 0);
+            self.data_stream
+                .truncate(self.data_stream.len() - self.prev_start);
+            self.current_start -= self.prev_start;
+            self.prev_start = 0;
+            self.debug_assert_invariants();
+        }
+
+        &mut self.data_stream
+    }
+
+    /// Runs `unfilter` on the current row, and then shifts rows so that the current row becomes the previous row.
+    ///
+    /// Will panic if `self.curr_row_len() < rowlen`.
+    pub fn unfilter_curr_row(
+        &mut self,
+        rowlen: usize,
+        bpp: BytesPerPixel,
+    ) -> Result<(), DecodingError> {
+        debug_assert!(rowlen >= 2); // 1 byte for `FilterType` and at least 1 byte of pixel data.
+
+        let (prev, row) = self.data_stream.split_at_mut(self.current_start);
+        let prev: &[u8] = prev; // `prev` is immutable
+        let prev = &prev[self.prev_start..];
+        debug_assert!(prev.is_empty() || prev.len() == (rowlen - 1));
+
+        // Get the filter type.
+        let filter = FilterType::from_u8(row[0]).ok_or(DecodingError::Format(
+            FormatErrorInner::UnknownFilterMethod(row[0]).into(),
+        ))?;
+        let row = &mut row[1..rowlen];
+
+        unfilter(filter, bpp, prev, row);
+
+        self.prev_start = self.current_start + 1;
+        self.current_start += rowlen;
+        self.debug_assert_invariants();
+
+        Ok(())
+    }
+}
