Retain LZW decoder across tiles

src/decoder/image.rs

@@ -3,6 +3,7 @@ use super::stream::{PackBitsReader, ReadSamples};
 use super::tag_reader::TagReader;
 use super::ChunkType;
 use super::{predict_f16, predict_f32, predict_f64, ValueReader};
+use crate::decoder::stream::CompressionEngines;
 use crate::tags::{
     ByteOrder, CompressionMethod, ExtraSamples, PhotometricInterpretation, PlanarConfiguration,
     Predictor, SampleFormat, Tag,
@@ -698,6 +699,7 @@ impl Image {
         dimensions: (u32, u32),
         samples: u16,
         #[cfg_attr(not(feature = "fax"), allow(unused_variables))] fill_order: u16,
+        engines: &mut CompressionEngines,
     ) -> TiffResult<Box<dyn ReadSamples + 'r>> {
         Ok(match compression_method {
             CompressionMethod::None => Box::new(reader),
@@ -729,6 +731,7 @@ impl Image {
             CompressionMethod::LZW => Box::new(super::stream::LZWReader::new(
                 reader,
                 usize::try_from(compressed_length)?,
+                engines,
             )),
             #[cfg(feature = "zstd")]
             CompressionMethod::ZSTD => Box::new(zstd::Decoder::new(reader)?),
@@ -1033,6 +1036,7 @@ impl Image {
         layout: &ReadoutLayout,
         chunk_index: u32,
         scratch: &mut Vec<u8>,
+        engines: &mut CompressionEngines,
     ) -> TiffResult<()> {
         let ValueReader {
             reader,
@@ -1179,6 +1183,7 @@ impl Image {
             chunk_dims,
             self.samples,
             self.fill_order,
+            engines,
         )?;
 
         // Extended bit depths (9-15, 17-31): packed N-bit samples must be unpacked to
@@ -1403,6 +1408,7 @@ impl Image {
             }
         }
 
+        reader.cache_or_destroy(engines);
         Ok(())
     }
 

src/decoder/logluv.rs

@@ -83,6 +83,14 @@ impl<R: Read> Read for LogLuv24<R> {
     }
 }
 
+impl<R: Read> super::stream::ReadSamples for LogLuv24<R> {
+    fn read_padded_samples(&mut self, buf: &mut [u8], discard: u64) -> std::io::Result<()> {
+        self.read_exact(buf)?;
+        std::io::copy(&mut self.by_ref().take(discard), &mut std::io::sink())?;
+        Ok(())
+    }
+}
+
 fn packbits_decode<R: Read, const N: usize>(
     reader: &mut R,
     buf: &mut [[u8; N]],
@@ -212,6 +220,14 @@ impl<R: Read> Read for LogLuv32<R> {
     }
 }
 
+impl<R: Read> super::stream::ReadSamples for LogLuv32<R> {
+    fn read_padded_samples(&mut self, buf: &mut [u8], discard: u64) -> std::io::Result<()> {
+        self.read_exact(buf)?;
+        std::io::copy(&mut self.by_ref().take(discard), &mut std::io::sink())?;
+        Ok(())
+    }
+}
+
 /// Pack-Bits like encoding for a single 16-bit log-L sample.
 pub struct LogLuv16<R> {
     reader: Take<R>,
@@ -278,6 +294,14 @@ impl<R: Read> Read for LogLuv16<R> {
     }
 }
 
+impl<R: Read> super::stream::ReadSamples for LogLuv16<R> {
+    fn read_padded_samples(&mut self, buf: &mut [u8], discard: u64) -> std::io::Result<()> {
+        self.read_exact(buf)?;
+        std::io::copy(&mut self.by_ref().take(discard), &mut std::io::sink())?;
+        Ok(())
+    }
+}
+
 // Table information from: `libtiff`, 1997 by Greg Ward Larson, SGI
 #[doc(alias = "UV_SQSIZ")]
 const UV_QUANTIZATION_WIDTH: f32 = 0.003500;

src/decoder/mod.rs

@@ -259,6 +259,7 @@ pub struct Decoder<R> {
     scratch: Vec<u8>,
     /// Leniency flag that allows libtiff-like TIFF color interpretation behavior
     lenient: bool,
+    compression: stream::CompressionEngines,
 }
 
 /// All the information needed to read and interpret byte slices from the underlying file, i.e. to
@@ -853,6 +854,7 @@ impl TiffHeader {
             image: None,
             scratch: Vec::new(),
             lenient: false,
+            compression: stream::CompressionEngines::default(),
         }
     }
 }
@@ -949,12 +951,18 @@ impl<R: Read + Seek> Decoder<R> {
     #[expect(clippy::unnecessary_unwrap)] // Lifetimes would overlap due to return.
     fn ensure_image_and_reader(
         &mut self,
-    ) -> TiffResult<(&mut Image, &mut ValueReader<R>, &mut Vec<u8>)> {
+    ) -> TiffResult<(
+        &mut Image,
+        &mut ValueReader<R>,
+        &mut Vec<u8>,
+        &mut stream::CompressionEngines,
+    )> {
         if self.image.is_some() {
             return Ok((
                 self.image.as_mut().unwrap(),
                 &mut self.value_reader,
                 &mut self.scratch,
+                &mut self.compression,
             ));
         }
 
@@ -963,6 +971,7 @@ impl<R: Read + Seek> Decoder<R> {
             self.image.insert(image),
             &mut self.value_reader,
             &mut self.scratch,
+            &mut self.compression,
         ))
     }
 
@@ -1095,6 +1104,7 @@ impl<R: Read + Seek> Decoder<R> {
             image: core::mem::take(&mut self.image),
             scratch: core::mem::take(&mut self.scratch),
             lenient: self.lenient,
+            compression: core::mem::take(&mut self.compression),
         }
     }
 
