Skip to content

Commit

Permalink
Enable customizing Zstd decoding parameters.
Browse files Browse the repository at this point in the history
  • Loading branch information
quantatic committed Jul 14, 2024
1 parent 7c2d530 commit c829e5d
Show file tree
Hide file tree
Showing 6 changed files with 86 additions and 1 deletion.
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ memchr = "2"
pin-project-lite = "0.2"
tokio = { version = "1.24.2", optional = true, default-features = false }
xz2 = { version = "0.1.6", optional = true }
zstd-safe = { version = "7", optional = true, default-features = false }
zstd-safe = { version = "7.1", optional = true, default-features = false }
deflate64 = { version = "0.1.5", optional = true }

[dev-dependencies]
Expand Down
10 changes: 10 additions & 0 deletions src/codec/zstd/decoder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,16 @@ impl ZstdDecoder {
}
}

pub(crate) fn new_with_params(params: &[crate::zstd::DParameter]) -> Self {
let mut decoder = Decoder::new().unwrap();
for param in params {
decoder.set_parameter(param.as_zstd()).unwrap();
}
Self {
decoder: Unshared::new(decoder),
}
}

pub(crate) fn new_with_dict(dictionary: &[u8]) -> io::Result<Self> {
let mut decoder = Decoder::with_dictionary(dictionary)?;
Ok(Self {
Expand Down
11 changes: 11 additions & 0 deletions src/macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,17 @@ macro_rules! algos {
}
}
{ @dec
/// Creates a new decoder, using the specified parameters, which will read compressed
/// data from the given stream and emit a decompressed stream.
pub fn with_params(inner: $inner, params: &[crate::zstd::DParameter]) -> Self {
Self {
inner: crate::$($mod::)+generic::Decoder::new(
inner,
crate::codec::ZstdDecoder::new_with_params(params),
),
}
}

/// Creates a new decoder, using the specified compression level and pre-trained
/// dictionary, which will read compressed data from the given stream and emit an
/// uncompressed stream.
Expand Down
20 changes: 20 additions & 0 deletions src/zstd.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
//! This module contains zstd-specific types for async-compression.
use libzstd::stream::raw::CParameter::*;
use libzstd::stream::raw::DParameter::*;

/// A compression parameter for zstd. This is a stable wrapper around zstd's own `CParameter`
/// type, to abstract over different versions of the zstd library.
Expand Down Expand Up @@ -110,3 +111,22 @@ impl CParameter {
self.0
}
}

/// A decompression parameter for zstd. This is a stable wrapper around zstd's own `DParameter`
/// type, to abstract over different versions of the zstd library.
///
/// See the [zstd documentation](https://facebook.github.io/zstd/zstd_manual.html) for more
/// information on these parameters.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct DParameter(libzstd::stream::raw::DParameter);

impl DParameter {
/// Maximum window size in bytes (as a power of two)
pub fn window_log_max(value: u32) -> Self {
Self(WindowLogMax(value))
}

pub(crate) fn as_zstd(&self) -> libzstd::stream::raw::DParameter {
self.0
}
}
Binary file added tests/artifacts/long-window-size-lib.rs.zst
Binary file not shown.
44 changes: 44 additions & 0 deletions tests/zstd.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,3 +2,47 @@
mod utils;

test_cases!(zstd);

use async_compression::zstd::DParameter;
use tokio::io::AsyncWriteExt as _;

Check failure on line 7 in tests/zstd.rs

View workflow job for this annotation

GitHub Actions / cargo hack check --all-targets --feature-powerset

unused import: `tokio::io::AsyncWriteExt`

#[tokio::test]
async fn zstd_decode_large_window_size_default() {
let compressed = include_bytes!("./artifacts/long-window-size-lib.rs.zst");

// Default decoder should throw with an error, window size maximum is too low.
let mut decoder = async_compression::tokio::write::ZstdDecoder::new(Vec::new());

Check failure on line 14 in tests/zstd.rs

View workflow job for this annotation

GitHub Actions / cargo hack check --all-targets --feature-powerset

failed to resolve: could not find `tokio` in `async_compression`
decoder.write_all(compressed).await.unwrap_err();
}

#[tokio::test]
async fn zstd_decode_large_window_size_explicit_small_window_size() {
let compressed = include_bytes!("./artifacts/long-window-size-lib.rs.zst");

// Short window decoder should throw with an error, window size maximum is too low.
let mut decoder = async_compression::tokio::write::ZstdDecoder::with_params(

Check failure on line 23 in tests/zstd.rs

View workflow job for this annotation

GitHub Actions / cargo hack check --all-targets --feature-powerset

failed to resolve: could not find `tokio` in `async_compression`
Vec::new(),
&[DParameter::window_log_max(16)],
);
decoder.write_all(compressed).await.unwrap_err();
}

#[tokio::test]
async fn zstd_decode_large_window_size_explicit_large_window_size() {
let compressed = include_bytes!("./artifacts/long-window-size-lib.rs.zst");
let source = include_bytes!("./artifacts/lib.rs");

// Long window decoder should succeed as the window size is large enough to decompress the given input.
let mut long_window_size_decoder = async_compression::tokio::write::ZstdDecoder::with_params(

Check failure on line 36 in tests/zstd.rs

View workflow job for this annotation

GitHub Actions / cargo hack check --all-targets --feature-powerset

failed to resolve: could not find `tokio` in `async_compression`
Vec::new(),
&[DParameter::window_log_max(31)],
);
// Long window size decoder should successfully decode the given input data.
long_window_size_decoder
.write_all(compressed)
.await
.unwrap();
long_window_size_decoder.shutdown().await.unwrap();

assert_eq!(long_window_size_decoder.into_inner(), source);
}

0 comments on commit c829e5d

Please sign in to comment.