80× faster palette reduction (#150)

* Faster palette reduction

* Simplified iterator

* Mutable scanline iterator

src/png/mod.rs

@@ -1,7 +1,7 @@
+use std::collections::hash_map::Entry::*;
 use rgb::RGBA8;
 use rgb::ComponentSlice;
 use atomicmin::AtomicMin;
-use bit_vec::BitVec;
 use byteorder::{BigEndian, WriteBytesExt};
 use colors::{AlphaOptim, BitDepth, ColorType};
 use crc::crc32;
@@ -10,7 +10,7 @@ use error::PngError;
 use filters::*;
 use headers::*;
 use interlace::{deinterlace_image, interlace_image};
-use itertools::{flatten, Itertools};
+use itertools::flatten;
 #[cfg(feature = "parallel")]
 use rayon::prelude::*;
 use reduction::bit_depth::*;
@@ -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
@@ -224,12 +224,13 @@ impl PngData {
     /// Return an iterator over the scanlines of the image
     #[inline]
     pub fn scan_lines(&self) -> ScanLines {
-        ScanLines {
+        ScanLines::new(self)
-            png: self,
+    }
-            start: 0,
+
-            end: 0,
+    /// Return an iterator over the scanlines of the image
-            pass: None,
+    #[inline]
-        }
+    pub fn scan_lines_mut(&mut self) -> ScanLinesMut {
+        ScanLinesMut::new(self)
     }
 
     /// Reverse all filters applied on the image, returning an unfiltered IDAT bytestream
@@ -358,171 +359,98 @@ impl PngData {
             return false;
         }
 
-        // A map of old indexes to new ones, for any moved
+        let mut palette_map = [0u8; 256];
-        let mut index_map: HashMap<u8, u8> = HashMap::new();
+        let mut used = [false; 256];
-
-        let mut palette = match self.palette {
-            Some(ref p) => p.clone(),
-            None => return false,
-        };
-
-        // A list of (original) indices that are duplicates and no longer needed
-        let mut duplicates: Vec<u8> = Vec::new();
         {
-            // Find duplicate entries in the palette
+            let palette = match self.palette {
-            let mut seen: HashMap<RGBA8, u8> = HashMap::with_capacity(palette.len());
+                Some(ref p) => p,
-            for (i, color) in palette.iter().cloned().enumerate() {
+                None => return false,
-                if seen.contains_key(&color) {
+            };
-                    let index = seen[&color];
+
-                    duplicates.push(i as u8);
+            // Find palette entries that are never used
-                    index_map.insert(i as u8, index);
+            for line in self.scan_lines() {
-                } else {
+                match self.ihdr_data.bit_depth {
-                    seen.insert(color, i as u8);
+                    BitDepth::Eight => for &byte in line.data {
+                        used[byte as usize] = true;
+                    },
+                    BitDepth::Four => for &byte in line.data {
+                        used[(byte & 0x0F) as usize] = true;
+                        used[(byte >> 4) as usize] = true;
+                    },
+                    BitDepth::Two => for &byte in line.data {
+                        used[(byte & 0x03) as usize] = true;
+                        used[((byte >> 2) & 0x03) as usize] = true;
+                        used[((byte >> 4) & 0x03) as usize] = true;
+                        used[(byte >> 6) as usize] = true;
+                    },
+                    _ => unreachable!(),
                 }
             }
+
+            let mut next_index = 0;
+            let mut seen = HashMap::with_capacity(palette.len());
+            for (i, (used, palette_map)) in used.iter().cloned().zip(palette_map.iter_mut()).enumerate() {
+                if !used {
+                    continue;
+                }
+                // There are invalid files that use pixel indices beyond palette size
+                let color = palette.get(i).cloned().unwrap_or(RGBA8::new(0,0,0,255));
+                match seen.entry(color) {
+                    Vacant(new) => {
+                        *palette_map = next_index;
+                        new.insert(next_index);
+                        next_index += 1;
+                    },
+                    Occupied(remap_to) => {
+                        *palette_map = *remap_to.get();
+                    },
+                }
+            }
+            if (0..palette.len()).all(|i| palette_map[i] == i as u8) {
+                return false;
+            }
         }
 
-        // Remove duplicates from the data
+        self.do_palette_reduction(&palette_map, &used);
-        if !duplicates.is_empty() {
-            self.do_palette_reduction(&duplicates, &mut index_map, &mut palette);
-        }
-
-        // Find palette entries that are never used
-        let mut seen = HashSet::with_capacity(palette.len());
-        for line in self.scan_lines() {
-            match self.ihdr_data.bit_depth {
-                BitDepth::Eight => for &byte in line.data {
-                    seen.insert(byte);
-                },
-                BitDepth::Four => {
-                    let bitvec = BitVec::from_bytes(&line.data);
-                    let mut current = 0u8;
-                    for (i, bit) in bitvec.iter().enumerate() {
-                        let mod_i = i % 4;
-                        if bit {
-                            current += 1u8 << (3 - mod_i);
-                        }
-                        if mod_i == 3 {
-                            seen.insert(current);
-                            current = 0;
-                        }
-                    }
-                }
-                BitDepth::Two => {
-                    let bitvec = BitVec::from_bytes(&line.data);
-                    let mut current = 0u8;
-                    for (i, bit) in bitvec.iter().enumerate() {
-                        let mod_i = i % 2;
-                        if bit {
-                            current += 1u8 << (1 - mod_i);
-                        }
-                        if mod_i == 1 {
-                            seen.insert(current);
-                            current = 0;
-                        }
-                    }
-                }
-                _ => unreachable!(),
-            }
-
-            if seen.len() == palette.len() {
-                // Exit early if no further possible optimizations
-                // Check at the end of each line
-                // Checking after every pixel would be overly expensive
-                return !duplicates.is_empty();
-            }
-        }
-
-        let unused: Vec<u8> = (0..palette.len() as u8)
-            .filter(|i| !seen.contains(i))
-            .collect();
-
-        // Remove unused palette indices
-        self.do_palette_reduction(&unused, &mut index_map, &mut palette);
-
         true
     }
 
-    fn do_palette_reduction(
+    fn do_palette_reduction(&mut self, palette_map: &[u8; 256], used: &[bool; 256]) {
-        &mut self,
+        let mut byte_map = *palette_map;
-        indices_to_remove: &[u8],
+
-        index_map: &mut HashMap<u8, u8>,
+        // low bit-depths can be pre-computed for every byte value
-        palette: &mut Vec<RGBA8>,
+        match self.ihdr_data.bit_depth {
-    ) {
+            BitDepth::Four => for byte in 0..=255 {
-        let mut new_data = Vec::with_capacity(self.raw_data.len());
+                byte_map[byte as usize] = palette_map[(byte & 0x0F) as usize] |
-        let original_len = palette.len();
+                    (palette_map[(byte >> 4) as usize] << 4);
-        for idx in indices_to_remove.iter().cloned().sorted_by(|a, b| b.cmp(a)) {
+            },
-            for i in (idx as usize + 1)..original_len {
+            BitDepth::Two => for byte in 0..=255 {
-                let existing = index_map.entry(i as u8).or_insert(i as u8);
+                byte_map[byte as usize] = palette_map[(byte & 0x03) as usize] |
-                if *existing >= idx {
+                    (palette_map[((byte >> 2) & 0x03) as usize] << 2) |
-                    *existing -= 1;
+                    (palette_map[((byte >> 4) & 0x03) as usize] << 4) |
-                }
+                    (palette_map[((byte >> 6)) as usize] << 6);
-            }
+            },
-            palette.remove(idx as usize);
+            _ => {}
         }
 
         // Reassign data bytes to new indices
-        for line in self.scan_lines() {
+        for line in self.scan_lines_mut() {
-            new_data.push(line.filter);
+            for byte in line.data {
-            match self.ihdr_data.bit_depth {
+                *byte = byte_map[*byte as usize];
-                BitDepth::Eight => for &byte in line.data {
-                    if let Some(&new_idx) = index_map.get(&byte) {
-                        new_data.push(new_idx);
-                    } else {
-                        new_data.push(byte);
-                    }
-                },
-                BitDepth::Four => for &byte in line.data {
-                    let upper = byte & 0b1111_0000;
-                    let lower = byte & 0b0000_1111;
-                    let mut new_byte = 0u8;
-                    new_byte |= if let Some(&new_idx) = index_map.get(&(upper >> 4)) {
-                        new_idx << 4
-                    } else {
-                        upper
-                    };
-                    new_byte |= if let Some(&new_idx) = index_map.get(&lower) {
-                        new_idx
-                    } else {
-                        lower
-                    };
-                    new_data.push(new_byte);
-                },
-                BitDepth::Two => for &byte in line.data {
-                    let one = byte & 0b1100_0000;
-                    let two = byte & 0b0011_0000;
-                    let three = byte & 0b0000_1100;
-                    let four = byte & 0b0000_0011;
-                    let mut new_byte = 0u8;
-                    new_byte |= if let Some(&new_idx) = index_map.get(&(one >> 6)) {
-                        new_idx << 6
-                    } else {
-                        one
-                    };
-                    new_byte |= if let Some(&new_idx) = index_map.get(&(two >> 4)) {
-                        new_idx << 4
-                    } else {
-                        two
-                    };
-                    new_byte |= if let Some(&new_idx) = index_map.get(&(three >> 2)) {
-                        new_idx << 2
-                    } else {
-                        three
-                    };
-                    new_byte |= if let Some(&new_idx) = index_map.get(&four) {
-                        new_idx
-                    } else {
-                        four
-                    };
-                    new_data.push(new_byte);
-                },
-                _ => unreachable!(),
             }
         }
-        index_map.clear();
+
-        self.raw_data = new_data;
         self.transparency_pixel = None;
-        self.palette = Some(palette.clone());
+        if let Some(palette) = self.palette.take() {
+            let max_index = palette_map.iter().max().cloned().unwrap_or(0) as usize;
+            let mut new_palette = vec![RGBA8::new(0,0,0,255); max_index+1];
+            for (color, (map_to, used)) in palette.into_iter().zip(palette_map.iter().cloned().zip(used.iter().cloned())) {
+                if used {
+                    new_palette[map_to as usize] = color;
+                }
+            }
+            self.palette = Some(new_palette);
+        }
     }
 
     /// Attempt to reduce the color type of the image

src/png/scan_lines.rs

@@ -3,129 +3,162 @@ 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
-    pub png: &'a PngData,
+    raw_data: &'a [u8],
-    pub start: usize,
+}
-    pub end: usize,
+
-    /// Current pass number, and 0-indexed row within the pass
+impl<'a> ScanLines<'a> {
-    pub pass: Option<(u8, u32)>,
+    pub fn new(png: &'a PngData) -> Self {
+        Self {
+            iter: ScanLineRanges::new(png),
+            raw_data: &png.raw_data,
+        }
+    }
 }
 
 impl<'a> Iterator for ScanLines<'a> {
     type Item = ScanLine<'a>;
+    #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        if self.end == self.png.raw_data.len() {
+        self.iter.next().map(|(len, pass)| {
-            None
+            let (data, rest) = self.raw_data.split_at(len);
-        } else if self.png.ihdr_data.interlaced == 1 {
+            self.raw_data = rest;
-            // Scanlines for interlaced PNG files
+            let (&filter, data) = data.split_first().unwrap();
-            if self.pass.is_none() {
+            ScanLine {
-                self.pass = Some((1, 0));
+                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.png.ihdr_data.width < 5 && self.pass.unwrap().0 == 2 {
+            if self.width < 5 && pass.0 == 2 {
-                if let Some(pass) = self.pass.as_mut() {
+                pass.0 = 3;
-                    pass.0 = 3;
+                pass.1 = 4;
-                    pass.1 = 4;
-                }
             }
             // Intentionally keep these separate so that they can be applied one after another
-            if self.png.ihdr_data.height < 5 && self.pass.unwrap().0 == 3 {
+            if self.height < 5 && pass.0 == 3 {
-                if let Some(pass) = self.pass.as_mut() {
+                pass.0 = 4;
-                    pass.0 = 4;
+                pass.1 = 0;
-                    pass.1 = 0;
-                }
             }
-            let bits_per_pixel = u32::from(self.png.ihdr_data.bit_depth.as_u8())
+            let (pixels_factor, y_steps) = match pass {
-                * u32::from(self.png.channels_per_pixel());
+                (1, _) | (2, _) => (8, 8),
-            let y_steps;
+                (3, _) => (4, 8),
-            let pixels_factor;
+                (4, _) => (4, 4),
-            match self.pass {
+                (5, _) => (2, 4),
-                Some((1, _)) | Some((2, _)) => {
+                (6, _) => (2, 2),
-                    pixels_factor = 8;
+                (7, _) => (1, 2),
-                    y_steps = 8;
-                }
-                Some((3, _)) => {
-                    pixels_factor = 4;
-                    y_steps = 8;
-                }
-                Some((4, _)) => {
-                    pixels_factor = 4;
-                    y_steps = 4;
-                }
-                Some((5, _)) => {
-                    pixels_factor = 2;
-                    y_steps = 4;
-                }
-                Some((6, _)) => {
-                    pixels_factor = 2;
-                    y_steps = 2;
-                }
-                Some((7, _)) => {
-                    pixels_factor = 1;
-                    y_steps = 2;
-                }
                 _ => unreachable!(),
-            }
-            let mut pixels_per_line = self.png.ihdr_data.width / pixels_factor as u32;
-            // Determine whether to add pixels if there is a final, incomplete 8x8 block
-            let gap = self.png.ihdr_data.width % pixels_factor;
-            if gap > 0 {
-                match self.pass.unwrap().0 {
-                    1 | 3 | 5 => {
-                        pixels_per_line += 1;
-                    }
-                    2 if gap >= 5 => {
-                        pixels_per_line += 1;
-                    }
-                    4 if gap >= 3 => {
-                        pixels_per_line += 1;
-                    }
-                    6 if gap >= 2 => {
-                        pixels_per_line += 1;
-                    }
-                    _ => (),
-                };
-            }
-            let current_pass = if let Some(pass) = self.pass {
-                Some(pass.0)
-            } else {
-                None
             };
-            let bytes_per_line = ((pixels_per_line * bits_per_pixel + 7) / 8) as usize;
+            let mut pixels_per_line = self.width / pixels_factor as u32;
-            self.start = self.end;
+            // Determine whether to add pixels if there is a final, incomplete 8x8 block
-            self.end = self.start + bytes_per_line + 1;
+            let gap = self.width % pixels_factor;
-            if let Some(pass) = self.pass.as_mut() {
+            match pass.0 {
-                if pass.1 + y_steps >= self.png.ihdr_data.height {
+                1 | 3 | 5 if gap > 0 => {
-                    pass.0 += 1;
+                    pixels_per_line += 1;
-                    pass.1 = match pass.0 {
-                        3 => 4,
-                        5 => 2,
-                        7 => 1,
-                        _ => 0,
-                    };
-                } else {
-                    pass.1 += y_steps;
                 }
+                2 if gap >= 5 => {
+                    pixels_per_line += 1;
+                }
+                4 if gap >= 3 => {
+                    pixels_per_line += 1;
+                }
+                6 if gap >= 2 => {
+                    pixels_per_line += 1;
+                }
+                _ => (),
+            };
+            let current_pass = Some(pass.0);
+            if pass.1 + y_steps >= self.height {
+                pass.0 += 1;
+                pass.1 = match pass.0 {
+                    3 => 4,
+                    5 => 2,
+                    7 => 1,
+                    _ => 0,
+                };
+            } else {
+                pass.1 += y_steps;
             }
-            Some(ScanLine {
+            (pixels_per_line, current_pass)
-                filter: self.png.raw_data[self.start],
-                data: &self.png.raw_data[(self.start + 1)..self.end],
-                pass: current_pass,
-            })
         } else {
             // Standard, non-interlaced PNG scanlines
-            let bits_per_line = self.png.ihdr_data.width as usize
+            (self.width, None)
-                * self.png.ihdr_data.bit_depth.as_u8() as usize
+        };
-                * self.png.channels_per_pixel() as usize;
+        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 bytes_per_line = ((bits_per_line + 7) / 8) as usize;
-            self.start = self.end;
+        let len = bytes_per_line + 1;
-            self.end = self.start + bytes_per_line + 1;
+        self.left -= len;
-            Some(ScanLine {
+        Some((len, current_pass))
-                filter: self.png.raw_data[self.start],
-                data: &self.png.raw_data[(self.start + 1)..self.end],
-                pass: None,
-            })
-        }
     }
 }
 
@@ -139,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>,
+}