Mutable scanline iterator

src/png/mod.rs

@@ -29,7 +29,7 @@ const STD_FILTERS: [u8; 2] = [0, 5];
 
 mod scan_lines;
 
-use self::scan_lines::{ScanLine, ScanLines};
+use self::scan_lines::{ScanLine, ScanLines, ScanLinesMut};
 
 #[derive(Debug, Clone)]
 /// Contains all data relevant to a PNG image
@@ -227,6 +227,12 @@ impl PngData {
         ScanLines::new(self)
     }
 
+    /// Return an iterator over the scanlines of the image
+    #[inline]
+    pub fn scan_lines_mut(&mut self) -> ScanLinesMut {
+        ScanLinesMut::new(self)
+    }
+
     /// Reverse all filters applied on the image, returning an unfiltered IDAT bytestream
     pub fn unfilter_image(&self) -> Vec<u8> {
         let mut unfiltered = Vec::with_capacity(self.raw_data.len());
@@ -410,7 +416,6 @@ impl PngData {
     }
 
     fn do_palette_reduction(&mut self, palette_map: &[u8; 256], used: &[bool; 256]) {
-        let mut new_data = Vec::with_capacity(self.raw_data.len());
         let mut byte_map = *palette_map;
 
         // low bit-depths can be pre-computed for every byte value
@@ -429,14 +434,12 @@ impl PngData {
         }
 
         // Reassign data bytes to new indices
-        for line in self.scan_lines() {
-            new_data.push(line.filter);
-            for &byte in line.data {
-                new_data.push(byte_map[byte as usize])
+        for line in self.scan_lines_mut() {
+            for byte in line.data {
+                *byte = byte_map[*byte as usize];
             }
         }
 
-        self.raw_data = new_data;
         self.transparency_pixel = None;
         if let Some(palette) = self.palette.take() {
             let max_index = palette_map.iter().max().cloned().unwrap_or(0) as usize;

src/png/scan_lines.rs

@@ -3,35 +3,102 @@ use super::PngData;
 #[derive(Debug, Clone)]
 /// An iterator over the scan lines of a PNG image
 pub struct ScanLines<'a> {
+    iter: ScanLineRanges,
     /// A reference to the PNG image being iterated upon
-    start: usize,
-    /// Current pass number, and 0-indexed row within the pass
-    pass: Option<(u8, u32)>,
-    bits_per_pixel: u8,
-    width: u32,
-    height: u32,
     raw_data: &'a [u8],
 }
 
 impl<'a> ScanLines<'a> {
     pub fn new(png: &'a PngData) -> Self {
         Self {
-            bits_per_pixel: png.ihdr_data.bit_depth.as_u8() * png.channels_per_pixel(),
-            width: png.ihdr_data.width,
-            height: png.ihdr_data.height,
+            iter: ScanLineRanges::new(png),
             raw_data: &png.raw_data,
-            start: 0,
-            pass: if png.ihdr_data.interlaced == 1 {Some((1, 0))} else {None},
-       }
+        }
     }
 }
 
 impl<'a> Iterator for ScanLines<'a> {
     type Item = ScanLine<'a>;
+    #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        if self.start >= self.raw_data.len() {
-            None
-        } else if let Some(ref mut pass) = self.pass {
+        self.iter.next().map(|(len, pass)| {
+            let (data, rest) = self.raw_data.split_at(len);
+            self.raw_data = rest;
+            let (&filter, data) = data.split_first().unwrap();
+            ScanLine {
+                filter,
+                data,
+                pass,
+            }
+        })
+    }
+}
+
+#[derive(Debug)]
+/// An iterator over the scan lines of a PNG image
+pub struct ScanLinesMut<'a> {
+    iter: ScanLineRanges,
+    /// A reference to the PNG image being iterated upon
+    raw_data: Option<&'a mut [u8]>,
+}
+
+impl<'a> ScanLinesMut<'a> {
+    pub fn new(png: &'a mut PngData) -> Self {
+        Self {
+            iter: ScanLineRanges::new(png),
+            raw_data: Some(&mut png.raw_data),
+        }
+    }
+}
+
+impl<'a> Iterator for ScanLinesMut<'a> {
+    type Item = ScanLineMut<'a>;
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        self.iter.next().map(|(len, pass)| {
+            let tmp = self.raw_data.take().unwrap();
+            let (data, rest) = tmp.split_at_mut(len);
+            self.raw_data = Some(rest);
+            let (&mut filter, data) = data.split_first_mut().unwrap();
+            ScanLineMut {
+                filter,
+                data,
+                pass,
+            }
+        })
+    }
+}
+
+#[derive(Debug, Clone)]
+/// An iterator over the scan line locations of a PNG image
+struct ScanLineRanges {
+    /// Current pass number, and 0-indexed row within the pass
+    pass: Option<(u8, u32)>,
+    bits_per_pixel: u8,
+    width: u32,
+    height: u32,
+    left: usize,
+}
+
+impl ScanLineRanges {
+    pub fn new(png: &PngData) -> Self {
+        Self {
+            bits_per_pixel: png.ihdr_data.bit_depth.as_u8() * png.channels_per_pixel(),
+            width: png.ihdr_data.width,
+            height: png.ihdr_data.height,
+            left: png.raw_data.len(),
+            pass: if png.ihdr_data.interlaced == 1 {Some((1, 0))} else {None},
+        }
+    }
+}
+
+impl Iterator for ScanLineRanges {
+    type Item = (usize, Option<u8>);
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.left == 0 {
+            return None;
+        }
+        let (pixels_per_line, current_pass) = if let Some(ref mut pass) = self.pass {
             // Scanlines for interlaced PNG files
             // Handle edge cases for images smaller than 5 pixels in either direction
             if self.width < 5 && pass.0 == 2 {
@@ -71,7 +138,6 @@ impl<'a> Iterator for ScanLines<'a> {
                 _ => (),
             };
             let current_pass = Some(pass.0);
-            let bytes_per_line = ((pixels_per_line * self.bits_per_pixel as u32 + 7) / 8) as usize;
             if pass.1 + y_steps >= self.height {
                 pass.0 += 1;
                 pass.1 = match pass.0 {
@@ -83,27 +149,16 @@ impl<'a> Iterator for ScanLines<'a> {
             } else {
                 pass.1 += y_steps;
             }
-            let start = self.start;
-            let len = bytes_per_line + 1;
-            self.start += len;
-            Some(ScanLine {
-                filter: self.raw_data[start],
-                data: &self.raw_data[(start + 1)..(start + len)],
-                pass: current_pass,
-            })
+            (pixels_per_line, current_pass)
         } else {
             // Standard, non-interlaced PNG scanlines
-            let bits_per_line = self.width * self.bits_per_pixel as u32;
-            let bytes_per_line = ((bits_per_line + 7) / 8) as usize;
-            let start = self.start;
-            let len = bytes_per_line + 1;
-            self.start += len;
-            Some(ScanLine {
-                filter: self.raw_data[start],
-                data: &self.raw_data[(start + 1)..(start + len)],
-                pass: None,
-            })
-        }
+            (self.width, None)
+        };
+        let bits_per_line = pixels_per_line * self.bits_per_pixel as u32;
+        let bytes_per_line = ((bits_per_line + 7) / 8) as usize;
+        let len = bytes_per_line + 1;
+        self.left -= len;
+        Some((len, current_pass))
     }
 }
 
@@ -117,3 +172,15 @@ pub struct ScanLine<'a> {
     /// The current pass if the image is interlaced
     pub pass: Option<u8>,
 }
+
+
+#[derive(Debug)]
+/// A scan line in a PNG image
+pub struct ScanLineMut<'a> {
+    /// The filter type used to encode the current scan line (0-4)
+    pub filter: u8,
+    /// The byte data for the current scan line, encoded with the filter specified in the `filter` field
+    pub data: &'a mut [u8],
+    /// The current pass if the image is interlaced
+    pub pass: Option<u8>,
+}