Add decode_all and decode_all_to_vec (#70)

Author: philipc

src/frame_decoder.rs

@@ -102,6 +102,7 @@ pub enum FrameDecoderError {
     NotYetInitialized,
     FailedToInitialize(frame::FrameHeaderError),
     FailedToDrainDecodebuffer(Error),
+    FailedToSkipFrame,
     TargetTooSmall,
     DictNotProvided { dict_id: u32 },
 }
@@ -164,6 +165,12 @@ impl core::fmt::Display for FrameDecoderError {
                     e,
                 )
             }
+            FrameDecoderError::FailedToSkipFrame => {
+                write!(
+                    f,
+                    "Failed to skip bytes for the length given in the frame header"
+                )
+            }
             FrameDecoderError::TargetTooSmall => {
                 write!(f, "Target must have at least as many bytes as the contentsize of the frame reports")
             }
@@ -605,6 +612,87 @@ impl FrameDecoder {
         let read_len = bytes_read_at_end - bytes_read_at_start;
         Ok((read_len as usize, result_len))
     }
+
+    /// Decode multiple frames into the output slice.
+    ///
+    /// `input` must contain an exact number of frames.
+    ///
+    /// `output` must be large enough to hold the decompressed data. If you don't know
+    /// how large the output will be, use [`FrameDecoder::decode_blocks`] instead.
+    ///
+    /// This calls [`FrameDecoder::init`], and all bytes currently in the decoder will be lost.
+    ///
+    /// Returns the number of bytes written to `output`.
+    pub fn decode_all(
+        &mut self,
+        mut input: &[u8],
+        mut output: &mut [u8],
+    ) -> Result<usize, FrameDecoderError> {
+        let mut total_bytes_written = 0;
+        while !input.is_empty() {
+            match self.init(&mut input) {
+                Ok(_) => {}
+                Err(FrameDecoderError::ReadFrameHeaderError(
+                    frame::ReadFrameHeaderError::SkipFrame { length, .. },
+                )) => {
+                    input = input
+                        .get(length as usize..)
+                        .ok_or(FrameDecoderError::FailedToSkipFrame)?;
+                    continue;
+                }
+                Err(e) => return Err(e),
+            };
+            loop {
+                self.decode_blocks(&mut input, BlockDecodingStrategy::UptoBlocks(1))?;
+                let bytes_written = self
+                    .read(output)
+                    .map_err(FrameDecoderError::FailedToDrainDecodebuffer)?;
+                output = &mut output[bytes_written..];
+                total_bytes_written += bytes_written;
+                if self.can_collect() != 0 {
+                    return Err(FrameDecoderError::TargetTooSmall);
+                }
+                if self.is_finished() {
+                    break;
+                }
+            }
+        }
+
+        Ok(total_bytes_written)
+    }
+
+    /// Decode multiple frames into the extra capacity of the output vector.
+    ///
+    /// `input` must contain an exact number of frames.
+    ///
+    /// `output` must have enough extra capacity to hold the decompressed data.
+    /// This function will not reallocate or grow the vector. If you don't know
+    /// how large the output will be, use [`FrameDecoder::decode_blocks`] instead.
+    ///
+    /// This calls [`FrameDecoder::init`], and all bytes currently in the decoder will be lost.
+    ///
+    /// The length of the output vector is updated to include the decompressed data.
+    /// The length is not changed if an error occurs.
+    pub fn decode_all_to_vec(
+        &mut self,
+        input: &[u8],
+        output: &mut Vec<u8>,
+    ) -> Result<(), FrameDecoderError> {
+        let len = output.len();
+        let cap = output.capacity();
+        output.resize(cap, 0);
+        match self.decode_all(input, &mut output[len..]) {
+            Ok(bytes_written) => {
+                let new_len = core::cmp::min(len + bytes_written, cap); // Sanitizes `bytes_written`.
+                output.resize(new_len, 0);
+                Ok(())
+            }
+            Err(e) => {
+                output.resize(len, 0);
+                Err(e)
+            }
+        }
+    }
 }
 
 /// Read bytes from the decode_buffer that are no longer needed. While the frame is not yet finished

src/tests/mod.rs

@@ -481,6 +481,91 @@ fn test_streaming_no_std() {
     }
 }
 
+#[test]
+fn test_decode_all() {
+    use crate::frame_decoder::{FrameDecoder, FrameDecoderError};
+
+    let skip_frame = |input: &mut Vec<u8>, length: usize| {
+        input.extend_from_slice(&0x184D2A50u32.to_le_bytes());
+        input.extend_from_slice(&(length as u32).to_le_bytes());
+        input.resize(input.len() + length, 0);
+    };
+
+    let mut original = Vec::new();
+    let mut input = Vec::new();
+
+    skip_frame(&mut input, 300);
+    input.extend_from_slice(include_bytes!("../../decodecorpus_files/z000089.zst"));
+    original.extend_from_slice(include_bytes!("../../decodecorpus_files/z000089"));
+    skip_frame(&mut input, 400);
+    input.extend_from_slice(include_bytes!("../../decodecorpus_files/z000090.zst"));
+    original.extend_from_slice(include_bytes!("../../decodecorpus_files/z000090"));
+    skip_frame(&mut input, 500);
+
+    let mut decoder = FrameDecoder::new();
+
+    // decode_all with correct buffers.
+    let mut output = vec![0; original.len()];
+    let result = decoder.decode_all(&input, &mut output).unwrap();
+    assert_eq!(result, original.len());
+    assert_eq!(output, original);
+
+    // decode_all with smaller output length.
+    let mut output = vec![0; original.len() - 1];
+    let result = decoder.decode_all(&input, &mut output);
+    assert!(
+        matches!(result, Err(FrameDecoderError::TargetTooSmall)),
+        "{:?}",
+        result
+    );
+
+    // decode_all with larger output length.
+    let mut output = vec![0; original.len() + 1];
+    let result = decoder.decode_all(&input, &mut output).unwrap();
+    assert_eq!(result, original.len());
+    assert_eq!(&output[..result], original);
+
+    // decode_all with truncated regular frame.
+    let mut output = vec![0; original.len()];
+    let result = decoder.decode_all(&input[..input.len() - 600], &mut output);
+    assert!(
+        matches!(result, Err(FrameDecoderError::FailedToReadBlockBody(_))),
+        "{:?}",
+        result
+    );
+
+    // decode_all with truncated skip frame.
+    let mut output = vec![0; original.len()];
+    let result = decoder.decode_all(&input[..input.len() - 1], &mut output);
+    assert!(
+        matches!(result, Err(FrameDecoderError::FailedToSkipFrame)),
+        "{:?}",
+        result
+    );
+
+    // decode_all_to_vec with correct output capacity.
+    let mut output = Vec::new();
+    output.reserve_exact(original.len());
+    decoder.decode_all_to_vec(&input, &mut output).unwrap();
+    assert_eq!(output, original);
+
+    // decode_all_to_vec with smaller output capacity.
+    let mut output = Vec::new();
+    output.reserve_exact(original.len() - 1);
+    let result = decoder.decode_all_to_vec(&input, &mut output);
+    assert!(
+        matches!(result, Err(FrameDecoderError::TargetTooSmall)),
+        "{:?}",
+        result
+    );
+
+    // decode_all_to_vec with larger output capacity.
+    let mut output = Vec::new();
+    output.reserve_exact(original.len() + 1);
+    decoder.decode_all_to_vec(&input, &mut output).unwrap();
+    assert_eq!(output, original);
+}
+
 pub mod bit_reader;
 pub mod decode_corpus;
 pub mod dict_test;