@@ -1578,10 +1588,10 @@ impl<R: Read + Seek> Decoder<R> {
         chunk_index: u32,
         layout: &image::ReadoutLayout,
     ) -> TiffResult<()> {
-        let (image, value_reader, scratch) = self.ensure_image_and_reader()?;
+        let (image, value_reader, scratch, engines) = self.ensure_image_and_reader()?;
         let offset = image.chunk_file_range(chunk_index)?.0;
         value_reader.reader.goto_offset(offset)?;
-        image.expand_chunk(value_reader, buffer, layout, chunk_index, scratch)?;
+        image.expand_chunk(value_reader, buffer, layout, chunk_index, scratch, engines)?;
         Ok(())
     }
 
@@ -1733,7 +1743,7 @@ impl<R: Read + Seek> Decoder<R> {
         slice: TiffCodingUnit,
         buffer: &mut [u8],
     ) -> TiffResult<()> {
-        let (image, value_reader, scratch) = self.ensure_image_and_reader()?;
+        let (image, value_reader, scratch, engines) = self.ensure_image_and_reader()?;
 
         let (width, height) = image.chunk_data_dimensions(slice.0)?;
         let readout = image.readout_for_size(width, height)?;
@@ -1766,6 +1776,7 @@ impl<R: Read + Seek> Decoder<R> {
                 readout,
                 chunk,
                 scratch,
+                engines,
             )?;
         }
 
