diff --git a/crates/clawhdf5-format/src/filters.rs b/crates/clawhdf5-format/src/filters.rs index 36f2e96..872b8d8 100644 --- a/crates/clawhdf5-format/src/filters.rs +++ b/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, + b: Vec, + #[cfg(feature = "deflate")] + inflater: Option, +} + +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) = 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, 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, +) -> 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(); - out.try_reserve_exact(size_hint.clamp(1, max_capacity)) + // A kept buffer may already be larger than `max_capacity`; the decoder + // 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( - &data[in_before as usize..], - &mut out, - FlushDecompress::Finish, - ) + .decompress_vec(&data[in_before as usize..], 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, 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, 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) { 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]; - result.reserve_exact(tail.len()); + // Every byte of `result[..whole]` is overwritten below, so a reused + // 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, Forma } } match element_size { - 2 => interleave::<2>(data, num_elements, &mut result), - 4 => interleave::<4>(data, num_elements, &mut result), - 8 => interleave::<8>(data, num_elements, &mut result), - 16 => interleave::<16>(data, num_elements, &mut result), + 2 => interleave::<2>(data, num_elements, result), + 4 => interleave::<4>(data, num_elements, result), + 8 => interleave::<8>(data, num_elements, 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, 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, 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 { 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, 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| 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 = (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]