crates/clawhdf5-format/src/filters.rs

@@ -148,6 +148,180 @@ pub fn decompress_chunk_exact(
     Ok(data)
 }
 
+/// Buffers a chunk decoder keeps between chunks, so decoding a dataset's
+/// chunks one after another reuses the same memory instead of allocating
+/// (and faulting in) fresh buffers for every chunk and every filter stage.
+///
+/// Use one per thread with [`decompress_chunk_exact_with`]. Buffers larger
+/// than [`DecodeScratch::RETAIN_BYTES`] are released by
+/// [`DecodeScratch::trim`], so a scratch kept for a long time (a
+/// thread-local, say) does not hold on to a huge chunk's memory.
+#[derive(Default)]
+pub struct DecodeScratch {
+    a: Vec<u8>,
+    b: Vec<u8>,
+    #[cfg(feature = "deflate")]
+    inflater: Option<flate2::Decompress>,
+}
+
+impl core::fmt::Debug for DecodeScratch {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("DecodeScratch")
+            .field("a_capacity", &self.a.capacity())
+            .field("b_capacity", &self.b.capacity())
+            .finish_non_exhaustive()
+    }
+}
+
+/// Which buffer holds the data between two filter stages.
+#[derive(Clone, Copy)]
+enum Stage {
+    /// `compressed[..len]`: still (a prefix of) the stored bytes.
+    Stored(usize),
+    A,
+    B,
+}
+
+impl DecodeScratch {
+    /// Largest buffer [`trim`](Self::trim) keeps (4 MiB).
+    pub const RETAIN_BYTES: usize = 4 << 20;
+
+    /// An empty scratch; buffers are allocated on first use.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Release any buffer larger than [`Self::RETAIN_BYTES`].
+    pub fn trim(&mut self) {
+        for buf in [&mut self.a, &mut self.b] {
+            if buf.capacity() > Self::RETAIN_BYTES {
+                *buf = Vec::new();
+            }
+        }
+    }
+}
+
+/// [`decompress_chunk_exact`] into reusable buffers: the decoded chunk is
+/// returned as a slice of `scratch` (or of `compressed`, when every filter
+/// that was applied only appended a checksum), valid until `scratch` is used
+/// again.
+///
+/// Deflate, shuffle and Fletcher32 — h5py's and libhdf5's usual pipeline —
+/// decode without allocating once `scratch` has grown to the chunk size: the
+/// inflater writes into a kept buffer (and its state is reset, not
+/// rebuilt), shuffle interleaves into the other buffer, and Fletcher32 checks
+/// the checksum and drops it in place. Every other filter goes through the
+/// filter registry as [`decompress_chunk_masked`] does, and its output
+/// replaces a scratch buffer. The result is byte for byte what
+/// [`decompress_chunk_exact`] returns, with the same errors.
+pub fn decompress_chunk_exact_with<'s>(
+    compressed: &'s [u8],
+    pipeline: &FilterPipeline,
+    chunk_size: usize,
+    element_size: u32,
+    filter_mask: u32,
+    coords: &[u64],
+    scratch: &'s mut DecodeScratch,
+) -> Result<&'s [u8], FormatError> {
+    // Same per-stage bounds as `decompress_chunk_masked`.
+    let mut bounds = [0usize; 32];
+    let mut bounds_vec = Vec::new();
+    let bounds: &mut [usize] = if pipeline.filters.len() <= bounds.len() {
+        &mut bounds[..pipeline.filters.len()]
+    } else {
+        bounds_vec.resize(pipeline.filters.len(), 0);
+        &mut bounds_vec
+    };
+    let mut size = chunk_size;
+    for (i, filter) in pipeline.filters.iter().enumerate() {
+        bounds[i] = size;
+        if !filter_skipped(filter_mask, i) {
+            size = filter_output_bound(filter.filter_id, size);
+        }
+    }
+
+    let mut stage = Stage::Stored(compressed.len());
+    for (i, filter) in pipeline.filters.iter().enumerate().rev() {
+        if filter_skipped(filter_mask, i) {
+            continue;
+        }
+        let ctx = FilterContext {
+            filter,
+            element_size: element_size as usize,
+            max_output: bounds[i],
+        };
+        // The stage's input, and the buffer its output goes to (the one
+        // not holding the input).
+        let (input, out): (&[u8], &mut Vec<u8>) = match stage {
+            Stage::Stored(len) => (&compressed[..len], &mut scratch.a),
+            Stage::A => (&scratch.a, &mut scratch.b),
+            Stage::B => (&scratch.b, &mut scratch.a),
+        };
+        let next = match stage {
+            Stage::Stored(_) | Stage::B => Stage::A,
+            Stage::A => Stage::B,
+        };
+        match filter.filter_id {
+            // Built in and never overridable (`register_filter` refuses
+            // built-in IDs), so the registry would pick exactly these.
+            FILTER_FLETCHER32 => {
+                // Check and drop the checksum where the data is.
+                let payload = fletcher32_payload(input)?;
+                stage = match stage {
+                    Stage::Stored(_) => Stage::Stored(payload),
+                    Stage::A => {
+                        scratch.a.truncate(payload);
+                        Stage::A
+                    }
+                    Stage::B => {
+                        scratch.b.truncate(payload);
+                        Stage::B
+                    }
+                };
+                continue;
+            }
+            FILTER_SHUFFLE => shuffle_decompress_into(input, ctx.element_size, out),
+            #[cfg(all(
+                feature = "deflate",
+                not(all(target_os = "macos", feature = "system-zlib-decompress"))
+            ))]
+            FILTER_DEFLATE => {
+                let limit = if ctx.max_output != 0 {
+                    ctx.max_output
+                } else {
+                    MAX_DECOMPRESS_SIZE
+                };
+                let size_hint = if ctx.max_output != 0 {
+                    ctx.max_output
+                } else {
+                    input.len().saturating_mul(4).min(1 << 20)
+                };
+                let inflater = scratch
+                    .inflater
+                    .get_or_insert_with(|| flate2::Decompress::new(true));
+                inflater.reset(true);
+                inflate_bounded_into(inflater, input, size_hint, limit, out)
+                    .map_err(FormatError::DecompressionError)?;
+            }
+            _ => *out = filter_registry::decode(input, &ctx)?,
+        }
+        stage = next;
+    }
+
+    let data: &[u8] = match stage {
+        Stage::Stored(len) => &compressed[..len],
+        Stage::A => &scratch.a,
+        Stage::B => &scratch.b,
+    };
+    if chunk_size != 0 && data.len() != chunk_size {
+        return Err(FormatError::ChunkedReadError(format!(
+            "chunk at {coords:?} decoded to {} bytes, expected {chunk_size}",
+            data.len()
+        )));
+    }
+    Ok(data)
+}
+
 /// Apply a filter pipeline to compress a chunk.
 /// Filters are applied in FORWARD order for compression.
 pub fn compress_chunk(
@@ -892,33 +1066,55 @@ pub(crate) fn inflate_bounded(
     size_hint: usize,
     limit: usize,
 ) -> Result<Vec<u8>, String> {
-    use flate2::{Decompress, FlushDecompress, Status};
+    let mut out = Vec::new();
+    inflate_bounded_into(
+        &mut flate2::Decompress::new(true),
+        data,
+        size_hint,
+        limit,
+        &mut out,
+    )?;
+    Ok(out)
+}
+
+/// [`inflate_bounded`] with a fresh or reset `inflater`, into `out`: its
+/// contents are replaced and its allocation reused.
+#[cfg(feature = "deflate")]
+fn inflate_bounded_into(
+    inflater: &mut flate2::Decompress,
+    data: &[u8],
+    size_hint: usize,
+    limit: usize,
+    out: &mut Vec<u8>,
+) -> Result<(), String> {
+    use flate2::{FlushDecompress, Status};
 
     // One byte of headroom past the limit distinguishes an over-size stream
     // from one that legitimately ends exactly at the limit.
     let max_capacity = limit.saturating_add(1);
-    let mut out = Vec::new();
+    // A kept buffer may already be larger than `max_capacity`; the decoder
-    out.try_reserve_exact(size_hint.clamp(1, max_capacity))
+    // can then write past the limit, which the check below still refuses.
+    out.clear();
+    let want = size_hint.clamp(1, max_capacity);
+    out.try_reserve_exact(want)
         .map_err(|e| format!("deflate: cannot allocate output: {e}"))?;
 
-    let mut inflater = Decompress::new(true);
     loop {
         let (in_before, out_before) = (inflater.total_in(), inflater.total_out());
         let status = inflater
-            .decompress_vec(
+            .decompress_vec(&data[in_before as usize..], out, FlushDecompress::Finish)
-                &data[in_before as usize..],
-                &mut out,
-                FlushDecompress::Finish,
-            )
             .map_err(|e| format!("deflate: {e}"))?;
         if out.len() > limit {
             return Err("deflate: output exceeds size limit".into());
         }
         match status {
-            Status::StreamEnd => return Ok(out),
+            Status::StreamEnd => return Ok(()),
             Status::Ok | Status::BufError if out.len() == out.capacity() => {
                 // Out of room: double, up to the limit.
-                let grow = out.capacity().min(max_capacity - out.capacity()).max(1);
+                let grow = out
+                    .capacity()
+                    .min(max_capacity.saturating_sub(out.capacity()))
+                    .max(1);
                 out.try_reserve_exact(grow)
                     .map_err(|e| format!("deflate: cannot allocate output: {e}"))?;
             }
@@ -1217,16 +1413,30 @@ fn zstd_compress(data: &[u8], level: u32) -> Result<Vec<u8>, FormatError> {
 /// On disk: all byte-0s of each element together, then all byte-1s, etc.
 /// Output: elements in natural order.
 fn shuffle_decompress(data: &[u8], element_size: usize) -> Result<Vec<u8>, FormatError> {
+    let mut result = Vec::new();
+    shuffle_decompress_into(data, element_size, &mut result);
+    Ok(result)
+}
+
+/// [`shuffle_decompress`] into `result`, replacing its contents and reusing
+/// its allocation.
+fn shuffle_decompress_into(data: &[u8], element_size: usize, result: &mut Vec<u8>) {
     if element_size <= 1 {
-        return Ok(data.to_vec());
+        result.clear();
+        result.extend_from_slice(data);
+        return;
     }
     // Like libhdf5, only whole elements are shuffled; trailing bytes (e.g. a
     // Fletcher32 checksum appended before the shuffle) are stored as-is.
     let whole = data.len() - data.len() % element_size;
     let (data, tail) = data.split_at(whole);
     let num_elements = data.len() / element_size;
-    let mut result = vec![0u8; whole];
+    // Every byte of `result[..whole]` is overwritten below, so a reused
-    result.reserve_exact(tail.len());
+    // buffer keeps its old bytes instead of being zeroed first; only growth
+    // is zero-filled.
+    result.truncate(whole);
+    result.reserve_exact(whole + tail.len() - result.len());
+    result.resize(whole, 0);
 
     // The shuffled stream is `element_size` byte planes of `num_elements`
     // bytes each; un-shuffling interleaves them. This is on the read path of
@@ -1245,10 +1455,10 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result<Vec<u8>, Forma
         }
     }
     match element_size {
-        2 => interleave::<2>(data, num_elements, &mut result),
+        2 => interleave::<2>(data, num_elements, result),
-        4 => interleave::<4>(data, num_elements, &mut result),
+        4 => interleave::<4>(data, num_elements, result),
-        8 => interleave::<8>(data, num_elements, &mut result),
+        8 => interleave::<8>(data, num_elements, result),
-        16 => interleave::<16>(data, num_elements, &mut result),
+        16 => interleave::<16>(data, num_elements, result),
         _ => {
             for (i, element) in result.chunks_exact_mut(element_size).enumerate() {
                 for (j, byte) in element.iter_mut().enumerate() {
@@ -1258,8 +1468,6 @@ fn shuffle_decompress(data: &[u8], element_size: usize) -> Result<Vec<u8>, Forma
         }
     }
     result.extend_from_slice(tail);
-
-    Ok(result)
 }
 
 /// Shuffle (compress direction): group bytes by position within each element.
@@ -1411,6 +1619,12 @@ fn fletcher32_compute(data: &[u8]) -> u32 {
 /// Verify Fletcher32 checksum and strip it from the data.
 /// The last 4 bytes are the stored checksum.
 fn fletcher32_verify(data: &[u8]) -> Result<Vec<u8>, FormatError> {
+    fletcher32_payload(data).map(|len| data[..len].to_vec())
+}
+
+/// Verify the Fletcher32 checksum that ends `data`; the length of the data
+/// before it.
+fn fletcher32_payload(data: &[u8]) -> Result<usize, FormatError> {
     if data.len() < 4 {
         return Err(FormatError::FilterError(
             "fletcher32: data too short for checksum".into(),
@@ -1430,7 +1644,7 @@ fn fletcher32_verify(data: &[u8]) -> Result<Vec<u8>, FormatError> {
             computed,
         });
     }
-    Ok(payload.to_vec())
+    Ok(payload.len())
 }
 
 /// Append Fletcher32 checksum to data.
@@ -1545,6 +1759,113 @@ fn pcodec_decompress(
 #[cfg(test)]
 mod tests {
 
+    /// `decompress_chunk_exact_with` returns exactly what
+    /// `decompress_chunk_exact` returns — data or error — for every pipeline
+    /// shape, filter mask and chunk size, with one scratch reused across all
+    /// of them in an order that grows, shrinks and swaps its buffers.
+    #[test]
+    fn decode_with_scratch_matches_the_allocating_decoder() {
+        let f = |filter_id: u16, client_data: Vec<u32>| FilterDescription {
+            filter_id,
+            name: None,
+            flags: 0,
+            client_data,
+        };
+        let mut pipelines = vec![
+            vec![f(FILTER_SHUFFLE, vec![4])],
+            vec![f(FILTER_FLETCHER32, vec![])],
+            // NetCDF-4's order: the checksum is taken before shuffle.
+            vec![f(FILTER_FLETCHER32, vec![]), f(FILTER_SHUFFLE, vec![4])],
+        ];
+        #[cfg(feature = "deflate")]
+        pipelines.extend([
+            vec![f(FILTER_DEFLATE, vec![4])],
+            vec![f(FILTER_SHUFFLE, vec![4]), f(FILTER_DEFLATE, vec![4])],
+            // h5py's order with `fletcher32=True`: checksum last.
+            vec![
+                f(FILTER_SHUFFLE, vec![4]),
+                f(FILTER_DEFLATE, vec![4]),
+                f(FILTER_FLETCHER32, vec![]),
+            ],
+            vec![
+                f(FILTER_FLETCHER32, vec![]),
+                f(FILTER_SHUFFLE, vec![4]),
+                f(FILTER_DEFLATE, vec![1]),
+            ],
+        ]);
+        #[cfg(feature = "lzf")]
+        pipelines.push(vec![
+            f(FILTER_SHUFFLE, vec![4]),
+            f(crate::filter_pipeline::FILTER_LZF, vec![]),
+        ]);
+
+        let mut scratch = DecodeScratch::new();
+        for elements in [1usize, 7, 4096, 3, 65536, 100] {
+            let data: Vec<u8> = (0..elements as u32)
+                .flat_map(|i| (i.wrapping_mul(2654435761) >> (i % 13)).to_le_bytes())
+                .collect();
+            for filters in &pipelines {
+                let pipeline = FilterPipeline {
+                    version: 2,
+                    filters: filters.clone(),
+                };
+                let n = filters.len() as u32;
+                for mask in 0..(1u32 << n) {
+                    // Encode only the filters the mask says were applied.
+                    let mut stored = data.clone();
+                    for (i, filter) in filters.iter().enumerate() {
+                        if mask & (1 << i) == 0 {
+                            let ctx = FilterContext {
+                                filter,
+                                element_size: 4,
+                                max_output: 0,
+                            };
+                            stored = filter_registry::encode(&stored, &ctx).unwrap();
+                        }
+                    }
+                    let mut cases = vec![(stored.clone(), data.len())];
+                    // Corrupt: last byte flipped, truncated, wrong size.
+                    let mut flipped = stored.clone();
+                    *flipped.last_mut().unwrap() ^= 0x5a;
+                    cases.push((flipped, data.len()));
+                    cases.push((stored[..stored.len() / 2].to_vec(), data.len()));
+                    cases.push((stored.clone(), data.len() + 4));
+                    cases.push((stored.clone(), 0));
+                    for (bytes, size) in cases {
+                        let want = decompress_chunk_exact(&bytes, &pipeline, size, 4, mask, &[3]);
+                        let got = decompress_chunk_exact_with(
+                            &bytes,
+                            &pipeline,
+                            size,
+                            4,
+                            mask,
+                            &[3],
+                            &mut scratch,
+                        )
+                        .map(<[u8]>::to_vec);
+                        match (&want, &got) {
+                            (Ok(w), Ok(g)) => assert_eq!(w, g, "{filters:?} mask {mask}"),
+                            (Err(w), Err(g)) => {
+                                assert_eq!(w.to_string(), g.to_string(), "{filters:?}")
+                            }
+                            _ => panic!("{filters:?} mask {mask} size {size}: {want:?} vs {got:?}"),
+                        }
+                    }
+                }
+            }
+        }
+        // Long-lived scratch gives back a huge chunk's buffers.
+        let big = vec![0u8; DecodeScratch::RETAIN_BYTES + 8];
+        let shuffle = FilterPipeline {
+            version: 2,
+            filters: vec![f(FILTER_SHUFFLE, vec![4])],
+        };
+        decompress_chunk_exact_with(&big, &shuffle, big.len(), 4, 0, &[0], &mut scratch).unwrap();
+        scratch.trim();
+        assert!(scratch.a.capacity() <= DecodeScratch::RETAIN_BYTES);
+        assert!(scratch.b.capacity() <= DecodeScratch::RETAIN_BYTES);
+    }
+
     /// A chunk whose pipeline decodes to fewer bytes than the chunk holds is
     /// an error naming the chunk, never a short buffer the reader pads.
     #[test]