@@ -1937,7 +1948,7 @@ impl<R: Read + Seek> Decoder<R> {
     /// Returns an error if the buffer fits less than one plane. In particular, for non-planar
     /// images returns an error if the buffer does not fit the required size.
     pub fn read_image_bytes(&mut self, buffer: &mut [u8]) -> TiffResult<()> {
-        let (image, value_reader, scratch) = self.ensure_image_and_reader()?;
+        let (image, value_reader, scratch, engines) = self.ensure_image_and_reader()?;
         let readout = image.readout_for_image()?;
 
         let ref layout @ image::PlaneLayout {
@@ -1971,6 +1982,7 @@ impl<R: Read + Seek> Decoder<R> {
                     readout,
                     chunk,
                     scratch,
+                    engines,
                 )?;
             }
         }

src/decoder/stream.rs

@@ -122,6 +122,21 @@ impl<R: Read> EndianReader<R> {
 // # READERS
 //
 
+#[derive(Default)]
+pub(crate) struct CompressionEngines {
+    #[cfg(feature = "lzw")]
+    lzw: Option<weezl::decode::Decoder>,
+}
+
+impl core::fmt::Debug for CompressionEngines {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut dbg = f.debug_struct("CompressionEngines");
+        #[cfg(feature = "lzw")]
+        dbg.field("lzw", &self.lzw.is_some().then_some("<weezl>"));
+        dbg.finish()
+    }
+}
+
 /// A variant of `io::Read` that takes one block to the end.
 ///
 /// This is derived from `Read` so that any special optimization done by `std` (or any nightly crate
@@ -142,19 +157,51 @@ pub trait ReadSamples: Read {
     ///
     /// This is automatically implemented for `Self: Read` as a copy into `io::Sink`.
     fn read_padded_samples(&mut self, buf: &mut [u8], discard: u64) -> io::Result<()>;
+
+    /// Put this back into our cache of compression engines.
+    fn cache_or_destroy(self: Box<Self>, _: &mut CompressionEngines) {}
 }
 
-impl<R: Read> ReadSamples for R {
+/// Note: the `Seek` bound ensures this does not overlap actual decompression.. This is an odd
+/// design problem born out of wanting both the underlying `Read` for uncompressed streams due to
+/// performance as well as the override to `cache_or_destroy` for all compressed streams.
+impl<R: Read + Seek> ReadSamples for R {
     fn read_padded_samples(&mut self, buf: &mut [u8], discard: u64) -> io::Result<()> {
         self.read_exact(buf)?;
-        std::io::copy(&mut self.by_ref().take(discard), &mut std::io::sink())?;
+        if discard > 0 {
+            std::io::copy(&mut self.by_ref().take(discard), &mut std::io::sink())?;
+        }
         Ok(())
     }
 }
 
 /// Type alias for the deflate Reader
 #[cfg(feature = "deflate")]
-pub type DeflateReader<R> = flate2::read::ZlibDecoder<R>;
+pub struct DeflateReader<R> {
+    inner: flate2::read::ZlibDecoder<R>,
+}
+
+impl<R: Read> DeflateReader<R> {
+    pub fn new(inner: R) -> Self {
+        DeflateReader {
+            inner: flate2::read::ZlibDecoder::new(inner),
+        }
+    }
+}
+
+impl<R: Read> Read for DeflateReader<R> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        self.inner.read(buf)
+    }
+}
+
+impl<R: Read> ReadSamples for DeflateReader<R> {
+    fn read_padded_samples(&mut self, buf: &mut [u8], discard: u64) -> io::Result<()> {
+        self.read_exact(buf)?;
+        std::io::copy(&mut self.by_ref().take(discard), &mut std::io::sink())?;
+        Ok(())
+    }
+}
 
 //
 // ## LZW Reader
