Refactor utils::unpack_bits for palette expanded images

src/decoder/mod.rs

@@ -619,12 +619,10 @@ impl<R: Read> Reader<R> {
         };
         match (color_type, trns) {
             (ColorType::Indexed, _) if expand => {
-                output_buffer[..row.len()].copy_from_slice(row);
-                expand_paletted(output_buffer, info, trns)?;
+                expand_paletted(row, output_buffer, info, trns)?;
             }
             (ColorType::Grayscale | ColorType::GrayscaleAlpha, _) if bit_depth < 8 && expand => {
-                output_buffer[..row.len()].copy_from_slice(row);
-                expand_gray_u8(output_buffer, info, trns)
+                expand_gray_u8(row, output_buffer, info, trns)
             }
             (ColorType::Grayscale | ColorType::Rgb, Some(trns)) if expand => {
                 let channels = color_type.samples();
@@ -811,6 +809,7 @@ impl SubframeInfo {
 }
 
 fn expand_paletted(
+    row: &[u8],
     buffer: &mut [u8],
     info: &Info,
     trns: Option<Option<&[u8]>>,
@@ -842,7 +841,7 @@ fn expand_paletted(
                     &[]
                 };
 
-                utils::unpack_bits(buffer, 4, info.bit_depth as u8, |i, chunk| {
+                utils::unpack_bits(row, buffer, 4, info.bit_depth as u8, |i, chunk| {
                     let (rgb, a) = (
                         palette
                             .get(3 * i as usize..3 * i as usize + 3)
@@ -855,7 +854,7 @@ fn expand_paletted(
                     chunk[3] = a;
                 });
             } else {
-                utils::unpack_bits(buffer, 3, info.bit_depth as u8, |i, chunk| {
+                utils::unpack_bits(row, buffer, 3, info.bit_depth as u8, |i, chunk| {
                     let rgb = palette
                         .get(3 * i as usize..3 * i as usize + 3)
                         .unwrap_or(&black);
@@ -873,15 +872,15 @@ fn expand_paletted(
     }
 }
 
-fn expand_gray_u8(buffer: &mut [u8], info: &Info, trns: Option<Option<&[u8]>>) {
+fn expand_gray_u8(row: &[u8], buffer: &mut [u8], info: &Info, trns: Option<Option<&[u8]>>) {
     let rescale = true;
     let scaling_factor = if rescale {
         (255) / ((1u16 << info.bit_depth as u8) - 1) as u8
     } else {
         1
     };
     if let Some(trns) = trns {
-        utils::unpack_bits(buffer, 2, info.bit_depth as u8, |pixel, chunk| {
+        utils::unpack_bits(row, buffer, 2, info.bit_depth as u8, |pixel, chunk| {
             chunk[1] = if let Some(trns) = trns {
                 if pixel == trns[0] {
                     0
@@ -894,7 +893,7 @@ fn expand_gray_u8(buffer: &mut [u8], info: &Info, trns: Option<Option<&[u8]>>) {
             chunk[0] = pixel * scaling_factor
         })
     } else {
-        utils::unpack_bits(buffer, 1, info.bit_depth as u8, |val, chunk| {
+        utils::unpack_bits(row, buffer, 1, info.bit_depth as u8, |val, chunk| {
             chunk[0] = val * scaling_factor
         })
     }

src/utils.rs

@@ -1,40 +1,52 @@
 //! Utility functions
-use std::iter::{repeat, StepBy};
+use std::iter::StepBy;
 use std::ops::Range;
 
 #[inline(always)]
-pub fn unpack_bits<F>(buf: &mut [u8], channels: usize, bit_depth: u8, func: F)
+pub fn unpack_bits<F>(input: &[u8], output: &mut [u8], channels: usize, bit_depth: u8, func: F)
 where
     F: Fn(u8, &mut [u8]),
 {
-    // Return early if empty. This enables to subtract `channels` later without overflow.
-    if buf.len() < channels {
-        return;
-    }
+    // Only [1, 2, 4, 8] are valid bit depths
+    assert!(matches!(bit_depth, 1 | 2 | 4 | 8));
+    // Check that `input` is capable of producing a buffer as long as `output`:
+    // number of shift lookups per bit depth * channels * input length
+    assert!((8 / bit_depth as usize * channels).saturating_mul(input.len()) >= output.len());
+
+    let mut buf_chunks = output.chunks_exact_mut(channels);
+    let mut iter = input.iter();
+
+    // `shift` iterates through the corresponding bit depth sequence:
+    // 1 => &[7, 6, 5, 4, 3, 2, 1, 0],
+    // 2 => &[6, 4, 2, 0],
+    // 4 => &[4, 0],
+    // 8 => &[0],
+    //
+    // `(0..8).step_by(bit_depth.into()).rev()` doesn't always optimize well so
+    // shifts are calculated instead. (2023-08, Rust 1.71)
+
+    if bit_depth == 8 {
+        for (&curr, chunk) in iter.zip(&mut buf_chunks) {
+            func(curr, chunk);
+        }
+    } else {
+        let mask = ((1u16 << bit_depth) - 1) as u8;
 
-    let bits = buf.len() / channels * bit_depth as usize;
-    let extra_bits = bits % 8;
-    let entries = bits / 8
-        + match extra_bits {
-            0 => 0,
-            _ => 1,
-        };
-    let skip = match extra_bits {
-        0 => 0,
-        n => (8 - n) / bit_depth as usize,
-    };
-    let mask = ((1u16 << bit_depth) - 1) as u8;
-    let i = (0..entries)
-        .rev() // reverse iterator
-        .flat_map(|idx|
-            // this has to be reversed too
-            (0..8).step_by(bit_depth.into())
-            .zip(repeat(idx)))
-        .skip(skip);
-    let j = (0..=buf.len() - channels).rev().step_by(channels);
-    for ((shift, i), j) in i.zip(j) {
-        let pixel = (buf[i] & (mask << shift)) >> shift;
-        func(pixel, &mut buf[j..(j + channels)])
+        // These variables are initialized in the loop
+        let mut shift = -1;
+        let mut curr = 0;
+
+        for chunk in buf_chunks {
+            if shift < 0 {
+                shift = 8 - bit_depth as i32;
+                curr = *iter.next().expect("input for unpack bits is not empty");
+            }
+
+            let pixel = (curr >> shift) & mask;
+            func(pixel, chunk);
+
+            shift -= bit_depth as i32;
+        }
     }
 }