@@ -165,12 +212,17 @@ pub type DeflateReader<R> = flate2::read::ZlibDecoder<R>;
 pub struct LZWReader<R: Read> {
     reader: BufReader<Take<R>>,
     decoder: weezl::decode::Decoder,
+    is_old_style: bool,
 }
 
 #[cfg(feature = "lzw")]
 impl<R: Read> LZWReader<R> {
     /// Wraps a reader
-    pub fn new(reader: R, compressed_length: usize) -> LZWReader<R> {
+    pub fn new(
+        reader: R,
+        compressed_length: usize,
+        engine: &mut CompressionEngines,
+    ) -> LZWReader<R> {
         let mut buffered = BufReader::with_capacity(
             (32 * 1024).min(compressed_length),
             reader.take(u64::try_from(compressed_length).unwrap()),
@@ -188,26 +240,75 @@ impl<R: Read> LZWReader<R> {
             Ok(head) if head.len() >= 2 && head[0] == 0x00 && head[1] & 0x01 != 0
         );
 
-        let configuration = if old_style {
+        let decoder = if old_style {
             weezl::decode::Configuration::new(weezl::BitOrder::Lsb, 8)
+                .with_yield_on_full_buffer(true)
+                .build()
         } else {
-            weezl::decode::Configuration::with_tiff_size_switch(weezl::BitOrder::Msb, 8)
-        }
-        .with_yield_on_full_buffer(true);
+            // All standard LZW decoders have the same configuration. We can thus reuse any that we
+            // have previously constructed, skipping internal buffer allocation. We only need to
+            // reset them to a new stream.
+            if let Some(mut lzw) = engine.lzw.take() {
+                lzw.reset();
+                lzw
+            } else {
+                weezl::decode::Configuration::with_tiff_size_switch(weezl::BitOrder::Msb, 8)
+                    .with_yield_on_full_buffer(true)
+                    .build()
+            }
+        };
 
         Self {
             reader: buffered,
-            decoder: configuration.build(),
+            decoder: decoder,
+            is_old_style: old_style,
         }
     }
+
+    fn discard_n(&mut self, mut n: u64) -> io::Result<()> {
+        let mut buf = [0u8; 2048];
+
+        while n > 0 {
+            let limit = n.min(2048) as usize;
+            let result = self
+                .decoder
+                .decode_bytes(self.reader.fill_buf()?, &mut buf[..limit]);
+
+            self.reader.consume(result.consumed_in);
+            n -= result.consumed_out as u64;
+
+            match result.status {
+                Ok(weezl::LzwStatus::Ok) => {
+                    continue;
+                }
+                Ok(weezl::LzwStatus::NoProgress) => {
+                    return Err(io::Error::new(
+                        io::ErrorKind::UnexpectedEof,
+                        "no lzw end code found",
+                    ));
+                }
+                Ok(weezl::LzwStatus::Done) => {
+                    if n != 0 {
+                        // stream is too short.
+                        return Err(std::io::ErrorKind::UnexpectedEof)?;
+                    }
+                }
+                Err(err) => return Err(io::Error::new(io::ErrorKind::InvalidData, err)),
+            }
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(feature = "lzw")]
 impl<R: Read> Read for LZWReader<R> {
-    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
         loop {
             let result = self.decoder.decode_bytes(self.reader.fill_buf()?, buf);
+
             self.reader.consume(result.consumed_in);
+            let _ = buf.split_off_mut(..result.consumed_out);
 
             match result.status {
                 Ok(weezl::LzwStatus::Ok) => {
@@ -232,6 +333,24 @@ impl<R: Read> Read for LZWReader<R> {
     }
 }
 
+impl<R: Read> ReadSamples for LZWReader<R> {
+    fn read_padded_samples(&mut self, buf: &mut [u8], discard: u64) -> io::Result<()> {
+        self.read_exact(buf)?;
+
+        if discard > 0 {
+            self.discard_n(discard)?;
+        }
+
+        Ok(())
+    }
+
+    fn cache_or_destroy(self: Box<Self>, engines: &mut CompressionEngines) {
+        if !self.is_old_style && engines.lzw.is_none() {
+            engines.lzw = Some(self.decoder);
+        }
+    }
+}
+
 //
 // ## PackBits Reader
 //
@@ -304,6 +423,14 @@ impl<R: Read> Read for PackBitsReader<R> {
     }
 }
 
+impl<R: Read> ReadSamples for PackBitsReader<R> {
+    fn read_padded_samples(&mut self, buf: &mut [u8], discard: u64) -> io::Result<()> {
+        self.read_exact(buf)?;
+        std::io::copy(&mut self.by_ref().take(discard), &mut std::io::sink())?;
+        Ok(())
+    }
+}
+
 /// Iterator adapter that reverses bits in each byte for FillOrder=2 (LSB-to-MSB).
 #[cfg(feature = "fax")]
 type FaxBytes<R> =
@@ -420,6 +547,14 @@ impl<R: Read> Read for Group4Reader<R> {
     }
 }
 
+impl<R: Read> ReadSamples for Group4Reader<R> {
+    fn read_padded_samples(&mut self, buf: &mut [u8], discard: u64) -> io::Result<()> {
+        self.read_exact(buf)?;
+        std::io::copy(&mut self.by_ref().take(discard), &mut std::io::sink())?;
+        Ok(())
+    }
+}
+
 /// Owning byte reader for Group 3 decoding: the compressed data is consumed
 /// into an iterator so the reader doesn't borrow the buffer.
 #[cfg(feature = "fax")]
@@ -558,6 +693,14 @@ impl Read for Group3Reader {
     }
 }
 
+impl ReadSamples for Group3Reader {
+    fn read_padded_samples(&mut self, buf: &mut [u8], discard: u64) -> io::Result<()> {
+        self.read_exact(buf)?;
+        std::io::copy(&mut self.by_ref().take(discard), &mut std::io::sink())?;
+        Ok(())
+    }
+}
+
 #[cfg(feature = "webp")]
 pub struct WebPReader {
     inner: Cursor<Vec<u8>>,
@@ -613,6 +756,15 @@ impl Read for WebPReader {
     }
 }
 
+#[cfg(feature = "webp")]
+impl ReadSamples for WebPReader {
+    fn read_padded_samples(&mut self, buf: &mut [u8], discard: u64) -> io::Result<()> {
+        self.read_exact(buf)?;
+        std::io::copy(&mut self.by_ref().take(discard), &mut std::io::sink())?;
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;