From 6de523783d21fe878e37ca98c2b406922363abf9 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sat, 23 May 2026 18:10:33 +0100 Subject: [PATCH 01/25] Add _zstd stdlib module backing compression.zstd - Implement ZstdDict, ZstdCompressor, and ZstdDecompressor on top of the zstd-safe crate - Wire the module into stdlib_module_defs (gated off Android/wasm32) - Mark test_zstd_multithread_compress as expected failure when libzstd lacks multi-threading support --- Cargo.toml | 2 + Lib/test/test_zstd.py | 7 +- crates/stdlib/Cargo.toml | 4 + crates/stdlib/src/lib.rs | 5 + crates/stdlib/src/zstd.rs | 1712 +++++++++++++++++++++++++++++++++++++ 5 files changed, 1728 insertions(+), 2 deletions(-) create mode 100644 crates/stdlib/src/zstd.rs diff --git a/Cargo.toml b/Cargo.toml index 4356442fada..5775e41b86c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -320,6 +320,8 @@ which = "8" x509-parser = "0.18" xml = "1.3" writeable = "0.6" +zstd = { version = "0.13", default-features = false, features = ["zdict_builder"] } +zstd-safe = { version = "7", default-features = false, features = ["zdict_builder"] } # Lints diff --git a/Lib/test/test_zstd.py b/Lib/test/test_zstd.py index cf618534add..1b4b819d2ef 100644 --- a/Lib/test/test_zstd.py +++ b/Lib/test/test_zstd.py @@ -321,8 +321,11 @@ def test_unknown_compression_parameter(self): with self.assertRaisesRegex(ValueError, pattern): ZstdCompressor(options=option) - @unittest.skipIf(not SUPPORT_MULTITHREADING, - "zstd build doesn't support multi-threaded compression") + # TODO: RUSTPYTHON + @unittest.expectedFailureIf( + not SUPPORT_MULTITHREADING, + "TODO: RUSTPYTHON; libzstd built without multi-threaded compression support", + ) def test_zstd_multithread_compress(self): size = 40*_1M b = THIS_FILE_BYTES * (size // len(THIS_FILE_BYTES)) diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index 5111380237f..2586138ac79 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -97,6 +97,10 @@ webpki-roots = { workspace = true, optional = true } [target.'cfg(not(any(target_os = "android", target_arch = "wasm32")))'.dependencies] libsqlite3-sys = { workspace = true, features = ["bundled"], optional = true } +xz = { workspace = true } +xz-sys = { workspace = true } +zstd = { workspace = true } +zstd-safe = { workspace = true } [target.'cfg(windows)'.dependencies] paste = { workspace = true } diff --git a/crates/stdlib/src/lib.rs b/crates/stdlib/src/lib.rs index 5eb8062ced7..1afd32b148d 100644 --- a/crates/stdlib/src/lib.rs +++ b/crates/stdlib/src/lib.rs @@ -28,6 +28,9 @@ mod lzma; mod zlib; +#[cfg(not(any(target_os = "android", target_arch = "wasm32")))] +mod zstd; + mod blake2; mod hashlib; mod md5; @@ -325,5 +328,7 @@ pub fn stdlib_module_defs(ctx: &Context) -> Vec<&'static builtins::PyModuleDef> ))] uuid::module_def(ctx), zlib::module_def(ctx), + #[cfg(not(any(target_os = "android", target_arch = "wasm32")))] + zstd::module_def(ctx), ] } diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs new file mode 100644 index 00000000000..6c9ad33396e --- /dev/null +++ b/crates/stdlib/src/zstd.rs @@ -0,0 +1,1712 @@ +// spell-checker:ignore cctx dctx CCTX DCTX ldm cdict ddict windowlog hashlog chainlog searchlog +// spell-checker:ignore minmatch dictid checksumflag dstream cstream pyobj zstandard btopt btultra +// spell-checker:ignore btlazy dfast nbworkers windowlogmax windowlog overlap targetcblock +// spell-checker:ignore srcsize zdict refprefix refcdict refddict pledgedsrcsize getframecontentsize +// spell-checker:ignore Zstd Zstandard pylib RFC + +//! The `_zstd` extension module. Backs the pure-Python `compression.zstd` +//! package by exposing the same classes, functions and constants that +//! CPython's `Modules/_zstd/` exposes. The Python wrapper at +//! `Lib/compression/zstd/__init__.py` imports from this module unconditionally, +//! so the names and call signatures here must stay in sync with CPython. +//! +//! Backend: the `zstd_safe` Rust crate (a thin safe wrapper over Facebook's +//! libzstd C library, which is what CPython links against). A handful of +//! routines that `zstd_safe` does not expose at a safe level (parameter +//! bounds, dictionary finalization) drop down to raw `zstd_sys` FFI calls. + +pub(crate) use _zstd::module_def; + +// The compression/decompression parameter and strategy constants below use +// CPython's `ZSTD_c_camelCase` / `ZSTD_d_camelCase` naming convention so the +// pure-Python `compression.zstd` package, which references them by those exact +// names, keeps working unchanged. +#[allow(non_upper_case_globals)] +#[pymodule] +mod _zstd { + use rustpython_common::lock::PyMutex; + use rustpython_vm::builtins::{ + PyBaseExceptionRef, PyBytesRef, PyDict, PyTupleRef, PyType, PyTypeRef, + }; + use rustpython_vm::function::{ArgBytesLike, OptionalArg}; + use rustpython_vm::types::{AsMapping, Constructor, Representable}; + use rustpython_vm::{ + AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, + }; + use std::ffi::c_int; + use zstd_safe::zstd_sys; + use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; + + /// Stores the `CompressionParameter` and `DecompressionParameter` enum + /// types that the pure-Python `compression.zstd` wrapper passes through + /// `set_parameter_types`. They let us produce CPython-compatible + /// `TypeError`s when callers pass the wrong enum kind to a (de)compressor + /// (a `CompressionParameter` key to a decompressor, or vice versa). Stored + /// in a `PyMutex` so the module init can set them once and constructors + /// can read them later from any thread. + static REGISTERED_PARAM_TYPES: PyMutex = + PyMutex::new(RegisteredParamTypes { + compression: None, + decompression: None, + }); + + struct RegisteredParamTypes { + compression: Option, + decompression: Option, + } + + // ========================================================================= + // Module-level constants + // ========================================================================= + + /// Default compression level used when none is requested. Matches + /// libzstd's `ZSTD_CLEVEL_DEFAULT`. + #[pyattr] + const ZSTD_CLEVEL_DEFAULT: i32 = zstd_sys::ZSTD_CLEVEL_DEFAULT as i32; + + // Compression parameter identifiers. Values match the `ZSTD_cParameter` + // enum in libzstd, which is what the public `CompressionParameter` IntEnum + // in `Lib/compression/zstd/__init__.py` derives its members from. + #[pyattr] + const ZSTD_c_compressionLevel: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_compressionLevel as i32; + #[pyattr] + const ZSTD_c_windowLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_windowLog as i32; + #[pyattr] + const ZSTD_c_hashLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_hashLog as i32; + #[pyattr] + const ZSTD_c_chainLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_chainLog as i32; + #[pyattr] + const ZSTD_c_searchLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_searchLog as i32; + #[pyattr] + const ZSTD_c_minMatch: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_minMatch as i32; + #[pyattr] + const ZSTD_c_targetLength: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_targetLength as i32; + #[pyattr] + const ZSTD_c_strategy: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_strategy as i32; + #[pyattr] + const ZSTD_c_enableLongDistanceMatching: i32 = + zstd_sys::ZSTD_cParameter::ZSTD_c_enableLongDistanceMatching as i32; + #[pyattr] + const ZSTD_c_ldmHashLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmHashLog as i32; + #[pyattr] + const ZSTD_c_ldmMinMatch: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmMinMatch as i32; + #[pyattr] + const ZSTD_c_ldmBucketSizeLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmBucketSizeLog as i32; + #[pyattr] + const ZSTD_c_ldmHashRateLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmHashRateLog as i32; + #[pyattr] + const ZSTD_c_contentSizeFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_contentSizeFlag as i32; + #[pyattr] + const ZSTD_c_checksumFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_checksumFlag as i32; + #[pyattr] + const ZSTD_c_dictIDFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_dictIDFlag as i32; + #[pyattr] + const ZSTD_c_nbWorkers: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_nbWorkers as i32; + #[pyattr] + const ZSTD_c_jobSize: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_jobSize as i32; + #[pyattr] + const ZSTD_c_overlapLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_overlapLog as i32; + + // Decompression parameter identifiers. libzstd only exposes one non- + // experimental decompression parameter. + #[pyattr] + const ZSTD_d_windowLogMax: i32 = zstd_sys::ZSTD_dParameter::ZSTD_d_windowLogMax as i32; + + // Strategy enum members ordered from fastest to strongest. These power + // the `Strategy` IntEnum in `Lib/compression/zstd/__init__.py`. + #[pyattr] + const ZSTD_fast: i32 = zstd_sys::ZSTD_strategy::ZSTD_fast as i32; + #[pyattr] + const ZSTD_dfast: i32 = zstd_sys::ZSTD_strategy::ZSTD_dfast as i32; + #[pyattr] + const ZSTD_greedy: i32 = zstd_sys::ZSTD_strategy::ZSTD_greedy as i32; + #[pyattr] + const ZSTD_lazy: i32 = zstd_sys::ZSTD_strategy::ZSTD_lazy as i32; + #[pyattr] + const ZSTD_lazy2: i32 = zstd_sys::ZSTD_strategy::ZSTD_lazy2 as i32; + #[pyattr] + const ZSTD_btlazy2: i32 = zstd_sys::ZSTD_strategy::ZSTD_btlazy2 as i32; + #[pyattr] + const ZSTD_btopt: i32 = zstd_sys::ZSTD_strategy::ZSTD_btopt as i32; + #[pyattr] + const ZSTD_btultra: i32 = zstd_sys::ZSTD_strategy::ZSTD_btultra as i32; + #[pyattr] + const ZSTD_btultra2: i32 = zstd_sys::ZSTD_strategy::ZSTD_btultra2 as i32; + + /// Runtime version string of the linked libzstd, like `"1.5.7"`. + #[pyattr(once, name = "zstd_version")] + fn zstd_version(_vm: &VirtualMachine) -> String { + zstd_safe::version_string().to_string() + } + + /// Packed integer version: `MAJOR * 10000 + MINOR * 100 + RELEASE`. + /// The Python wrapper unpacks this into a `zstd_version_info` tuple. + #[pyattr(once, name = "zstd_version_number")] + fn zstd_version_number(_vm: &VirtualMachine) -> u32 { + zstd_safe::version_number() + } + + /// Recommended output buffer size for decompression. The `ZstdFile.read1` + /// wrapper uses this as its default read size. + #[pyattr(once, name = "ZSTD_DStreamOutSize")] + fn zstd_dstream_out_size(_vm: &VirtualMachine) -> usize { + DCtx::out_size() + } + + // Dictionary load type markers. The `ZstdDict.as_*` properties wrap the + // dictionary in a `(zdict, marker)` tuple so the compressor or decompressor + // constructor knows which load mode to apply. Numbering matches CPython's + // `Modules/_zstd/_zstdmodule.h::DictType`. + const DICT_TYPE_DIGESTED: i32 = 0; + const DICT_TYPE_UNDIGESTED: i32 = 1; + const DICT_TYPE_PREFIX: i32 = 2; + + // ========================================================================= + // ZstdError exception + // ========================================================================= + + /// `_zstd.ZstdError`: the exception type raised for libzstd errors and for + /// invalid usage (calling `decompress` on a closed decompressor, etc.). + #[pyattr(once, name = "ZstdError")] + fn zstd_error(vm: &VirtualMachine) -> PyTypeRef { + vm.ctx.new_exception_type( + "_zstd", + "ZstdError", + Some(vec![vm.ctx.exceptions.exception_type.to_owned()]), + ) + } + + fn new_zstd_error(message: impl Into, vm: &VirtualMachine) -> PyBaseExceptionRef { + let msg: String = message.into(); + vm.new_exception_msg(vm.class("_zstd", "ZstdError"), msg.into()) + } + + /// Convert a libzstd error code (the `usize` returned by most ZSTD_* + /// functions when `ZSTD_isError(code)` is non-zero) into a `ZstdError` + /// carrying the human-readable message from `ZSTD_getErrorName`. + fn catch_zstd_error(code: usize, vm: &VirtualMachine) -> PyBaseExceptionRef { + new_zstd_error(zstd_safe::get_error_name(code).to_string(), vm) + } + + // ========================================================================= + // Parameter helpers + // ========================================================================= + + /// Convert any int-like Python value to an `i32`, rejecting floats. CPython + /// uses `PyLong_AsInt` here which only accepts integers; `try_index` is + /// the equivalent of CPython's `__index__` slot lookup (rejects floats). + fn pyobj_to_i32(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + obj.try_index(vm)?.try_to_primitive(vm) + } + + /// Normalize an `OptionalArg` constructor parameter to an + /// `Option`. Treats both "argument missing" and "argument + /// explicitly None" the same way (consistent with CPython, which uses + /// `None` defaults throughout the `_zstd` API). + fn arg_or_none(arg: OptionalArg, vm: &VirtualMachine) -> Option { + arg.into_option().filter(|o| !vm.is_none(o)) + } + + /// If `key` is a registered `CompressionParameter` or `DecompressionParameter` + /// enum member and its kind does not match the caller's context, return a + /// `TypeError`. Used by both (de)compressor option-dict loops so passing a + /// `CompressionParameter` to a `ZstdDecompressor` (or vice versa) is + /// rejected with the CPython-compatible "must not be a XParameter" + /// message. If `set_parameter_types` has not been called yet, this is a + /// no-op. + fn check_wrong_param_kind( + key: &PyObjectRef, + is_compress: bool, + vm: &VirtualMachine, + ) -> PyResult<()> { + let guard = REGISTERED_PARAM_TYPES.lock(); + let (forbidden_type, forbidden_name, kind) = if is_compress { + (&guard.decompression, "DecompressionParameter", "compression") + } else { + (&guard.compression, "CompressionParameter", "decompression") + }; + if let Some(ref t) = *forbidden_type + && key.fast_isinstance(t) + { + return Err(vm.new_type_error(format!( + "{kind} options dictionary key must not be a {forbidden_name}" + ))); + } + Ok(()) + } + + /// Map a compression parameter id to its raw libzstd `ZSTD_cParameter` + /// enum value. Returns `None` for unknown ids so callers can surface a + /// targeted `ValueError`. Done with an explicit match rather than + /// `mem::transmute` so passing junk like `ZSTD_cParameter(42)` cannot + /// be triggered from Python. + fn c_param_enum(param: i32) -> Option { + use zstd_sys::ZSTD_cParameter as P; + Some(match param { + ZSTD_c_compressionLevel => P::ZSTD_c_compressionLevel, + ZSTD_c_windowLog => P::ZSTD_c_windowLog, + ZSTD_c_hashLog => P::ZSTD_c_hashLog, + ZSTD_c_chainLog => P::ZSTD_c_chainLog, + ZSTD_c_searchLog => P::ZSTD_c_searchLog, + ZSTD_c_minMatch => P::ZSTD_c_minMatch, + ZSTD_c_targetLength => P::ZSTD_c_targetLength, + ZSTD_c_strategy => P::ZSTD_c_strategy, + ZSTD_c_enableLongDistanceMatching => P::ZSTD_c_enableLongDistanceMatching, + ZSTD_c_ldmHashLog => P::ZSTD_c_ldmHashLog, + ZSTD_c_ldmMinMatch => P::ZSTD_c_ldmMinMatch, + ZSTD_c_ldmBucketSizeLog => P::ZSTD_c_ldmBucketSizeLog, + ZSTD_c_ldmHashRateLog => P::ZSTD_c_ldmHashRateLog, + ZSTD_c_contentSizeFlag => P::ZSTD_c_contentSizeFlag, + ZSTD_c_checksumFlag => P::ZSTD_c_checksumFlag, + ZSTD_c_dictIDFlag => P::ZSTD_c_dictIDFlag, + ZSTD_c_nbWorkers => P::ZSTD_c_nbWorkers, + ZSTD_c_jobSize => P::ZSTD_c_jobSize, + ZSTD_c_overlapLog => P::ZSTD_c_overlapLog, + _ => return None, + }) + } + + /// Map a decompression parameter id to its raw libzstd `ZSTD_dParameter` + /// enum value. See [`c_param_enum`] for rationale. + fn d_param_enum(param: i32) -> Option { + use zstd_sys::ZSTD_dParameter as P; + match param { + ZSTD_d_windowLogMax => Some(P::ZSTD_d_windowLogMax), + _ => None, + } + } + + /// Map a compression-parameter id and integer value to a `CParameter` + /// variant. Used by the compressor's `options=` constructor argument. + fn cparameter_from_int(param: i32, value: i32, vm: &VirtualMachine) -> PyResult { + let p = match param { + ZSTD_c_compressionLevel => CParameter::CompressionLevel(value), + ZSTD_c_windowLog => CParameter::WindowLog(value as u32), + ZSTD_c_hashLog => CParameter::HashLog(value as u32), + ZSTD_c_chainLog => CParameter::ChainLog(value as u32), + ZSTD_c_searchLog => CParameter::SearchLog(value as u32), + ZSTD_c_minMatch => CParameter::MinMatch(value as u32), + ZSTD_c_targetLength => CParameter::TargetLength(value as u32), + ZSTD_c_strategy => CParameter::Strategy(strategy_from_int(value).ok_or_else( + || new_zstd_error(format!("invalid strategy value: {value}"), vm), + )?), + ZSTD_c_enableLongDistanceMatching => CParameter::EnableLongDistanceMatching(value != 0), + ZSTD_c_ldmHashLog => CParameter::LdmHashLog(value as u32), + ZSTD_c_ldmMinMatch => CParameter::LdmMinMatch(value as u32), + ZSTD_c_ldmBucketSizeLog => CParameter::LdmBucketSizeLog(value as u32), + ZSTD_c_ldmHashRateLog => CParameter::LdmHashRateLog(value as u32), + ZSTD_c_contentSizeFlag => CParameter::ContentSizeFlag(value != 0), + ZSTD_c_checksumFlag => CParameter::ChecksumFlag(value != 0), + ZSTD_c_dictIDFlag => CParameter::DictIdFlag(value != 0), + ZSTD_c_nbWorkers => CParameter::NbWorkers(value as u32), + ZSTD_c_jobSize => CParameter::JobSize(value as u32), + ZSTD_c_overlapLog => CParameter::OverlapSizeLog(value as u32), + _ => { + return Err(vm.new_value_error(format!( + "invalid compression parameter 'unknown parameter (key {param})'" + ))); + } + }; + Ok(p) + } + + /// Map a decompression-parameter id and integer value to a `DParameter`. + /// Used by the decompressor's `options=` constructor argument. + fn dparameter_from_int(param: i32, value: i32, vm: &VirtualMachine) -> PyResult { + match param { + ZSTD_d_windowLogMax => Ok(DParameter::WindowLogMax(value as u32)), + _ => Err(vm.new_value_error(format!( + "invalid decompression parameter 'unknown parameter (key {param})'" + ))), + } + } + + /// Build the CPython-compatible " parameter 'name' received an + /// illegal value V; the valid range is [lo, hi]" `ValueError` for a + /// parameter that was rejected by libzstd or that fell outside the + /// documented bounds. + fn param_value_error_for( + param: i32, + value: i32, + is_compress: bool, + vm: &VirtualMachine, + ) -> PyBaseExceptionRef { + let kind = if is_compress { "compression" } else { "decompression" }; + let name = parameter_name(param, is_compress); + match lookup_param_bounds(param, is_compress) { + Some((lo, hi)) => vm.new_value_error(format!( + "{kind} parameter '{name}' received an illegal value {value}; \ + the valid range is [{lo}, {hi}]" + )), + None => vm.new_value_error(format!( + "{kind} parameter '{name}' received an illegal value {value}" + )), + } + } + + /// Return the valid `(lower, upper)` bounds for the libzstd compression + /// level. Used when validating the `level=` argument upfront because + /// libzstd silently clamps out-of-range values rather than surfacing + /// them as errors. + fn level_bounds() -> (i32, i32) { + lookup_param_bounds(ZSTD_c_compressionLevel, true) + .expect("compressionLevel always has valid bounds") + } + + /// Look up parameter bounds for a known compression or decompression + /// parameter id. Returns `None` if the id is not recognized (callers + /// validate the id separately). + fn lookup_param_bounds(param: i32, is_compress: bool) -> Option<(i32, i32)> { + // SAFETY: `ZSTD_*Param_getBounds` reads no memory beyond the enum + // discriminant; the helpers above validated that `param` maps to a + // real enum variant. + let bounds = if is_compress { + let p = c_param_enum(param)?; + unsafe { zstd_sys::ZSTD_cParam_getBounds(p) } + } else { + let p = d_param_enum(param)?; + unsafe { zstd_sys::ZSTD_dParam_getBounds(p) } + }; + // SAFETY: ZSTD_isError just inspects the bounds.error integer. + if unsafe { zstd_sys::ZSTD_isError(bounds.error) } != 0 { + return None; + } + Some((bounds.lowerBound, bounds.upperBound)) + } + + /// Map a parameter integer id back to the Python-visible enum member + /// name. Used for error messages that pin-point the parameter that went + /// out of range. Returns `"unknown"` for unrecognized ids. + fn parameter_name(param: i32, is_compress: bool) -> &'static str { + if is_compress { + match param { + ZSTD_c_compressionLevel => "compression_level", + ZSTD_c_windowLog => "window_log", + ZSTD_c_hashLog => "hash_log", + ZSTD_c_chainLog => "chain_log", + ZSTD_c_searchLog => "search_log", + ZSTD_c_minMatch => "min_match", + ZSTD_c_targetLength => "target_length", + ZSTD_c_strategy => "strategy", + ZSTD_c_enableLongDistanceMatching => "enable_long_distance_matching", + ZSTD_c_ldmHashLog => "ldm_hash_log", + ZSTD_c_ldmMinMatch => "ldm_min_match", + ZSTD_c_ldmBucketSizeLog => "ldm_bucket_size_log", + ZSTD_c_ldmHashRateLog => "ldm_hash_rate_log", + ZSTD_c_contentSizeFlag => "content_size_flag", + ZSTD_c_checksumFlag => "checksum_flag", + ZSTD_c_dictIDFlag => "dict_id_flag", + ZSTD_c_nbWorkers => "nb_workers", + ZSTD_c_jobSize => "job_size", + ZSTD_c_overlapLog => "overlap_log", + _ => "unknown", + } + } else { + match param { + ZSTD_d_windowLogMax => "window_log_max", + _ => "unknown", + } + } + } + + /// Map a strategy integer (as exposed via the `Strategy` IntEnum) back + /// to the underlying `ZSTD_strategy` C enum value. Done via an explicit + /// match for the same reason as [`c_param_enum`]: an untrusted int + /// might not correspond to any real enum variant. + fn strategy_from_int(v: i32) -> Option { + use zstd_sys::ZSTD_strategy as S; + Some(match v { + ZSTD_fast => S::ZSTD_fast, + ZSTD_dfast => S::ZSTD_dfast, + ZSTD_greedy => S::ZSTD_greedy, + ZSTD_lazy => S::ZSTD_lazy, + ZSTD_lazy2 => S::ZSTD_lazy2, + ZSTD_btlazy2 => S::ZSTD_btlazy2, + ZSTD_btopt => S::ZSTD_btopt, + ZSTD_btultra => S::ZSTD_btultra, + ZSTD_btultra2 => S::ZSTD_btultra2, + _ => return None, + }) + } + + /// Decode the `zstd_dict=` constructor argument. Accepts either a + /// `ZstdDict` instance (treated as the default digested form) or a + /// `(ZstdDict, marker)` tuple produced by one of `ZstdDict.as_*`. + fn parse_zstd_dict_arg( + obj: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<(PyRef, i32)> { + if let Ok(d) = obj.clone().downcast::() { + return Ok((d, DICT_TYPE_DIGESTED)); + } + if let Ok(tuple) = obj.clone().downcast::() { + let items = tuple.as_slice(); + // Reject any tuple shape that is not (ZstdDict, int_marker) so the + // test suite's bad-args coverage (`(zd, 1.0)`, `(zd,)`, `(zd, 3)`, + // etc.) raises TypeError. Marker bounds match the three documented + // `as_*` properties. + if items.len() != 2 { + return Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")); + } + let d = items[0] + .clone() + .downcast::() + .map_err(|_| vm.new_type_error("zstd_dict argument should be a ZstdDict object"))?; + // The marker must be a plain int (not float/etc); overflow on + // `2**1000` propagates as OverflowError via `try_index`. + let marker_obj = &items[1]; + let marker = pyobj_to_i32(marker_obj, vm).map_err(|e| { + // Preserve OverflowError; everything else becomes TypeError so + // callers see a consistent "should be a ZstdDict" message. + if e.fast_isinstance(vm.ctx.exceptions.overflow_error) { + e + } else { + vm.new_type_error("zstd_dict argument should be a ZstdDict object") + } + })?; + if !(DICT_TYPE_DIGESTED..=DICT_TYPE_PREFIX).contains(&marker) { + return Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")); + } + return Ok((d, marker)); + } + Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")) + } + + // ========================================================================= + // ZstdDict + // ========================================================================= + + #[derive(FromArgs)] + pub(super) struct ZstdDictArgs { + #[pyarg(positional)] + dict_content: ArgBytesLike, + #[pyarg(named, default = false)] + is_raw: bool, + } + + /// ZstdDict(dict_content, /, *, is_raw=False) + /// + /// Represents a Zstandard dictionary. Hold the raw dictionary bytes and + /// the parsed `dict_id`. The same instance may be shared by many + /// compressors and decompressors. + /// + /// *dict_content* is a bytes-like object containing the dictionary + /// content. If *is_raw* is False (the default), the content is expected + /// to start with the Zstandard dictionary magic number and have a valid + /// header. If *is_raw* is True, arbitrary bytes are accepted and the + /// dictionary is treated as raw back-reference content. + #[pyattr] + #[pyclass(name = "ZstdDict")] + #[derive(PyPayload)] + struct ZstdDict { + dict_content: PyBytesRef, + dict_id: u32, + } + + impl std::fmt::Debug for ZstdDict { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "", + self.dict_id, + self.dict_content.len() + ) + } + } + + impl Constructor for ZstdDict { + type Args = ZstdDictArgs; + + fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { + let dict_content = args.dict_content.with_ref(|b| b.to_vec()); + // libzstd's `get_dict_id_from_dict` returns `None` either when + // the content is too small to contain a valid header or when it + // does not carry the dictionary magic. Both are runtime errors + // when `is_raw=False`, matching CPython's behavior of raising + // `ValueError` on a non-conformant dictionary. + let parsed_id = zstd_safe::get_dict_id_from_dict(&dict_content) + .map(|n| n.get()) + .unwrap_or(0); + if !args.is_raw && parsed_id == 0 { + return Err(vm.new_value_error( + "ZSTD_DICT_MAGIC_NUMBER not found, dict_content cannot be a 'raw content' \ + dictionary. To create a raw content dictionary, pass is_raw=True.", + )); + } + // Raw dictionaries still get a non-zero `dict_id` whenever their + // contents happen to look like a valid dict (this is the + // documented behavior tested in `test_is_raw`). + Ok(Self { + dict_content: vm.ctx.new_bytes(dict_content), + dict_id: parsed_id, + }) + } + } + + impl Representable for ZstdDict { + #[inline] + fn repr_str(zelf: &Py, _vm: &VirtualMachine) -> PyResult { + Ok(format!( + "", + zelf.dict_id, + zelf.dict_content.len() + )) + } + } + + impl rustpython_vm::types::AsMapping for ZstdDict { + fn as_mapping() -> &'static rustpython_vm::protocol::PyMappingMethods { + static AS_MAPPING: rustpython_vm::protocol::PyMappingMethods = + rustpython_vm::protocol::PyMappingMethods { + length: Some(|mapping, _vm| { + Ok(ZstdDict::mapping_downcast(mapping).dict_content.len()) + }), + subscript: None, + ass_subscript: None, + }; + &AS_MAPPING + } + } + + #[pyclass(with(Constructor, Representable, AsMapping))] + impl ZstdDict { + /// The content of a Zstandard dictionary, as a bytes object. + #[pygetset] + fn dict_content(&self) -> PyBytesRef { + self.dict_content.clone() + } + + /// The Zstandard dictionary ID. + /// + /// A non-zero value identifies a standard-format dictionary. A value + /// of zero either means the dictionary is in raw-content form or + /// that the dictionary ID was deliberately omitted from the header. + /// Dictionary IDs fall in the range 0 to 2**32 - 1. + #[pygetset] + fn dict_id(&self) -> u32 { + self.dict_id + } + + /// Load this ZstdDict as a digested dictionary. + /// + /// Returns a `(self, DICT_TYPE_DIGESTED)` tuple that can be passed + /// as the `zstd_dict=` argument of `ZstdCompressor` or + /// `ZstdDecompressor`. Digested form lets libzstd cache the parsed + /// dictionary for reuse across compression contexts. + #[pygetset] + fn as_digested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { + vm.ctx + .new_tuple(vec![zelf.into(), vm.ctx.new_int(DICT_TYPE_DIGESTED).into()]) + } + + /// Load this ZstdDict as an undigested dictionary. + /// + /// Returns a `(self, DICT_TYPE_UNDIGESTED)` tuple. In undigested + /// form the raw dictionary bytes are copied into each + /// (de)compressor, which is slower than digested but preserves the + /// original dictionary parameters across loads. + #[pygetset] + fn as_undigested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { + vm.ctx.new_tuple(vec![ + zelf.into(), + vm.ctx.new_int(DICT_TYPE_UNDIGESTED).into(), + ]) + } + + /// Load this ZstdDict as a prefix dictionary. + /// + /// Returns a `(self, DICT_TYPE_PREFIX)` tuple. Prefix loading treats + /// the dictionary bytes as a one-shot back-reference window for the + /// next frame only; subsequent frames will not see the prefix. + #[pygetset] + fn as_prefix(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { + vm.ctx + .new_tuple(vec![zelf.into(), vm.ctx.new_int(DICT_TYPE_PREFIX).into()]) + } + } + + // ========================================================================= + // ZstdCompressor + // ========================================================================= + + // The three flush modes for `ZstdCompressor.compress()`, mirrored as + // class attributes via `extend_class` below. Values are positional and + // chosen to match what CPython exposes. + const COMP_MODE_CONTINUE: i32 = 0; + const COMP_MODE_FLUSH_BLOCK: i32 = 1; + const COMP_MODE_FLUSH_FRAME: i32 = 2; + + /// Internal state of a `ZstdCompressor`. Holds the libzstd context, the + /// last mode used (for `last_mode` and `set_pledged_input_size` + /// validation), and the dictionary handles that the context may reference + /// internally. Field order matters here: Rust drops in declaration order, + /// so `cctx` is freed first; the held `CDict` (if any) and the source + /// `PyRef` go away afterwards, which is the safe order for + /// teardown. + struct CompressorState { + cctx: CCtx<'static>, + /// Cached digested dictionary. The CCtx references this via + /// `ref_cdict`, so it must outlive the CCtx (handled by Rust's + /// field drop order: `cctx` drops first). + _cdict: Option>, + /// Keeps the ZstdDict's bytes alive for `ref_prefix` mode. + _dict: Option>, + last_mode: i32, + } + + /// ZstdCompressor(level=None, options=None, zstd_dict=None) + /// + /// Create a compressor object for incremental compression of data into + /// the Zstandard format. + /// + /// *level* is an int specifying the compression level. Valid range is + /// roughly [-(1<<17), 22]; 0 means "use the default level". Cannot be + /// combined with *options*. + /// + /// *options* is a dict mapping `CompressionParameter` enum members to + /// integer values for fine-grained control. Cannot be combined with + /// *level*. + /// + /// *zstd_dict* is an optional `ZstdDict` (or one of its `as_*` + /// variants) for dictionary-assisted compression. + /// + /// Use `compress()` to feed data into the compressor and `flush()` to + /// finalize the current frame. + #[pyattr] + #[pyclass(name = "ZstdCompressor")] + #[derive(PyPayload)] + struct ZstdCompressor { + state: PyMutex, + } + + impl std::fmt::Debug for ZstdCompressor { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "_zstd.ZstdCompressor") + } + } + + #[derive(FromArgs)] + pub(super) struct ZstdCompressorArgs { + #[pyarg(any, optional)] + level: OptionalArg, + #[pyarg(any, optional)] + options: OptionalArg, + #[pyarg(any, optional)] + zstd_dict: OptionalArg, + } + + /// Translate the public `mode` int to the libzstd `ZSTD_EndDirective` + /// the streaming API takes. + fn end_directive_from_mode( + mode: i32, + vm: &VirtualMachine, + ) -> PyResult { + match mode { + COMP_MODE_CONTINUE => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_continue), + COMP_MODE_FLUSH_BLOCK => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_flush), + COMP_MODE_FLUSH_FRAME => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_end), + _ => Err(vm.new_value_error(format!( + "mode argument wrong value, it should be one of \ + ZstdCompressor.CONTINUE ({COMP_MODE_CONTINUE}), \ + ZstdCompressor.FLUSH_BLOCK ({COMP_MODE_FLUSH_BLOCK}), or \ + ZstdCompressor.FLUSH_FRAME ({COMP_MODE_FLUSH_FRAME})" + ))), + } + } + + impl Constructor for ZstdCompressor { + type Args = ZstdCompressorArgs; + + fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { + let level_opt = arg_or_none(args.level, vm); + let options_opt = arg_or_none(args.options, vm); + let dict_opt = arg_or_none(args.zstd_dict, vm); + + if level_opt.is_some() && options_opt.is_some() { + return Err(vm.new_runtime_error("Only one of level or options should be used")); + } + + let mut cctx = CCtx::<'static>::create(); + + if let Some(level_obj) = level_opt { + let level = parse_compression_level(&level_obj, vm)?; + cctx.set_parameter(CParameter::CompressionLevel(level)) + .map_err(|_| param_value_error_for(ZSTD_c_compressionLevel, level, true, vm))?; + } + + if let Some(options_obj) = options_opt { + apply_options(&mut cctx, options_obj, true, vm)?; + } + + let (held_cdict, held_dict) = load_compressor_dict(&mut cctx, dict_opt, vm)?; + + Ok(Self { + state: PyMutex::new(CompressorState { + cctx, + _cdict: held_cdict, + _dict: held_dict, + last_mode: COMP_MODE_FLUSH_FRAME, + }), + }) + } + } + + /// Parse and validate a compression `level` argument. libzstd silently + /// clamps out-of-range levels rather than erroring, but CPython surfaces + /// them as `ValueError`, and bigints become `ValueError` (not + /// `OverflowError`) for the same reason. This helper centralizes both + /// conversions so the constructor stays linear. + fn parse_compression_level(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + let (lo, hi) = level_bounds(); + let level = pyobj_to_i32(obj, vm).map_err(|e| { + if e.fast_isinstance(vm.ctx.exceptions.overflow_error) { + vm.new_value_error(format!( + "illegal compression level; the valid range is [{lo}, {hi}]" + )) + } else { + e + } + })?; + if level < lo || level > hi { + return Err(vm.new_value_error(format!( + "illegal compression level {level}; the valid range is [{lo}, {hi}]" + ))); + } + Ok(level) + } + + /// Drain an `options=` dict onto either a `CCtx` or a `DCtx`. Validates + /// each key/value pair (rejects wrong enum kind, rejects floats, rejects + /// out-of-range values) so the constructor's flow stays a single line. + fn apply_options( + ctx: &mut dyn ParamSetter, + options_obj: PyObjectRef, + is_compress: bool, + vm: &VirtualMachine, + ) -> PyResult<()> { + let dict = options_obj + .downcast::() + .map_err(|_| vm.new_type_error("options must be a dict"))?; + for (k, v) in dict { + // Reject the wrong enum kind ("CompressionParameter" vs + // "DecompressionParameter") before doing any numeric coercion, + // so the error carries the documented type name rather than a + // generic out-of-range message. + check_wrong_param_kind(&k, is_compress, vm)?; + let key_int = pyobj_to_i32(&k, vm)?; + let val_int = pyobj_to_i32(&v, vm)?; + // libzstd silently clamps out-of-range values for some + // parameters (notably compression_level) rather than rejecting + // them, so validate against the documented bounds upfront. + if let Some((lo, hi)) = lookup_param_bounds(key_int, is_compress) + && (val_int < lo || val_int > hi) + { + return Err(param_value_error_for(key_int, val_int, is_compress, vm)); + } + ctx.apply(key_int, val_int, vm)?; + } + Ok(()) + } + + /// Trait wrapper over `CCtx::set_parameter` and `DCtx::set_parameter` so + /// `apply_options` can drive either context without duplicated code. + /// `set_parameter` translates the (id, value) pair to the appropriate + /// `CParameter` / `DParameter` enum variant, then forwards to libzstd. + trait ParamSetter { + fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()>; + } + + impl ParamSetter for CCtx<'static> { + fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()> { + let p = cparameter_from_int(param, value, vm)?; + self.set_parameter(p) + .map_err(|_| param_value_error_for(param, value, true, vm))?; + Ok(()) + } + } + + impl ParamSetter for DCtx<'static> { + fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()> { + let p = dparameter_from_int(param, value, vm)?; + self.set_parameter(p) + .map_err(|_| param_value_error_for(param, value, false, vm))?; + Ok(()) + } + } + + /// Trait that captures the only differences between how the compressor + /// and decompressor consume a dictionary: the name of the type that + /// appears in error messages, the eager-validation constructor for the + /// digested variant, and the three ways of attaching it to the context. + trait DictLoader<'a> { + type Digested; + const KIND_NAME: &'static str; + fn try_create_digested(bytes: &[u8]) -> Option; + fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult; + fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult; + fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult; + } + + impl DictLoader<'static> for CCtx<'static> { + type Digested = zstd_safe::CDict<'static>; + const KIND_NAME: &'static str = "ZSTD_CDict"; + fn try_create_digested(bytes: &[u8]) -> Option { + zstd_safe::CDict::try_create(bytes, ZSTD_CLEVEL_DEFAULT) + } + fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult { + self.ref_cdict(dict) + } + fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult { + self.load_dictionary(bytes) + } + fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult { + self.ref_prefix(bytes) + } + } + + impl DictLoader<'static> for DCtx<'static> { + type Digested = zstd_safe::DDict<'static>; + const KIND_NAME: &'static str = "ZSTD_DDict"; + fn try_create_digested(bytes: &[u8]) -> Option { + zstd_safe::DDict::try_create(bytes) + } + fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult { + self.ref_ddict(dict) + } + fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult { + self.load_dictionary(bytes) + } + fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult { + self.ref_prefix(bytes) + } + } + + /// Common path for attaching a dictionary to either context type. Returns + /// the digested `CDict`/`DDict` (if the caller used digested mode) plus + /// the `PyRef` we hold to keep the bytes alive for `ref_prefix`. + fn load_dict>( + ctx: &mut L, + dict_obj: Option, + vm: &VirtualMachine, + ) -> PyResult<(Option, Option>)> { + let Some(dict_obj) = dict_obj else { + return Ok((None, None)); + }; + let (zdict, marker) = parse_zstd_dict_arg(dict_obj, vm)?; + let bad_dict_err = || -> PyBaseExceptionRef { + new_zstd_error( + format!( + "Failed to load the {} instance from corrupted Zstandard dictionary content.", + L::KIND_NAME + ), + vm, + ) + }; + let dict_bytes = zdict.dict_content.as_bytes(); + let mut digested = None; + match marker { + DICT_TYPE_PREFIX => { + // SAFETY: the returned `PyRef` keeps `zdict` alive + // (and therefore `dict_bytes`) for the lifetime of the + // context. libzstd `ref_prefix` stores a raw pointer into + // `dict_bytes`; as long as the (de)compressor outlives no + // longer than the held PyRef, the pointer stays valid. + let static_bytes: &'static [u8] = unsafe { + std::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) + }; + ctx.ref_prefix_static(static_bytes) + .map_err(|_| bad_dict_err())?; + } + DICT_TYPE_DIGESTED => { + // Build the digested dict eagerly so a corrupted dictionary + // surfaces as a `ZstdError` at construction time, not when + // the first compress/decompress call runs. + let d = L::try_create_digested(dict_bytes).ok_or_else(bad_dict_err)?; + ctx.ref_digested(&d).map_err(|_| bad_dict_err())?; + digested = Some(d); + } + _ => { + // Undigested: copy the bytes into the context. Validation + // happens lazily at the first stream call in this mode. + ctx.load_undigested(dict_bytes).map_err(|_| bad_dict_err())?; + } + } + Ok((digested, Some(zdict))) + } + + fn load_compressor_dict( + cctx: &mut CCtx<'static>, + dict_obj: Option, + vm: &VirtualMachine, + ) -> PyResult<(Option>, Option>)> { + load_dict::>(cctx, dict_obj, vm) + } + + fn load_decompressor_dict( + dctx: &mut DCtx<'static>, + dict_obj: Option, + vm: &VirtualMachine, + ) -> PyResult<(Option>, Option>)> { + load_dict::>(dctx, dict_obj, vm) + } + + + /// Drive `compress_stream2` until the input is fully consumed and, for + /// flush/end directives, the internal buffers report zero remaining bytes. + /// Grows the output `Vec` by `CStreamOutSize` chunks as needed. + fn do_compress( + state: &mut CompressorState, + data: &[u8], + end_op: zstd_sys::ZSTD_EndDirective, + vm: &VirtualMachine, + ) -> PyResult> { + let mut output = Vec::new(); + let mut input = InBuffer::around(data); + let chunk_size = CCtx::out_size().max(1); + let is_end = end_op != zstd_sys::ZSTD_EndDirective::ZSTD_e_continue; + + loop { + let prev_len = output.len(); + output.reserve(chunk_size); + let remaining = { + let mut out_buf = OutBuffer::around_pos(&mut output, prev_len); + state.cctx.compress_stream2(&mut out_buf, &mut input, end_op) + } + .map_err(|c| catch_zstd_error(c, vm))?; + let consumed_all = input.pos == input.src.len(); + // Stop when input is fully consumed and, for flush/end directives, + // libzstd reports that all internal buffers have been drained + // (remaining == 0). Otherwise loop; the next `reserve` will grow + // the output buffer if we hit the previous capacity. + if consumed_all && (!is_end || remaining == 0) { + break; + } + } + Ok(output) + } + + #[pyclass(with(Constructor))] + impl ZstdCompressor { + /// compress(data, mode=ZstdCompressor.CONTINUE) + /// + /// Provide data to the compressor and return a chunk of compressed + /// output bytes (which may be empty if not enough data has been + /// buffered yet). + /// + /// *mode* controls when to flush the output: + /// * `CONTINUE` (default) keeps buffering, producing the densest + /// compression. The compressor may return an empty byte string. + /// * `FLUSH_BLOCK` closes the current zstd block so all data + /// provided so far can be immediately decompressed. Frequent + /// use reduces the compression ratio. + /// * `FLUSH_FRAME` closes the current frame entirely. The next + /// `compress()` call starts a new frame. + #[pymethod] + fn compress(&self, args: CompressMethodArgs, vm: &VirtualMachine) -> PyResult> { + let mode = args.mode.unwrap_or(COMP_MODE_CONTINUE); + let end_op = end_directive_from_mode(mode, vm)?; + let data = args.data.with_ref(|b| b.to_vec()); + let mut state = self.state.lock(); + let out = do_compress(&mut state, &data, end_op, vm)?; + state.last_mode = mode; + Ok(out) + } + + /// flush(mode=ZstdCompressor.FLUSH_FRAME) + /// + /// Finish the compression process and return all the remaining + /// compressed bytes. *mode* must be one of `FLUSH_FRAME` (close the + /// current frame; new frames may be started afterward) or + /// `FLUSH_BLOCK` (close the current block but stay inside the + /// frame). `CONTINUE` is not a valid value for `flush()`. + #[pymethod] + fn flush(&self, args: FlushMethodArgs, vm: &VirtualMachine) -> PyResult> { + let mode = args.mode.unwrap_or(COMP_MODE_FLUSH_FRAME); + if mode != COMP_MODE_FLUSH_BLOCK && mode != COMP_MODE_FLUSH_FRAME { + return Err(vm.new_value_error(format!( + "mode argument wrong value, it should be \ + ZstdCompressor.FLUSH_FRAME ({COMP_MODE_FLUSH_FRAME}) or \ + ZstdCompressor.FLUSH_BLOCK ({COMP_MODE_FLUSH_BLOCK})" + ))); + } + let end_op = end_directive_from_mode(mode, vm)?; + let mut state = self.state.lock(); + let out = do_compress(&mut state, &[], end_op, vm)?; + state.last_mode = mode; + Ok(out) + } + + /// set_pledged_input_size(size, /) + /// + /// Set the uncompressed content size to be written into the frame + /// header of the next frame. Useful when the size is known up front: + /// libzstd writes it into the header (unless the content-size flag + /// is disabled) and verifies that the actual input matches at frame + /// end. Pass `None` to mean "unknown". + /// + /// May only be called when `last_mode == FLUSH_FRAME` (i.e. between + /// frames; not in the middle of one). + #[pymethod] + fn set_pledged_input_size(&self, size: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + let mut state = self.state.lock(); + if state.last_mode != COMP_MODE_FLUSH_FRAME { + return Err(vm.new_value_error( + "set_pledged_input_size() method must be called when last_mode == FLUSH_FRAME", + )); + } + // Python passes `None` to mean "unknown"; libzstd represents that + // internally as `ZSTD_CONTENTSIZE_UNKNOWN` (`u64::MAX`), and + // `zstd_safe` translates `None` accordingly. Reject any concrete + // int outside the valid range up front so callers see the + // documented `ValueError`, not a libzstd-level error. The valid + // upper bound is `ZSTD_CONTENTSIZE_ERROR - 1`, i.e. `u64::MAX - 2`. + let pledged: Option = if vm.is_none(&size) { + None + } else { + const UPPER_BOUND: u64 = u64::MAX - 2; + let int_ref = size.try_index(vm)?; + // Try fitting into u64; OverflowError covers negatives and + // values above `u64::MAX`. + let v: u64 = int_ref.try_to_primitive(vm).map_err(|_| { + vm.new_value_error(format!( + "size argument should be a positive int less than {UPPER_BOUND}" + )) + })?; + if v > UPPER_BOUND { + return Err(vm.new_value_error(format!( + "size argument should be a positive int less than {UPPER_BOUND}" + ))); + } + Some(v) + }; + state + .cctx + .set_pledged_src_size(pledged) + .map_err(|c| catch_zstd_error(c, vm))?; + Ok(()) + } + + /// The mode argument passed to the most recent `compress()` or + /// `flush()` call. + /// + /// The value is one of `CONTINUE`, `FLUSH_BLOCK`, or `FLUSH_FRAME`. + /// The initial value is `FLUSH_FRAME`, which signals that no frame + /// is currently in progress (it is safe to call + /// `set_pledged_input_size()` in this state). + #[pygetset] + fn last_mode(&self) -> i32 { + self.state.lock().last_mode + } + + /// Install class-level constants `CONTINUE`, `FLUSH_BLOCK`, and + /// `FLUSH_FRAME` so callers can reference them as + /// `ZstdCompressor.FLUSH_FRAME` (as the Python `ZstdFile` wrapper + /// does). + #[extend_class] + fn extend_class(ctx: &Context, class: &'static Py) { + class.set_attr( + ctx.intern_str("CONTINUE"), + ctx.new_int(COMP_MODE_CONTINUE).into(), + ); + class.set_attr( + ctx.intern_str("FLUSH_BLOCK"), + ctx.new_int(COMP_MODE_FLUSH_BLOCK).into(), + ); + class.set_attr( + ctx.intern_str("FLUSH_FRAME"), + ctx.new_int(COMP_MODE_FLUSH_FRAME).into(), + ); + } + } + + #[derive(FromArgs)] + pub(super) struct CompressMethodArgs { + #[pyarg(positional)] + data: ArgBytesLike, + #[pyarg(any, optional)] + mode: Option, + } + + #[derive(FromArgs)] + pub(super) struct FlushMethodArgs { + #[pyarg(any, optional)] + mode: Option, + } + + // ========================================================================= + // ZstdDecompressor + // ========================================================================= + + /// Internal state of a `ZstdDecompressor`. The CPython decompressor is + /// single-frame: once we hit end-of-frame, additional bytes go into + /// `unused_data` and further `decompress` calls raise `EOFError`. Field + /// drop order matters here for the same reason as in `CompressorState`: + /// the `dctx` is freed first and must give up its internal pointers + /// before any referenced `DDict`/`PyRef` is dropped. + struct DecompressorState { + dctx: DCtx<'static>, + /// Cached decompression dictionary referenced by the DCtx. + _ddict: Option>, + _dict: Option>, + eof: bool, + needs_input: bool, + /// Bytes that arrived after the end of the first frame. + unused_data: PyBytesRef, + /// Input bytes buffered because the previous `decompress` call ran + /// into its `max_length` cap before consuming them all. + input_buffer: Vec, + } + + /// ZstdDecompressor(zstd_dict=None, options=None) + /// + /// Create a decompressor object for incremental decompression of a + /// single Zstandard frame. Once the frame is fully decompressed the + /// decompressor is at EOF; trailing bytes are exposed via + /// `unused_data`. For multi-frame input use the module-level + /// `compression.zstd.decompress()` or `ZstdFile`. + /// + /// *zstd_dict* is an optional `ZstdDict` (or one of its `as_*` + /// variants) matching the dictionary used during compression. + /// + /// *options* is a dict mapping `DecompressionParameter` enum members + /// to integer values for fine-grained control. + #[pyattr] + #[pyclass(name = "ZstdDecompressor")] + #[derive(PyPayload)] + struct ZstdDecompressor { + state: PyMutex, + } + + impl std::fmt::Debug for ZstdDecompressor { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "_zstd.ZstdDecompressor") + } + } + + #[derive(FromArgs)] + pub(super) struct ZstdDecompressorArgs { + #[pyarg(any, optional)] + zstd_dict: OptionalArg, + #[pyarg(any, optional)] + options: OptionalArg, + } + + impl Constructor for ZstdDecompressor { + type Args = ZstdDecompressorArgs; + + fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { + let dict_opt = arg_or_none(args.zstd_dict, vm); + let options_opt = arg_or_none(args.options, vm); + + let mut dctx = DCtx::<'static>::create(); + + if let Some(options_obj) = options_opt { + apply_options(&mut dctx, options_obj, false, vm)?; + } + + let (held_ddict, held_dict) = load_decompressor_dict(&mut dctx, dict_opt, vm)?; + + Ok(Self { + state: PyMutex::new(DecompressorState { + dctx, + _ddict: held_ddict, + _dict: held_dict, + eof: false, + needs_input: true, + unused_data: vm.ctx.empty_bytes.clone(), + input_buffer: Vec::new(), + }), + }) + } + } + + #[derive(FromArgs)] + pub(super) struct DecompressMethodArgs { + #[pyarg(positional)] + data: ArgBytesLike, + #[pyarg(any, default = -1)] + max_length: isize, + } + + /// Drive `decompress_stream` until either the frame ends, the input is + /// exhausted with no more output coming, or `max_length` bytes have been + /// produced. Sets the various `state` flags to reflect the new situation. + /// + /// Loop control: we always keep going while either the input still has + /// bytes to feed OR the previous call filled the output buffer (which + /// indicates libzstd had more to emit but ran out of room). We only stop + /// short of a frame boundary when the input is exhausted AND libzstd had + /// room left in the output buffer, which means it is genuinely waiting + /// for more compressed bytes. + fn do_decompress( + state: &mut DecompressorState, + new_data: &[u8], + max_length: Option, + vm: &VirtualMachine, + ) -> PyResult> { + // Combine any buffered leftover input with the new data so the + // decompressor sees one contiguous stream. `Cow` avoids the + // allocation when there is no leftover. + let work_data: std::borrow::Cow<'_, [u8]> = if state.input_buffer.is_empty() { + std::borrow::Cow::Borrowed(new_data) + } else { + let mut combined = Vec::with_capacity(state.input_buffer.len() + new_data.len()); + combined.extend_from_slice(&state.input_buffer); + combined.extend_from_slice(new_data); + std::borrow::Cow::Owned(combined) + }; + + let mut input = InBuffer::around(&*work_data); + let mut output: Vec = Vec::new(); + let chunk_size = DCtx::out_size().max(1); + // Reusable scratch buffer for each decompress_stream call. We need an + // exact-size output buffer because `Vec::reserve` may over-allocate; + // OutBuffer reports the full Vec capacity to libzstd, which would + // then happily write past `max_length`. + let mut scratch: Vec = vec![0u8; chunk_size]; + let mut hit_max = false; + let mut iteration = 0usize; + + loop { + iteration += 1; + // Honor `max_length`: stop growing the output buffer once we have + // produced enough. Special-case the first iteration when the cap + // is zero so a zero-output frame (skippable frame, empty content + // frame) can still complete; we hand libzstd a 1-byte scratch + // and discard the byte if it ends up writing one. + let grow = match max_length { + Some(maxl) if output.len() >= maxl && iteration > 1 => { + hit_max = true; + break; + } + Some(maxl) if output.len() >= maxl => 1, + Some(maxl) => (maxl - output.len()).min(chunk_size), + None => chunk_size, + }; + let result; + let written; + { + let slot = &mut scratch[..grow]; + let mut out_buf = OutBuffer::around(slot as &mut [u8]); + result = state.dctx.decompress_stream(&mut out_buf, &mut input); + written = out_buf.pos(); + } + output.extend_from_slice(&scratch[..written]); + match result { + Ok(0) => { + // Frame fully decompressed; the decompressor is at EOF. + state.eof = true; + break; + } + Ok(_) => { + let output_was_full = written == grow; + let input_consumed = input.pos == input.src.len(); + + if let Some(maxl) = max_length + && output.len() >= maxl + && iteration > 1 + { + hit_max = true; + break; + } + + // Input is gone and libzstd had room to write but did not, + // which means the frame is incomplete and the caller has + // to supply more input. + if input_consumed && !output_was_full { + break; + } + } + Err(code) => return Err(catch_zstd_error(code, vm)), + } + } + + // If `max_length == 0` opened a courtesy iteration that produced more + // bytes than the caller asked for, truncate. Should not happen with + // the scratch slicing above, but keep the safety net. + if let Some(maxl) = max_length + && output.len() > maxl + { + output.truncate(maxl); + hit_max = true; + } + + let consumed = input.pos; + let remaining = &work_data[consumed..]; + + if state.eof { + if !remaining.is_empty() { + state.unused_data = vm.ctx.new_bytes(remaining.to_vec()); + } + state.input_buffer.clear(); + state.needs_input = false; + } else if hit_max { + // Output cap reached with input still pending. Buffer the rest + // and report `needs_input == false` so the caller knows to call + // `decompress(b'', max_length=...)` to drain it. + state.input_buffer = remaining.to_vec(); + state.needs_input = false; + } else if max_length == Some(0) { + // Caller explicitly asked for zero output bytes. Keep whatever + // input is left around for the next call and signal that they + // do not need to feed more right now. CPython's decompressor + // treats `max_length=0` as "stop here without losing state". + state.input_buffer = remaining.to_vec(); + state.needs_input = false; + } else { + // All input consumed but the frame is not complete; the caller + // should provide more data on the next call. + state.input_buffer.clear(); + state.needs_input = true; + } + + Ok(output) + } + + #[pyclass(with(Constructor))] + impl ZstdDecompressor { + /// decompress(data, max_length=-1) + /// + /// Decompress *data* and return at most *max_length* bytes of + /// uncompressed output. A *max_length* of -1 (the default) means + /// "no limit". After the current frame is fully decompressed, `eof` + /// becomes `True` and any further `decompress()` calls raise + /// `EOFError`; bytes that arrived after the frame go into + /// `unused_data` for the caller to forward. + /// + /// If the output cap is reached with input still pending, the + /// remaining input is buffered internally and `needs_input` is set + /// to `False`, signalling that the next call can be made with an + /// empty *data* argument to keep draining. + #[pymethod] + fn decompress(&self, args: DecompressMethodArgs, vm: &VirtualMachine) -> PyResult> { + let data_vec = args.data.with_ref(|b| b.to_vec()); + let max_length = if args.max_length < 0 { + None + } else { + Some(args.max_length as usize) + }; + let mut state = self.state.lock(); + if state.eof { + return Err(vm.new_exception_msg( + vm.ctx.exceptions.eof_error.to_owned(), + "Already at the end of a Zstandard frame.".to_owned().into(), + )); + } + do_decompress(&mut state, &data_vec, max_length, vm) + } + + /// True if the end-of-frame marker has been reached. + #[pygetset] + fn eof(&self) -> bool { + self.state.lock().eof + } + + /// True when the decompressor expects more compressed input. + /// + /// Becomes False after a `decompress()` call that hits its + /// `max_length` cap with input still buffered internally; in that + /// state the caller should call `decompress(b'', max_length=...)` + /// to drain the buffered output before providing more input. + #[pygetset] + fn needs_input(&self) -> bool { + self.state.lock().needs_input + } + + /// Bytes left over after the end-of-frame marker. + /// + /// Empty until the decompressor reaches EOF. Read-only. + #[pygetset] + fn unused_data(&self) -> PyBytesRef { + self.state.lock().unused_data.clone() + } + } + + // ========================================================================= + // Module-level functions + // ========================================================================= + + /// get_frame_size(frame_buffer, /) + /// + /// Return the size in bytes of the first complete Zstandard frame in + /// *frame_buffer* (which must contain the entire frame, header through + /// trailer). Raises `ZstdError` if *frame_buffer* is shorter than the + /// complete frame. + #[pyfunction] + fn get_frame_size(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { + let buf = frame_buffer.with_ref(|b| b.to_vec()); + zstd_safe::find_frame_compressed_size(&buf).map_err(|_| { + new_zstd_error( + "Error when finding the compressed size of a Zstandard frame. \ + Ensure the frame_buffer argument starts from the beginning of a frame, \ + and its length not less than this complete frame.", + vm, + ) + }) + } + + /// get_frame_info(frame_buffer, /) + /// + /// Return `(decompressed_size, dictionary_id)` for the first Zstandard + /// frame in *frame_buffer*. `decompressed_size` is `None` when the + /// frame header does not record it (small or streaming frames). + /// `dictionary_id` is 0 if no dictionary was used. *frame_buffer* must + /// contain at least the 6-18 byte frame header. The pure-Python + /// wrapper boxes this tuple into a `FrameInfo` named tuple. + #[pyfunction] + fn get_frame_info( + frame_buffer: ArgBytesLike, + vm: &VirtualMachine, + ) -> PyResult<(PyObjectRef, u32)> { + let buf = frame_buffer.with_ref(|b| b.to_vec()); + let content_size = zstd_safe::get_frame_content_size(&buf).map_err(|_| { + new_zstd_error( + "Error when getting information from the header of a Zstandard frame. \ + Ensure the frame_buffer argument starts from the beginning of a frame, \ + and its length not less than the frame header (6~18 bytes).", + vm, + ) + })?; + let content_size_obj: PyObjectRef = match content_size { + Some(n) => vm.ctx.new_int(n).into(), + None => vm.ctx.none(), + }; + let dict_id = zstd_safe::get_dict_id_from_frame(&buf) + .map(|n| n.get()) + .unwrap_or(0); + Ok((content_size_obj, dict_id)) + } + + #[derive(FromArgs)] + pub(super) struct TrainDictArgs { + /// Concatenated sample bytes. Must be a `bytes` object, not + /// `bytearray` or another buffer type, to match CPython's strict + /// type-checking on this argument. + #[pyarg(positional)] + samples_bytes: PyBytesRef, + /// A tuple of integer sample sizes that partition `samples_bytes`. + /// Lists and other iterables are not accepted. + #[pyarg(positional)] + samples_sizes: PyTupleRef, + /// Maximum size of the returned dictionary, in bytes. Must be a + /// positive `int`. + #[pyarg(positional)] + dict_size: PyObjectRef, + } + + /// Collect the elements of `tuple` into a `Vec`, validating that + /// each element is a non-negative int that fits in `usize`. Used by both + /// `train_dict` and `finalize_dict` for the `samples_sizes` argument. + /// + /// Floats (and any object whose `__index__` slot is missing) raise + /// `TypeError`; values that do not fit `usize` raise `ValueError` so + /// the test suite's `(2**1000,)` / `(-1,)` coverage holds. + fn parse_sample_sizes(tuple: PyTupleRef, vm: &VirtualMachine) -> PyResult> { + let items = tuple.as_slice(); + let mut out = Vec::with_capacity(items.len()); + for item in items { + let idx = item.try_index(vm)?; + let v: usize = idx.try_to_primitive(vm).map_err(|_| { + vm.new_value_error("sample size out of range for size_t".to_owned()) + })?; + out.push(v); + } + Ok(out) + } + + /// Convert a Python `int` to a positive `isize`. Rejects floats (via + /// `try_index`) and non-positive values; bigints that don't fit `isize` + /// propagate as `OverflowError`. Used for the `dict_size` argument of + /// `train_dict` and `finalize_dict`, which must always be a strictly + /// positive int. + fn parse_positive_dict_size(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { + let idx = obj.try_index(vm)?; + // `try_to_primitive::` raises `OverflowError` on bigints out of + // range; pass that through verbatim so the test suite's + // `assertRaises(OverflowError)` coverage matches. + let v: isize = idx.try_to_primitive(vm)?; + if v <= 0 { + return Err(vm.new_value_error("dict_size must be positive")); + } + Ok(v as usize) + } + + /// train_dict(samples_bytes, samples_sizes, dict_size, /) + /// + /// Train a Zstandard dictionary from sample data and return the raw + /// dictionary bytes. The pure-Python `train_dict` wrapper then wraps + /// the result in a `ZstdDict`. + /// + /// `samples_bytes` is the concatenation of all sample contents. + /// `samples_sizes` is a tuple giving the length of each sample, in the + /// same order they appear in `samples_bytes`. `dict_size` is the + /// maximum size of the dictionary in bytes. + #[pyfunction] + fn train_dict(args: TrainDictArgs, vm: &VirtualMachine) -> PyResult> { + let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; + let samples_buffer = args.samples_bytes.as_bytes().to_vec(); + let sizes = parse_sample_sizes(args.samples_sizes, vm)?; + let total: usize = sizes.iter().sum(); + if total != samples_buffer.len() { + return Err( + vm.new_value_error("The samples size tuple doesn't match the concatenation's size") + ); + } + let mut dict_buffer: Vec = Vec::with_capacity(dict_size); + zstd_safe::train_from_buffer(&mut dict_buffer, &samples_buffer, &sizes) + .map_err(|c| catch_zstd_error(c, vm))?; + Ok(dict_buffer) + } + + #[derive(FromArgs)] + pub(super) struct FinalizeDictArgs { + /// Raw "starting" dictionary content to finalize. Must be `bytes` + /// (not `bytearray`) to match CPython. + #[pyarg(positional)] + custom_dict_bytes: PyBytesRef, + /// Concatenated sample bytes used to derive the dictionary's + /// statistics tables. Must be `bytes`. + #[pyarg(positional)] + samples_bytes: PyBytesRef, + /// Tuple of integer sample sizes partitioning `samples_bytes`. + #[pyarg(positional)] + samples_sizes: PyTupleRef, + /// Maximum size of the finalized dictionary, in bytes. Positive int. + #[pyarg(positional)] + dict_size: PyObjectRef, + /// Compression level the dictionary will be tuned for. Must be int. + #[pyarg(positional)] + compression_level: PyObjectRef, + } + + /// finalize_dict(custom_dict_bytes, samples_bytes, samples_sizes, dict_size, compression_level, /) + /// + /// Take a hand-crafted custom dictionary plus a set of samples and + /// produce a standard-format Zstandard dictionary by appending the + /// header and statistics tables. `zstd_safe` does not wrap this + /// function, so we drop down to raw `zstd_sys` FFI for it. + #[pyfunction] + fn finalize_dict(args: FinalizeDictArgs, vm: &VirtualMachine) -> PyResult> { + let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; + let compression_level = pyobj_to_i32(&args.compression_level, vm)?; + let custom_dict = args.custom_dict_bytes.as_bytes().to_vec(); + let samples_buffer = args.samples_bytes.as_bytes().to_vec(); + let sizes = parse_sample_sizes(args.samples_sizes, vm)?; + let total: usize = sizes.iter().sum(); + if total != samples_buffer.len() { + return Err( + vm.new_value_error("The samples size tuple doesn't match the concatenation's size") + ); + } + + let mut dict_buffer: Vec = vec![0u8; dict_size]; + let params = zstd_sys::ZDICT_params_t { + compressionLevel: compression_level, + notificationLevel: 0, + dictID: 0, + }; + + // SAFETY: All pointers point into Rust-owned, properly sized buffers + // that outlive the FFI call. ZDICT_finalizeDictionary just reads from + // the sample/dict buffers and writes into `dict_buffer`. + let written = unsafe { + zstd_sys::ZDICT_finalizeDictionary( + dict_buffer.as_mut_ptr() as *mut _, + dict_buffer.len(), + custom_dict.as_ptr() as *const _, + custom_dict.len(), + samples_buffer.as_ptr() as *const _, + sizes.as_ptr(), + sizes.len() as u32, + params, + ) + }; + // SAFETY: ZDICT_isError just inspects the integer return code. + if unsafe { zstd_sys::ZDICT_isError(written) } != 0 { + // SAFETY: ZDICT_getErrorName returns a static NUL-terminated + // C string from libzstd's internal error table. + let err_ptr = unsafe { zstd_sys::ZDICT_getErrorName(written) }; + let msg = if err_ptr.is_null() { + "zstd dictionary finalization failed".to_string() + } else { + unsafe { std::ffi::CStr::from_ptr(err_ptr) } + .to_string_lossy() + .into_owned() + }; + return Err(new_zstd_error(msg, vm)); + } + dict_buffer.truncate(written); + Ok(dict_buffer) + } + + #[derive(FromArgs)] + pub(super) struct ParamBoundsArgs { + #[pyarg(positional)] + parameter: i32, + #[pyarg(named)] + is_compress: bool, + } + + /// get_param_bounds(parameter, /, *, is_compress) + /// + /// Return the inclusive `(lower, upper)` integer bounds for a + /// compression parameter (when *is_compress* is True) or a + /// decompression parameter (when False). Used by + /// `CompressionParameter.bounds()` and `DecompressionParameter.bounds()` + /// in the pure-Python wrapper. + #[pyfunction] + fn get_param_bounds(args: ParamBoundsArgs, vm: &VirtualMachine) -> PyResult<(c_int, c_int)> { + let bounds = if args.is_compress { + // Validate by mapping the int to a CParameter variant first; if + // the int is not a known parameter id we surface a clear error + // instead of letting libzstd return a generic "parameter unsupported". + let _ = cparameter_from_int(args.parameter, 0, vm)?; + // SAFETY: we just confirmed `args.parameter` is a value of the + // `ZSTD_cParameter` enum. + let param_enum = unsafe { + std::mem::transmute::(args.parameter as u32) + }; + unsafe { zstd_sys::ZSTD_cParam_getBounds(param_enum) } + } else { + let _ = dparameter_from_int(args.parameter, 0, vm)?; + // SAFETY: same justification as the compress branch above. + let param_enum = unsafe { + std::mem::transmute::(args.parameter as u32) + }; + unsafe { zstd_sys::ZSTD_dParam_getBounds(param_enum) } + }; + // SAFETY: ZSTD_isError just inspects the integer error code. + if unsafe { zstd_sys::ZSTD_isError(bounds.error) } != 0 { + return Err(catch_zstd_error(bounds.error, vm)); + } + Ok((bounds.lowerBound, bounds.upperBound)) + } + + /// set_parameter_types(c_parameter_type, d_parameter_type, /) + /// + /// Register the `CompressionParameter` and `DecompressionParameter` + /// `IntEnum` types defined in the pure-Python wrapper. Once registered, + /// the (de)compressor constructors validate that callers do not mix the + /// two enum kinds (e.g. passing a `CompressionParameter` key to a + /// `ZstdDecompressor`), raising `TypeError` if they do. + /// + /// The Python wrapper at `Lib/compression/zstd/__init__.py:242` calls + /// this exactly once at import time. + #[pyfunction] + fn set_parameter_types( + c_parameter_type: PyObjectRef, + d_parameter_type: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<()> { + let c_type = c_parameter_type + .downcast::() + .map_err(|_| vm.new_type_error("c_parameter_type must be a type object"))?; + let d_type = d_parameter_type + .downcast::() + .map_err(|_| vm.new_type_error("d_parameter_type must be a type object"))?; + let mut guard = REGISTERED_PARAM_TYPES.lock(); + guard.compression = Some(c_type); + guard.decompression = Some(d_type); + Ok(()) + } +} From 78c37029b2b7b6bea09deed22ada6dfc7f950143 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sat, 23 May 2026 18:33:36 +0100 Subject: [PATCH 02/25] Fix iOS build by removing static PyMutex `PyMutex` is built on `RawCellMutex` on single-threaded targets (iOS, Android, wasm32) and is not `Sync`, so a `static PyMutex<...>` fails to compile. `PyTypeRef` is also not `Send` by default, which compounds the issue. Replace the static parameter-type registry with class-name comparison (`"CompressionParameter"` / `"DecompressionParameter"`). The two names originate from the `compression.zstd` module we ship, so identity vs. name comparison is equivalent in practice. `set_parameter_types` now just validates that its arguments are type objects. Verified locally: - `cargo build -p rustpython-stdlib` (default features): clean. - `cargo check -p rustpython-stdlib --no-default-features --features host_env` (simulates iOS by dropping the `threading` feature): clean. - `cargo run -- -m test test_zstd`: 119 tests pass, 1 expected failure. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/stdlib/src/zstd.rs | 81 ++++++++++++++++++--------------------- 1 file changed, 37 insertions(+), 44 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 6c9ad33396e..7f0ef4990f5 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -37,23 +37,15 @@ mod _zstd { use zstd_safe::zstd_sys; use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; - /// Stores the `CompressionParameter` and `DecompressionParameter` enum - /// types that the pure-Python `compression.zstd` wrapper passes through - /// `set_parameter_types`. They let us produce CPython-compatible - /// `TypeError`s when callers pass the wrong enum kind to a (de)compressor - /// (a `CompressionParameter` key to a decompressor, or vice versa). Stored - /// in a `PyMutex` so the module init can set them once and constructors - /// can read them later from any thread. - static REGISTERED_PARAM_TYPES: PyMutex = - PyMutex::new(RegisteredParamTypes { - compression: None, - decompression: None, - }); - - struct RegisteredParamTypes { - compression: Option, - decompression: Option, - } + /// Class names of the `CompressionParameter` and `DecompressionParameter` + /// `IntEnum`s defined by the pure-Python wrapper at + /// `Lib/compression/zstd/__init__.py`. We identify these by name rather + /// than identity because storing `PyTypeRef` in a Rust `static` requires + /// `Sync`, which `PyMutex` does not provide on single-threaded targets + /// (Android, iOS, wasm32). Name comparison is sufficient since both names + /// originate from a module we ship and control. + const C_PARAMETER_TYPE_NAME: &str = "CompressionParameter"; + const D_PARAMETER_TYPE_NAME: &str = "DecompressionParameter"; // ========================================================================= // Module-level constants @@ -207,29 +199,29 @@ mod _zstd { arg.into_option().filter(|o| !vm.is_none(o)) } - /// If `key` is a registered `CompressionParameter` or `DecompressionParameter` - /// enum member and its kind does not match the caller's context, return a - /// `TypeError`. Used by both (de)compressor option-dict loops so passing a - /// `CompressionParameter` to a `ZstdDecompressor` (or vice versa) is - /// rejected with the CPython-compatible "must not be a XParameter" - /// message. If `set_parameter_types` has not been called yet, this is a - /// no-op. + /// If `key`'s class is the wrong kind of parameter enum for the caller's + /// context (a `CompressionParameter` passed to a decompressor, or vice + /// versa), return a `TypeError`. Used by both (de)compressor option-dict + /// loops so the resulting message names the actual type rather than the + /// generic "out of range" wording. + /// + /// We identify the enums by class name (see [`C_PARAMETER_TYPE_NAME`]), + /// which avoids storing a `PyTypeRef` in a Rust `static` (impossible on + /// single-threaded builds whose `PyMutex` is not `Sync`). fn check_wrong_param_kind( key: &PyObjectRef, is_compress: bool, vm: &VirtualMachine, ) -> PyResult<()> { - let guard = REGISTERED_PARAM_TYPES.lock(); - let (forbidden_type, forbidden_name, kind) = if is_compress { - (&guard.decompression, "DecompressionParameter", "compression") + let class_name = key.class().name(); + let (forbidden, kind) = if is_compress { + (D_PARAMETER_TYPE_NAME, "compression") } else { - (&guard.compression, "CompressionParameter", "decompression") + (C_PARAMETER_TYPE_NAME, "decompression") }; - if let Some(ref t) = *forbidden_type - && key.fast_isinstance(t) - { + if &*class_name == forbidden { return Err(vm.new_type_error(format!( - "{kind} options dictionary key must not be a {forbidden_name}" + "{kind} options dictionary key must not be a {forbidden}" ))); } Ok(()) @@ -1684,29 +1676,30 @@ mod _zstd { /// set_parameter_types(c_parameter_type, d_parameter_type, /) /// - /// Register the `CompressionParameter` and `DecompressionParameter` - /// `IntEnum` types defined in the pure-Python wrapper. Once registered, - /// the (de)compressor constructors validate that callers do not mix the - /// two enum kinds (e.g. passing a `CompressionParameter` key to a - /// `ZstdDecompressor`), raising `TypeError` if they do. + /// Validate that the `CompressionParameter` and `DecompressionParameter` + /// types passed in by the pure-Python wrapper are actual type objects. + /// The (de)compressor constructors use the *names* of these classes + /// (see [`check_wrong_param_kind`]) rather than identity, so this + /// function does not need to retain the type objects beyond the call. /// - /// The Python wrapper at `Lib/compression/zstd/__init__.py:242` calls - /// this exactly once at import time. + /// CPython retains them and uses them for nicer error messages naming + /// the specific enum member; matching that behavior would require a + /// `Sync` cell-of-`PyTypeRef`, which is not available on single-threaded + /// targets. The pure-Python wrapper at + /// `Lib/compression/zstd/__init__.py:242` calls this exactly once at + /// import time. #[pyfunction] fn set_parameter_types( c_parameter_type: PyObjectRef, d_parameter_type: PyObjectRef, vm: &VirtualMachine, ) -> PyResult<()> { - let c_type = c_parameter_type + let _ = c_parameter_type .downcast::() .map_err(|_| vm.new_type_error("c_parameter_type must be a type object"))?; - let d_type = d_parameter_type + let _ = d_parameter_type .downcast::() .map_err(|_| vm.new_type_error("d_parameter_type must be a type object"))?; - let mut guard = REGISTERED_PARAM_TYPES.lock(); - guard.compression = Some(c_type); - guard.decompression = Some(d_type); Ok(()) } } From cb70b822e4ce1fa8e2cb3941ba62b9067f77d334 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sat, 23 May 2026 18:56:31 +0100 Subject: [PATCH 03/25] Fix clippy errors in zstd module RustPython's clippy lint set rejects `std::` imports for items that also live in `core::` or `alloc::`. Switch every such reference in `zstd.rs`: - `std::ffi::{c_int, CStr}` -> `core::ffi::{c_int, CStr}` - `std::fmt::*` -> `core::fmt::*` - `std::slice::from_raw_parts` -> `core::slice::from_raw_parts` - `std::borrow::Cow` -> `alloc::borrow::Cow` Also clear up other clippy findings under `-D warnings`: - Drop the redundant `obj.clone()` on the last use of `obj` in `parse_zstd_dict_arg`. - Collapse `.map(|n| n.get()).unwrap_or(0)` to `.map_or(0, |n| n.get())` in the two dict-id readers. - Replace `&*work_data` with `&work_data` (auto-deref). - Factor the `(Option, Option>)` return shape into a `DictLoadResult` type alias to satisfy `type_complexity`. - Replace the remaining `std::mem::transmute` / `ZSTD_dParameter` in `get_param_bounds` with the existing safe `c_param_enum` / `d_param_enum` helpers; surfaces a clear "invalid parameter" `ValueError` instead of relying on UB-adjacent transmutes for unknown ints. Verified locally: - `cargo clippy -p rustpython-stdlib --no-deps -- -Dwarnings`: clean (only the pre-existing `socket.rs::sock_wait` warning remains). - `cargo run -- -m test test_zstd`: 119 pass, 1 expected failure. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/stdlib/src/zstd.rs | 84 +++++++++++++++++++++------------------ 1 file changed, 45 insertions(+), 39 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 7f0ef4990f5..95793ccee5b 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -33,7 +33,7 @@ mod _zstd { use rustpython_vm::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, }; - use std::ffi::c_int; + use core::ffi::c_int; use zstd_safe::zstd_sys; use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; @@ -428,10 +428,13 @@ mod _zstd { obj: PyObjectRef, vm: &VirtualMachine, ) -> PyResult<(PyRef, i32)> { + // The first downcast clones `obj` because we fall through to the + // tuple branch if it fails. The second downcast (the tuple one) is + // the last use of `obj`, so we let it move directly. if let Ok(d) = obj.clone().downcast::() { return Ok((d, DICT_TYPE_DIGESTED)); } - if let Ok(tuple) = obj.clone().downcast::() { + if let Ok(tuple) = obj.downcast::() { let items = tuple.as_slice(); // Reject any tuple shape that is not (ZstdDict, int_marker) so the // test suite's bad-args coverage (`(zd, 1.0)`, `(zd,)`, `(zd, 3)`, @@ -495,8 +498,8 @@ mod _zstd { dict_id: u32, } - impl std::fmt::Debug for ZstdDict { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + impl core::fmt::Debug for ZstdDict { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { write!( f, "", @@ -516,9 +519,8 @@ mod _zstd { // does not carry the dictionary magic. Both are runtime errors // when `is_raw=False`, matching CPython's behavior of raising // `ValueError` on a non-conformant dictionary. - let parsed_id = zstd_safe::get_dict_id_from_dict(&dict_content) - .map(|n| n.get()) - .unwrap_or(0); + let parsed_id = + zstd_safe::get_dict_id_from_dict(&dict_content).map_or(0, |n| n.get()); if !args.is_raw && parsed_id == 0 { return Err(vm.new_value_error( "ZSTD_DICT_MAGIC_NUMBER not found, dict_content cannot be a 'raw content' \ @@ -671,8 +673,8 @@ mod _zstd { state: PyMutex, } - impl std::fmt::Debug for ZstdCompressor { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + impl core::fmt::Debug for ZstdCompressor { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { write!(f, "_zstd.ZstdCompressor") } } @@ -873,6 +875,11 @@ mod _zstd { } } + /// Return value of `load_dict`: the digested `CDict`/`DDict` (if any) + /// and the `PyRef` we hold to keep the dictionary bytes alive + /// while `ref_prefix` may point into them. + type DictLoadResult = PyResult<(Option, Option>)>; + /// Common path for attaching a dictionary to either context type. Returns /// the digested `CDict`/`DDict` (if the caller used digested mode) plus /// the `PyRef` we hold to keep the bytes alive for `ref_prefix`. @@ -880,7 +887,7 @@ mod _zstd { ctx: &mut L, dict_obj: Option, vm: &VirtualMachine, - ) -> PyResult<(Option, Option>)> { + ) -> DictLoadResult { let Some(dict_obj) = dict_obj else { return Ok((None, None)); }; @@ -904,7 +911,7 @@ mod _zstd { // `dict_bytes`; as long as the (de)compressor outlives no // longer than the held PyRef, the pointer stays valid. let static_bytes: &'static [u8] = unsafe { - std::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) + core::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) }; ctx.ref_prefix_static(static_bytes) .map_err(|_| bad_dict_err())?; @@ -930,7 +937,7 @@ mod _zstd { cctx: &mut CCtx<'static>, dict_obj: Option, vm: &VirtualMachine, - ) -> PyResult<(Option>, Option>)> { + ) -> DictLoadResult> { load_dict::>(cctx, dict_obj, vm) } @@ -938,7 +945,7 @@ mod _zstd { dctx: &mut DCtx<'static>, dict_obj: Option, vm: &VirtualMachine, - ) -> PyResult<(Option>, Option>)> { + ) -> DictLoadResult> { load_dict::>(dctx, dict_obj, vm) } @@ -1169,8 +1176,8 @@ mod _zstd { state: PyMutex, } - impl std::fmt::Debug for ZstdDecompressor { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + impl core::fmt::Debug for ZstdDecompressor { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { write!(f, "_zstd.ZstdDecompressor") } } @@ -1239,16 +1246,16 @@ mod _zstd { // Combine any buffered leftover input with the new data so the // decompressor sees one contiguous stream. `Cow` avoids the // allocation when there is no leftover. - let work_data: std::borrow::Cow<'_, [u8]> = if state.input_buffer.is_empty() { - std::borrow::Cow::Borrowed(new_data) + let work_data: alloc::borrow::Cow<'_, [u8]> = if state.input_buffer.is_empty() { + alloc::borrow::Cow::Borrowed(new_data) } else { let mut combined = Vec::with_capacity(state.input_buffer.len() + new_data.len()); combined.extend_from_slice(&state.input_buffer); combined.extend_from_slice(new_data); - std::borrow::Cow::Owned(combined) + alloc::borrow::Cow::Owned(combined) }; - let mut input = InBuffer::around(&*work_data); + let mut input = InBuffer::around(&work_data); let mut output: Vec = Vec::new(); let chunk_size = DCtx::out_size().max(1); // Reusable scratch buffer for each decompress_stream call. We need an @@ -1463,9 +1470,7 @@ mod _zstd { Some(n) => vm.ctx.new_int(n).into(), None => vm.ctx.none(), }; - let dict_id = zstd_safe::get_dict_id_from_frame(&buf) - .map(|n| n.get()) - .unwrap_or(0); + let dict_id = zstd_safe::get_dict_id_from_frame(&buf).map_or(0, |n| n.get()); Ok((content_size_obj, dict_id)) } @@ -1621,7 +1626,7 @@ mod _zstd { let msg = if err_ptr.is_null() { "zstd dictionary finalization failed".to_string() } else { - unsafe { std::ffi::CStr::from_ptr(err_ptr) } + unsafe { core::ffi::CStr::from_ptr(err_ptr) } .to_string_lossy() .into_owned() }; @@ -1648,24 +1653,25 @@ mod _zstd { /// in the pure-Python wrapper. #[pyfunction] fn get_param_bounds(args: ParamBoundsArgs, vm: &VirtualMachine) -> PyResult<(c_int, c_int)> { + let unknown = || -> PyBaseExceptionRef { + let kind = if args.is_compress { "compression" } else { "decompression" }; + vm.new_value_error(format!( + "invalid {kind} parameter 'unknown parameter (key {})'", + args.parameter + )) + }; + // Validate the id via the same safe enum-lookup helpers used in + // `lookup_param_bounds`, then call libzstd directly so we can + // distinguish a libzstd-reported error from our own "unknown". let bounds = if args.is_compress { - // Validate by mapping the int to a CParameter variant first; if - // the int is not a known parameter id we surface a clear error - // instead of letting libzstd return a generic "parameter unsupported". - let _ = cparameter_from_int(args.parameter, 0, vm)?; - // SAFETY: we just confirmed `args.parameter` is a value of the - // `ZSTD_cParameter` enum. - let param_enum = unsafe { - std::mem::transmute::(args.parameter as u32) - }; - unsafe { zstd_sys::ZSTD_cParam_getBounds(param_enum) } + let p = c_param_enum(args.parameter).ok_or_else(unknown)?; + // SAFETY: `c_param_enum` returned `Some`, so `p` is a real + // `ZSTD_cParameter` discriminant. + unsafe { zstd_sys::ZSTD_cParam_getBounds(p) } } else { - let _ = dparameter_from_int(args.parameter, 0, vm)?; - // SAFETY: same justification as the compress branch above. - let param_enum = unsafe { - std::mem::transmute::(args.parameter as u32) - }; - unsafe { zstd_sys::ZSTD_dParam_getBounds(param_enum) } + let p = d_param_enum(args.parameter).ok_or_else(unknown)?; + // SAFETY: same as above. + unsafe { zstd_sys::ZSTD_dParam_getBounds(p) } }; // SAFETY: ZSTD_isError just inspects the integer error code. if unsafe { zstd_sys::ZSTD_isError(bounds.error) } != 0 { From fb78befdce7ac94e275d179dbb3e0646382f775a Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sat, 23 May 2026 18:59:59 +0100 Subject: [PATCH 04/25] Retrigger CI The previous run's `Lint` job failed with a `401 Unauthorized` from reviewdog hitting the GitHub API (a transient CI/permissions issue, unrelated to this PR's diff). Pushing an empty commit to re-run CI. Co-Authored-By: Claude Opus 4.7 (1M context) From 97ba0b433e3ce68653c6050d9efc80fa395d3141 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sat, 23 May 2026 19:09:49 +0100 Subject: [PATCH 05/25] chore: format zstd module {"subject": "Format zstd module and extend spell-check ignores", "body": "- Apply rustfmt formatting to zstd.rs\n- Add additional spell-checker ignore entries for zstd identifiers"} --- crates/stdlib/src/zstd.rs | 40 +++++++++++++++++++++++++-------------- 1 file changed, 26 insertions(+), 14 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 95793ccee5b..dc5a3db73a7 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -3,6 +3,8 @@ // spell-checker:ignore btlazy dfast nbworkers windowlogmax windowlog overlap targetcblock // spell-checker:ignore srcsize zdict refprefix refcdict refddict pledgedsrcsize getframecontentsize // spell-checker:ignore Zstd Zstandard pylib RFC +// spell-checker:ignore CLEVEL zstdmodule cparameter dparameter maxl +// spell-checker:ignore cctx dctx CCTX DCTX ldm cdict ddict windowlog hashlog chainlog searchlog CLEVEL //! The `_zstd` extension module. Backs the pure-Python `compression.zstd` //! package by exposing the same classes, functions and constants that @@ -24,6 +26,7 @@ pub(crate) use _zstd::module_def; #[allow(non_upper_case_globals)] #[pymodule] mod _zstd { + use core::ffi::c_int; use rustpython_common::lock::PyMutex; use rustpython_vm::builtins::{ PyBaseExceptionRef, PyBytesRef, PyDict, PyTupleRef, PyType, PyTypeRef, @@ -33,7 +36,6 @@ mod _zstd { use rustpython_vm::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, }; - use core::ffi::c_int; use zstd_safe::zstd_sys; use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; @@ -279,9 +281,11 @@ mod _zstd { ZSTD_c_searchLog => CParameter::SearchLog(value as u32), ZSTD_c_minMatch => CParameter::MinMatch(value as u32), ZSTD_c_targetLength => CParameter::TargetLength(value as u32), - ZSTD_c_strategy => CParameter::Strategy(strategy_from_int(value).ok_or_else( - || new_zstd_error(format!("invalid strategy value: {value}"), vm), - )?), + ZSTD_c_strategy => { + CParameter::Strategy(strategy_from_int(value).ok_or_else(|| { + new_zstd_error(format!("invalid strategy value: {value}"), vm) + })?) + } ZSTD_c_enableLongDistanceMatching => CParameter::EnableLongDistanceMatching(value != 0), ZSTD_c_ldmHashLog => CParameter::LdmHashLog(value as u32), ZSTD_c_ldmMinMatch => CParameter::LdmMinMatch(value as u32), @@ -323,7 +327,11 @@ mod _zstd { is_compress: bool, vm: &VirtualMachine, ) -> PyBaseExceptionRef { - let kind = if is_compress { "compression" } else { "decompression" }; + let kind = if is_compress { + "compression" + } else { + "decompression" + }; let name = parameter_name(param, is_compress); match lookup_param_bounds(param, is_compress) { Some((lo, hi)) => vm.new_value_error(format!( @@ -519,8 +527,7 @@ mod _zstd { // does not carry the dictionary magic. Both are runtime errors // when `is_raw=False`, matching CPython's behavior of raising // `ValueError` on a non-conformant dictionary. - let parsed_id = - zstd_safe::get_dict_id_from_dict(&dict_content).map_or(0, |n| n.get()); + let parsed_id = zstd_safe::get_dict_id_from_dict(&dict_content).map_or(0, |n| n.get()); if !args.is_raw && parsed_id == 0 { return Err(vm.new_value_error( "ZSTD_DICT_MAGIC_NUMBER not found, dict_content cannot be a 'raw content' \ @@ -910,9 +917,8 @@ mod _zstd { // context. libzstd `ref_prefix` stores a raw pointer into // `dict_bytes`; as long as the (de)compressor outlives no // longer than the held PyRef, the pointer stays valid. - let static_bytes: &'static [u8] = unsafe { - core::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) - }; + let static_bytes: &'static [u8] = + unsafe { core::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) }; ctx.ref_prefix_static(static_bytes) .map_err(|_| bad_dict_err())?; } @@ -927,7 +933,8 @@ mod _zstd { _ => { // Undigested: copy the bytes into the context. Validation // happens lazily at the first stream call in this mode. - ctx.load_undigested(dict_bytes).map_err(|_| bad_dict_err())?; + ctx.load_undigested(dict_bytes) + .map_err(|_| bad_dict_err())?; } } Ok((digested, Some(zdict))) @@ -949,7 +956,6 @@ mod _zstd { load_dict::>(dctx, dict_obj, vm) } - /// Drive `compress_stream2` until the input is fully consumed and, for /// flush/end directives, the internal buffers report zero remaining bytes. /// Grows the output `Vec` by `CStreamOutSize` chunks as needed. @@ -969,7 +975,9 @@ mod _zstd { output.reserve(chunk_size); let remaining = { let mut out_buf = OutBuffer::around_pos(&mut output, prev_len); - state.cctx.compress_stream2(&mut out_buf, &mut input, end_op) + state + .cctx + .compress_stream2(&mut out_buf, &mut input, end_op) } .map_err(|c| catch_zstd_error(c, vm))?; let consumed_all = input.pos == input.src.len(); @@ -1654,7 +1662,11 @@ mod _zstd { #[pyfunction] fn get_param_bounds(args: ParamBoundsArgs, vm: &VirtualMachine) -> PyResult<(c_int, c_int)> { let unknown = || -> PyBaseExceptionRef { - let kind = if args.is_compress { "compression" } else { "decompression" }; + let kind = if args.is_compress { + "compression" + } else { + "decompression" + }; vm.new_value_error(format!( "invalid {kind} parameter 'unknown parameter (key {})'", args.parameter From 31c4296d54425e34da9aeb5f3a2e1d924ae8fb18 Mon Sep 17 00:00:00 2001 From: James Clarke Date: Sun, 24 May 2026 20:09:51 +0100 Subject: [PATCH 06/25] Update crates/stdlib/src/zstd.rs Co-authored-by: Shahar Naveh <50263213+ShaharNaveh@users.noreply.github.com> --- crates/stdlib/src/zstd.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index dc5a3db73a7..339e9e3aa6c 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -1512,7 +1512,7 @@ mod _zstd { for item in items { let idx = item.try_index(vm)?; let v: usize = idx.try_to_primitive(vm).map_err(|_| { - vm.new_value_error("sample size out of range for size_t".to_owned()) + vm.new_value_error("sample size out of range for size_t") })?; out.push(v); } From ca664fda545c84ed1f2cdb94e8a008ef9abb4c54 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sun, 24 May 2026 20:51:35 +0100 Subject: [PATCH 07/25] chore: remove rustdoc comments from zstd module {"subject": "Remove rustdoc comments from zstd module", "body": "- Strip doc comments from internal items in _zstd module\n- Convert a few API-doc comments into regular implementation notes where appropriate"} --- crates/stdlib/src/zstd.rs | 203 +++----------------------------------- 1 file changed, 13 insertions(+), 190 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 339e9e3aa6c..d6a16efb6dd 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -53,8 +53,6 @@ mod _zstd { // Module-level constants // ========================================================================= - /// Default compression level used when none is requested. Matches - /// libzstd's `ZSTD_CLEVEL_DEFAULT`. #[pyattr] const ZSTD_CLEVEL_DEFAULT: i32 = zstd_sys::ZSTD_CLEVEL_DEFAULT as i32; @@ -127,21 +125,16 @@ mod _zstd { #[pyattr] const ZSTD_btultra2: i32 = zstd_sys::ZSTD_strategy::ZSTD_btultra2 as i32; - /// Runtime version string of the linked libzstd, like `"1.5.7"`. #[pyattr(once, name = "zstd_version")] fn zstd_version(_vm: &VirtualMachine) -> String { zstd_safe::version_string().to_string() } - /// Packed integer version: `MAJOR * 10000 + MINOR * 100 + RELEASE`. - /// The Python wrapper unpacks this into a `zstd_version_info` tuple. #[pyattr(once, name = "zstd_version_number")] fn zstd_version_number(_vm: &VirtualMachine) -> u32 { zstd_safe::version_number() } - /// Recommended output buffer size for decompression. The `ZstdFile.read1` - /// wrapper uses this as its default read size. #[pyattr(once, name = "ZSTD_DStreamOutSize")] fn zstd_dstream_out_size(_vm: &VirtualMachine) -> usize { DCtx::out_size() @@ -159,8 +152,6 @@ mod _zstd { // ZstdError exception // ========================================================================= - /// `_zstd.ZstdError`: the exception type raised for libzstd errors and for - /// invalid usage (calling `decompress` on a closed decompressor, etc.). #[pyattr(once, name = "ZstdError")] fn zstd_error(vm: &VirtualMachine) -> PyTypeRef { vm.ctx.new_exception_type( @@ -487,17 +478,6 @@ mod _zstd { is_raw: bool, } - /// ZstdDict(dict_content, /, *, is_raw=False) - /// - /// Represents a Zstandard dictionary. Hold the raw dictionary bytes and - /// the parsed `dict_id`. The same instance may be shared by many - /// compressors and decompressors. - /// - /// *dict_content* is a bytes-like object containing the dictionary - /// content. If *is_raw* is False (the default), the content is expected - /// to start with the Zstandard dictionary magic number and have a valid - /// header. If *is_raw* is True, arbitrary bytes are accepted and the - /// dictionary is treated as raw back-reference content. #[pyattr] #[pyclass(name = "ZstdDict")] #[derive(PyPayload)] @@ -571,41 +551,22 @@ mod _zstd { #[pyclass(with(Constructor, Representable, AsMapping))] impl ZstdDict { - /// The content of a Zstandard dictionary, as a bytes object. #[pygetset] fn dict_content(&self) -> PyBytesRef { self.dict_content.clone() } - /// The Zstandard dictionary ID. - /// - /// A non-zero value identifies a standard-format dictionary. A value - /// of zero either means the dictionary is in raw-content form or - /// that the dictionary ID was deliberately omitted from the header. - /// Dictionary IDs fall in the range 0 to 2**32 - 1. #[pygetset] fn dict_id(&self) -> u32 { self.dict_id } - /// Load this ZstdDict as a digested dictionary. - /// - /// Returns a `(self, DICT_TYPE_DIGESTED)` tuple that can be passed - /// as the `zstd_dict=` argument of `ZstdCompressor` or - /// `ZstdDecompressor`. Digested form lets libzstd cache the parsed - /// dictionary for reuse across compression contexts. #[pygetset] fn as_digested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { vm.ctx .new_tuple(vec![zelf.into(), vm.ctx.new_int(DICT_TYPE_DIGESTED).into()]) } - /// Load this ZstdDict as an undigested dictionary. - /// - /// Returns a `(self, DICT_TYPE_UNDIGESTED)` tuple. In undigested - /// form the raw dictionary bytes are copied into each - /// (de)compressor, which is slower than digested but preserves the - /// original dictionary parameters across loads. #[pygetset] fn as_undigested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { vm.ctx.new_tuple(vec![ @@ -614,11 +575,6 @@ mod _zstd { ]) } - /// Load this ZstdDict as a prefix dictionary. - /// - /// Returns a `(self, DICT_TYPE_PREFIX)` tuple. Prefix loading treats - /// the dictionary bytes as a one-shot back-reference window for the - /// next frame only; subsequent frames will not see the prefix. #[pygetset] fn as_prefix(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { vm.ctx @@ -655,24 +611,6 @@ mod _zstd { last_mode: i32, } - /// ZstdCompressor(level=None, options=None, zstd_dict=None) - /// - /// Create a compressor object for incremental compression of data into - /// the Zstandard format. - /// - /// *level* is an int specifying the compression level. Valid range is - /// roughly [-(1<<17), 22]; 0 means "use the default level". Cannot be - /// combined with *options*. - /// - /// *options* is a dict mapping `CompressionParameter` enum members to - /// integer values for fine-grained control. Cannot be combined with - /// *level*. - /// - /// *zstd_dict* is an optional `ZstdDict` (or one of its `as_*` - /// variants) for dictionary-assisted compression. - /// - /// Use `compress()` to feed data into the compressor and `flush()` to - /// finalize the current frame. #[pyattr] #[pyclass(name = "ZstdCompressor")] #[derive(PyPayload)] @@ -994,20 +932,6 @@ mod _zstd { #[pyclass(with(Constructor))] impl ZstdCompressor { - /// compress(data, mode=ZstdCompressor.CONTINUE) - /// - /// Provide data to the compressor and return a chunk of compressed - /// output bytes (which may be empty if not enough data has been - /// buffered yet). - /// - /// *mode* controls when to flush the output: - /// * `CONTINUE` (default) keeps buffering, producing the densest - /// compression. The compressor may return an empty byte string. - /// * `FLUSH_BLOCK` closes the current zstd block so all data - /// provided so far can be immediately decompressed. Frequent - /// use reduces the compression ratio. - /// * `FLUSH_FRAME` closes the current frame entirely. The next - /// `compress()` call starts a new frame. #[pymethod] fn compress(&self, args: CompressMethodArgs, vm: &VirtualMachine) -> PyResult> { let mode = args.mode.unwrap_or(COMP_MODE_CONTINUE); @@ -1019,13 +943,6 @@ mod _zstd { Ok(out) } - /// flush(mode=ZstdCompressor.FLUSH_FRAME) - /// - /// Finish the compression process and return all the remaining - /// compressed bytes. *mode* must be one of `FLUSH_FRAME` (close the - /// current frame; new frames may be started afterward) or - /// `FLUSH_BLOCK` (close the current block but stay inside the - /// frame). `CONTINUE` is not a valid value for `flush()`. #[pymethod] fn flush(&self, args: FlushMethodArgs, vm: &VirtualMachine) -> PyResult> { let mode = args.mode.unwrap_or(COMP_MODE_FLUSH_FRAME); @@ -1043,16 +960,6 @@ mod _zstd { Ok(out) } - /// set_pledged_input_size(size, /) - /// - /// Set the uncompressed content size to be written into the frame - /// header of the next frame. Useful when the size is known up front: - /// libzstd writes it into the header (unless the content-size flag - /// is disabled) and verifies that the actual input matches at frame - /// end. Pass `None` to mean "unknown". - /// - /// May only be called when `last_mode == FLUSH_FRAME` (i.e. between - /// frames; not in the middle of one). #[pymethod] fn set_pledged_input_size(&self, size: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { let mut state = self.state.lock(); @@ -1093,13 +1000,6 @@ mod _zstd { Ok(()) } - /// The mode argument passed to the most recent `compress()` or - /// `flush()` call. - /// - /// The value is one of `CONTINUE`, `FLUSH_BLOCK`, or `FLUSH_FRAME`. - /// The initial value is `FLUSH_FRAME`, which signals that no frame - /// is currently in progress (it is safe to call - /// `set_pledged_input_size()` in this state). #[pygetset] fn last_mode(&self) -> i32 { self.state.lock().last_mode @@ -1164,19 +1064,6 @@ mod _zstd { input_buffer: Vec, } - /// ZstdDecompressor(zstd_dict=None, options=None) - /// - /// Create a decompressor object for incremental decompression of a - /// single Zstandard frame. Once the frame is fully decompressed the - /// decompressor is at EOF; trailing bytes are exposed via - /// `unused_data`. For multi-frame input use the module-level - /// `compression.zstd.decompress()` or `ZstdFile`. - /// - /// *zstd_dict* is an optional `ZstdDict` (or one of its `as_*` - /// variants) matching the dictionary used during compression. - /// - /// *options* is a dict mapping `DecompressionParameter` enum members - /// to integer values for fine-grained control. #[pyattr] #[pyclass(name = "ZstdDecompressor")] #[derive(PyPayload)] @@ -1372,19 +1259,6 @@ mod _zstd { #[pyclass(with(Constructor))] impl ZstdDecompressor { - /// decompress(data, max_length=-1) - /// - /// Decompress *data* and return at most *max_length* bytes of - /// uncompressed output. A *max_length* of -1 (the default) means - /// "no limit". After the current frame is fully decompressed, `eof` - /// becomes `True` and any further `decompress()` calls raise - /// `EOFError`; bytes that arrived after the frame go into - /// `unused_data` for the caller to forward. - /// - /// If the output cap is reached with input still pending, the - /// remaining input is buffered internally and `needs_input` is set - /// to `False`, signalling that the next call can be made with an - /// empty *data* argument to keep draining. #[pymethod] fn decompress(&self, args: DecompressMethodArgs, vm: &VirtualMachine) -> PyResult> { let data_vec = args.data.with_ref(|b| b.to_vec()); @@ -1403,26 +1277,16 @@ mod _zstd { do_decompress(&mut state, &data_vec, max_length, vm) } - /// True if the end-of-frame marker has been reached. #[pygetset] fn eof(&self) -> bool { self.state.lock().eof } - /// True when the decompressor expects more compressed input. - /// - /// Becomes False after a `decompress()` call that hits its - /// `max_length` cap with input still buffered internally; in that - /// state the caller should call `decompress(b'', max_length=...)` - /// to drain the buffered output before providing more input. #[pygetset] fn needs_input(&self) -> bool { self.state.lock().needs_input } - /// Bytes left over after the end-of-frame marker. - /// - /// Empty until the decompressor reaches EOF. Read-only. #[pygetset] fn unused_data(&self) -> PyBytesRef { self.state.lock().unused_data.clone() @@ -1433,12 +1297,6 @@ mod _zstd { // Module-level functions // ========================================================================= - /// get_frame_size(frame_buffer, /) - /// - /// Return the size in bytes of the first complete Zstandard frame in - /// *frame_buffer* (which must contain the entire frame, header through - /// trailer). Raises `ZstdError` if *frame_buffer* is shorter than the - /// complete frame. #[pyfunction] fn get_frame_size(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { let buf = frame_buffer.with_ref(|b| b.to_vec()); @@ -1452,14 +1310,6 @@ mod _zstd { }) } - /// get_frame_info(frame_buffer, /) - /// - /// Return `(decompressed_size, dictionary_id)` for the first Zstandard - /// frame in *frame_buffer*. `decompressed_size` is `None` when the - /// frame header does not record it (small or streaming frames). - /// `dictionary_id` is 0 if no dictionary was used. *frame_buffer* must - /// contain at least the 6-18 byte frame header. The pure-Python - /// wrapper boxes this tuple into a `FrameInfo` named tuple. #[pyfunction] fn get_frame_info( frame_buffer: ArgBytesLike, @@ -1511,9 +1361,9 @@ mod _zstd { let mut out = Vec::with_capacity(items.len()); for item in items { let idx = item.try_index(vm)?; - let v: usize = idx.try_to_primitive(vm).map_err(|_| { - vm.new_value_error("sample size out of range for size_t") - })?; + let v: usize = idx + .try_to_primitive(vm) + .map_err(|_| vm.new_value_error("sample size out of range for size_t"))?; out.push(v); } Ok(out) @@ -1536,16 +1386,6 @@ mod _zstd { Ok(v as usize) } - /// train_dict(samples_bytes, samples_sizes, dict_size, /) - /// - /// Train a Zstandard dictionary from sample data and return the raw - /// dictionary bytes. The pure-Python `train_dict` wrapper then wraps - /// the result in a `ZstdDict`. - /// - /// `samples_bytes` is the concatenation of all sample contents. - /// `samples_sizes` is a tuple giving the length of each sample, in the - /// same order they appear in `samples_bytes`. `dict_size` is the - /// maximum size of the dictionary in bytes. #[pyfunction] fn train_dict(args: TrainDictArgs, vm: &VirtualMachine) -> PyResult> { let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; @@ -1584,12 +1424,8 @@ mod _zstd { compression_level: PyObjectRef, } - /// finalize_dict(custom_dict_bytes, samples_bytes, samples_sizes, dict_size, compression_level, /) - /// - /// Take a hand-crafted custom dictionary plus a set of samples and - /// produce a standard-format Zstandard dictionary by appending the - /// header and statistics tables. `zstd_safe` does not wrap this - /// function, so we drop down to raw `zstd_sys` FFI for it. + // `zstd_safe` does not wrap `ZDICT_finalizeDictionary`, so we drop down to + // raw `zstd_sys` FFI for it. #[pyfunction] fn finalize_dict(args: FinalizeDictArgs, vm: &VirtualMachine) -> PyResult> { let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; @@ -1652,13 +1488,6 @@ mod _zstd { is_compress: bool, } - /// get_param_bounds(parameter, /, *, is_compress) - /// - /// Return the inclusive `(lower, upper)` integer bounds for a - /// compression parameter (when *is_compress* is True) or a - /// decompression parameter (when False). Used by - /// `CompressionParameter.bounds()` and `DecompressionParameter.bounds()` - /// in the pure-Python wrapper. #[pyfunction] fn get_param_bounds(args: ParamBoundsArgs, vm: &VirtualMachine) -> PyResult<(c_int, c_int)> { let unknown = || -> PyBaseExceptionRef { @@ -1692,20 +1521,14 @@ mod _zstd { Ok((bounds.lowerBound, bounds.upperBound)) } - /// set_parameter_types(c_parameter_type, d_parameter_type, /) - /// - /// Validate that the `CompressionParameter` and `DecompressionParameter` - /// types passed in by the pure-Python wrapper are actual type objects. - /// The (de)compressor constructors use the *names* of these classes - /// (see [`check_wrong_param_kind`]) rather than identity, so this - /// function does not need to retain the type objects beyond the call. - /// - /// CPython retains them and uses them for nicer error messages naming - /// the specific enum member; matching that behavior would require a - /// `Sync` cell-of-`PyTypeRef`, which is not available on single-threaded - /// targets. The pure-Python wrapper at - /// `Lib/compression/zstd/__init__.py:242` calls this exactly once at - /// import time. + // The (de)compressor constructors identify the parameter enums by class + // *name* (see [`check_wrong_param_kind`]) rather than identity, so this + // function does not need to retain the type objects beyond the call. + // CPython retains them to produce nicer error messages naming the specific + // enum member; matching that would require a `Sync` cell-of-`PyTypeRef`, + // which is not available on single-threaded targets. The pure-Python + // wrapper at `Lib/compression/zstd/__init__.py:242` calls this exactly + // once at import time. #[pyfunction] fn set_parameter_types( c_parameter_type: PyObjectRef, From a9ce31b832ab8525e3a142d0fc194e73b6e30812 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Mon, 25 May 2026 09:28:38 +0100 Subject: [PATCH 08/25] commit message {"subject": "Mark zstd load_dict as unsafe and document invariant", "body": "- Add # Safety docs requiring PyRef to outlive the context\n- Wrap calls in load_compressor_dict and load_decompressor_dict with SAFETY comments"} --- crates/stdlib/src/zstd.rs | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index d6a16efb6dd..de220ecce6a 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -828,7 +828,13 @@ mod _zstd { /// Common path for attaching a dictionary to either context type. Returns /// the digested `CDict`/`DDict` (if the caller used digested mode) plus /// the `PyRef` we hold to keep the bytes alive for `ref_prefix`. - fn load_dict>( + /// # Safety + /// + /// The caller must ensure that the returned `PyRef` (if any) + /// outlives the context `ctx`. In `ref_prefix` mode, libzstd stores a + /// raw pointer into the dictionary bytes; the `PyRef` keeps those bytes + /// alive for the required lifetime. + unsafe fn load_dict>( ctx: &mut L, dict_obj: Option, vm: &VirtualMachine, @@ -883,7 +889,9 @@ mod _zstd { dict_obj: Option, vm: &VirtualMachine, ) -> DictLoadResult> { - load_dict::>(cctx, dict_obj, vm) + // SAFETY: The returned `PyRef` is stored alongside the + // `CCtx` in `CompressorState`, so it outlives the context. + unsafe { load_dict::>(cctx, dict_obj, vm) } } fn load_decompressor_dict( @@ -891,7 +899,9 @@ mod _zstd { dict_obj: Option, vm: &VirtualMachine, ) -> DictLoadResult> { - load_dict::>(dctx, dict_obj, vm) + // SAFETY: The returned `PyRef` is stored alongside the + // `DCtx` in `DecompressorState`, so it outlives the context. + unsafe { load_dict::>(dctx, dict_obj, vm) } } /// Drive `compress_stream2` until the input is fully consumed and, for From b77ba8b606a98a5b1cd3858a7ffa2d40e4b54d0e Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Mon, 25 May 2026 12:02:48 +0100 Subject: [PATCH 09/25] commit message { "subject": "refactor(zstd): release GIL during (de)compress and tighten dict-load safety", "body": "- Wrap compress/decompress loops in vm.allow_threads to release the GIL\n- Replace load_*_dict helpers with build_*_state that assemble the full state, making the load_dict safety invariants structural\n- Switch constructor args from OptionalArg to OptionalOption and use try_to_value instead of the pyobj_to_i32/arg_or_none helpers\n- Add check_sample_sizes_match with overflow-safe summing for train_dict/finalize_dict" } --- crates/stdlib/src/zstd.rs | 375 +++++++++++++++++++++----------------- 1 file changed, 208 insertions(+), 167 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index de220ecce6a..58f5f807996 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -31,7 +31,7 @@ mod _zstd { use rustpython_vm::builtins::{ PyBaseExceptionRef, PyBytesRef, PyDict, PyTupleRef, PyType, PyTypeRef, }; - use rustpython_vm::function::{ArgBytesLike, OptionalArg}; + use rustpython_vm::function::{ArgBytesLike, OptionalOption}; use rustpython_vm::types::{AsMapping, Constructor, Representable}; use rustpython_vm::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, @@ -177,21 +177,6 @@ mod _zstd { // Parameter helpers // ========================================================================= - /// Convert any int-like Python value to an `i32`, rejecting floats. CPython - /// uses `PyLong_AsInt` here which only accepts integers; `try_index` is - /// the equivalent of CPython's `__index__` slot lookup (rejects floats). - fn pyobj_to_i32(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { - obj.try_index(vm)?.try_to_primitive(vm) - } - - /// Normalize an `OptionalArg` constructor parameter to an - /// `Option`. Treats both "argument missing" and "argument - /// explicitly None" the same way (consistent with CPython, which uses - /// `None` defaults throughout the `_zstd` API). - fn arg_or_none(arg: OptionalArg, vm: &VirtualMachine) -> Option { - arg.into_option().filter(|o| !vm.is_none(o)) - } - /// If `key`'s class is the wrong kind of parameter enum for the caller's /// context (a `CompressionParameter` passed to a decompressor, or vice /// versa), return a `TypeError`. Used by both (de)compressor option-dict @@ -449,7 +434,7 @@ mod _zstd { // The marker must be a plain int (not float/etc); overflow on // `2**1000` propagates as OverflowError via `try_index`. let marker_obj = &items[1]; - let marker = pyobj_to_i32(marker_obj, vm).map_err(|e| { + let marker: i32 = marker_obj.try_to_value(vm).map_err(|e| { // Preserve OverflowError; everything else becomes TypeError so // callers see a consistent "should be a ZstdDict" message. if e.fast_isinstance(vm.ctx.exceptions.overflow_error) { @@ -627,11 +612,11 @@ mod _zstd { #[derive(FromArgs)] pub(super) struct ZstdCompressorArgs { #[pyarg(any, optional)] - level: OptionalArg, + level: OptionalOption, #[pyarg(any, optional)] - options: OptionalArg, + options: OptionalOption, #[pyarg(any, optional)] - zstd_dict: OptionalArg, + zstd_dict: OptionalOption, } /// Translate the public `mode` int to the libzstd `ZSTD_EndDirective` @@ -657,9 +642,9 @@ mod _zstd { type Args = ZstdCompressorArgs; fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { - let level_opt = arg_or_none(args.level, vm); - let options_opt = arg_or_none(args.options, vm); - let dict_opt = arg_or_none(args.zstd_dict, vm); + let level_opt = args.level.flatten(); + let options_opt = args.options.flatten(); + let dict_opt = args.zstd_dict.flatten(); if level_opt.is_some() && options_opt.is_some() { return Err(vm.new_runtime_error("Only one of level or options should be used")); @@ -677,15 +662,9 @@ mod _zstd { apply_options(&mut cctx, options_obj, true, vm)?; } - let (held_cdict, held_dict) = load_compressor_dict(&mut cctx, dict_opt, vm)?; - + let state = build_compressor_state(cctx, dict_opt, COMP_MODE_FLUSH_FRAME, vm)?; Ok(Self { - state: PyMutex::new(CompressorState { - cctx, - _cdict: held_cdict, - _dict: held_dict, - last_mode: COMP_MODE_FLUSH_FRAME, - }), + state: PyMutex::new(state), }) } } @@ -697,7 +676,7 @@ mod _zstd { /// conversions so the constructor stays linear. fn parse_compression_level(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { let (lo, hi) = level_bounds(); - let level = pyobj_to_i32(obj, vm).map_err(|e| { + let level: i32 = obj.try_to_value(vm).map_err(|e| { if e.fast_isinstance(vm.ctx.exceptions.overflow_error) { vm.new_value_error(format!( "illegal compression level; the valid range is [{lo}, {hi}]" @@ -732,8 +711,8 @@ mod _zstd { // so the error carries the documented type name rather than a // generic out-of-range message. check_wrong_param_kind(&k, is_compress, vm)?; - let key_int = pyobj_to_i32(&k, vm)?; - let val_int = pyobj_to_i32(&v, vm)?; + let key_int: i32 = k.try_to_value(vm)?; + let val_int: i32 = v.try_to_value(vm)?; // libzstd silently clamps out-of-range values for some // parameters (notably compression_level) rather than rejecting // them, so validate against the documented bounds upfront. @@ -827,13 +806,23 @@ mod _zstd { /// Common path for attaching a dictionary to either context type. Returns /// the digested `CDict`/`DDict` (if the caller used digested mode) plus - /// the `PyRef` we hold to keep the bytes alive for `ref_prefix`. + /// the `PyRef` whose bytes libzstd's `ref_prefix` may point into. + /// /// # Safety /// - /// The caller must ensure that the returned `PyRef` (if any) - /// outlives the context `ctx`. In `ref_prefix` mode, libzstd stores a - /// raw pointer into the dictionary bytes; the `PyRef` keeps those bytes - /// alive for the required lifetime. + /// libzstd stores the dictionary as a raw pointer that bypasses Rust's + /// lifetime tracking. The caller must keep both returned values alive at + /// least as long as `ctx`: + /// + /// - In `digested` mode, `ctx` holds a raw pointer to the returned + /// `L::Digested`; dropping it before `ctx` is use-after-free. + /// - In `prefix` mode, `ctx` holds a raw pointer into the bytes owned by + /// the returned `PyRef`; dropping the `PyRef` before `ctx` + /// is use-after-free. + /// + /// In `undigested` mode the bytes are copied into `ctx`, so neither + /// return value carries a safety obligation — but the caller cannot tell + /// the modes apart, so it must keep both alive regardless. unsafe fn load_dict>( ctx: &mut L, dict_obj: Option, @@ -856,11 +845,13 @@ mod _zstd { let mut digested = None; match marker { DICT_TYPE_PREFIX => { - // SAFETY: the returned `PyRef` keeps `zdict` alive - // (and therefore `dict_bytes`) for the lifetime of the - // context. libzstd `ref_prefix` stores a raw pointer into - // `dict_bytes`; as long as the (de)compressor outlives no - // longer than the held PyRef, the pointer stays valid. + // SAFETY: we extend `dict_bytes`' lifetime to `'static` only + // to thread it through `ref_prefix_static`'s signature. + // `load_dict`'s own safety contract requires the caller to + // keep the returned `PyRef` (which owns + // `dict_bytes`) alive at least as long as `ctx`, so the + // raw pointer libzstd stores stays valid for the required + // window. let static_bytes: &'static [u8] = unsafe { core::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) }; ctx.ref_prefix_static(static_bytes) @@ -884,24 +875,53 @@ mod _zstd { Ok((digested, Some(zdict))) } - fn load_compressor_dict( - cctx: &mut CCtx<'static>, + /// Build a fully-initialized `CompressorState` from a freshly-created + /// `CCtx` and an optional dictionary argument. This is the safe interface + /// that `unsafe fn load_dict` was waiting for: by assembling the struct + /// here, both invariants `load_dict` documents become structural and a + /// safe-Rust caller cannot split the pieces apart. + fn build_compressor_state( + mut cctx: CCtx<'static>, dict_obj: Option, + last_mode: i32, vm: &VirtualMachine, - ) -> DictLoadResult> { - // SAFETY: The returned `PyRef` is stored alongside the - // `CCtx` in `CompressorState`, so it outlives the context. - unsafe { load_dict::>(cctx, dict_obj, vm) } + ) -> PyResult { + // SAFETY: `load_dict` requires its two return values to outlive `ctx`. + // We satisfy that by moving `cctx` and both return values into + // `CompressorState` in one expression — Rust drops the struct's + // fields in declaration order, so on teardown `cctx` is dropped + // first, releasing its raw pointers before `_cdict` (digested mode) + // and `_dict` (prefix mode) are freed. `CompressorState` is private + // to this module and is never destructured, so no safe caller can + // reorder the drops. + let (cdict, dict) = unsafe { load_dict::>(&mut cctx, dict_obj, vm) }?; + Ok(CompressorState { + cctx, + _cdict: cdict, + _dict: dict, + last_mode, + }) } - fn load_decompressor_dict( - dctx: &mut DCtx<'static>, + /// Build a fully-initialized `DecompressorState`. See + /// [`build_compressor_state`] for the safety reasoning; + /// `DecompressorState`'s field order plays the same role here. + fn build_decompressor_state( + mut dctx: DCtx<'static>, dict_obj: Option, vm: &VirtualMachine, - ) -> DictLoadResult> { - // SAFETY: The returned `PyRef` is stored alongside the - // `DCtx` in `DecompressorState`, so it outlives the context. - unsafe { load_dict::>(dctx, dict_obj, vm) } + ) -> PyResult { + // SAFETY: see [`build_compressor_state`]. + let (ddict, dict) = unsafe { load_dict::>(&mut dctx, dict_obj, vm) }?; + Ok(DecompressorState { + dctx, + _ddict: ddict, + _dict: dict, + eof: false, + needs_input: true, + unused_data: vm.ctx.empty_bytes.clone(), + input_buffer: Vec::new(), + }) } /// Drive `compress_stream2` until the input is fully consumed and, for @@ -913,31 +933,39 @@ mod _zstd { end_op: zstd_sys::ZSTD_EndDirective, vm: &VirtualMachine, ) -> PyResult> { - let mut output = Vec::new(); - let mut input = InBuffer::around(data); - let chunk_size = CCtx::out_size().max(1); + // Release the GIL for the duration of the compression loop. Safety: + // `data` is an immutable borrow of a local `Vec` in the caller, + // `state` is held under the compressor's `PyMutex` (no other Python + // thread can touch it), the `_dict` bytes referenced by libzstd are + // an immutable `Vec` inside a `PyRef` (other readers fine), + // and the output `Vec` is local to this function. No Python object + // access happens inside the closure — error codes are surfaced as + // `usize` and converted into exceptions after re-attaching. let is_end = end_op != zstd_sys::ZSTD_EndDirective::ZSTD_e_continue; - - loop { - let prev_len = output.len(); - output.reserve(chunk_size); - let remaining = { - let mut out_buf = OutBuffer::around_pos(&mut output, prev_len); - state - .cctx - .compress_stream2(&mut out_buf, &mut input, end_op) - } - .map_err(|c| catch_zstd_error(c, vm))?; - let consumed_all = input.pos == input.src.len(); - // Stop when input is fully consumed and, for flush/end directives, - // libzstd reports that all internal buffers have been drained - // (remaining == 0). Otherwise loop; the next `reserve` will grow - // the output buffer if we hit the previous capacity. - if consumed_all && (!is_end || remaining == 0) { - break; + let chunk_size = CCtx::out_size().max(1); + let result: Result, usize> = vm.allow_threads(|| { + let mut output = Vec::new(); + let mut input = InBuffer::around(data); + loop { + let prev_len = output.len(); + output.reserve(chunk_size); + let remaining = { + let mut out_buf = OutBuffer::around_pos(&mut output, prev_len); + state + .cctx + .compress_stream2(&mut out_buf, &mut input, end_op) + }?; + let consumed_all = input.pos == input.src.len(); + // Stop when input is fully consumed and, for flush/end + // directives, libzstd reports that all internal buffers have + // been drained (remaining == 0). Otherwise loop; the next + // `reserve` will grow the output if we hit the previous cap. + if consumed_all && (!is_end || remaining == 0) { + break Ok(output); + } } - } - Ok(output) + }); + result.map_err(|c| catch_zstd_error(c, vm)) } #[pyclass(with(Constructor))] @@ -1090,17 +1118,17 @@ mod _zstd { #[derive(FromArgs)] pub(super) struct ZstdDecompressorArgs { #[pyarg(any, optional)] - zstd_dict: OptionalArg, + zstd_dict: OptionalOption, #[pyarg(any, optional)] - options: OptionalArg, + options: OptionalOption, } impl Constructor for ZstdDecompressor { type Args = ZstdDecompressorArgs; fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { - let dict_opt = arg_or_none(args.zstd_dict, vm); - let options_opt = arg_or_none(args.options, vm); + let dict_opt = args.zstd_dict.flatten(); + let options_opt = args.options.flatten(); let mut dctx = DCtx::<'static>::create(); @@ -1108,18 +1136,9 @@ mod _zstd { apply_options(&mut dctx, options_obj, false, vm)?; } - let (held_ddict, held_dict) = load_decompressor_dict(&mut dctx, dict_opt, vm)?; - + let state = build_decompressor_state(dctx, dict_opt, vm)?; Ok(Self { - state: PyMutex::new(DecompressorState { - dctx, - _ddict: held_ddict, - _dict: held_dict, - eof: false, - needs_input: true, - unused_data: vm.ctx.empty_bytes.clone(), - input_buffer: Vec::new(), - }), + state: PyMutex::new(state), }) } } @@ -1160,70 +1179,81 @@ mod _zstd { alloc::borrow::Cow::Owned(combined) }; - let mut input = InBuffer::around(&work_data); - let mut output: Vec = Vec::new(); let chunk_size = DCtx::out_size().max(1); - // Reusable scratch buffer for each decompress_stream call. We need an - // exact-size output buffer because `Vec::reserve` may over-allocate; - // OutBuffer reports the full Vec capacity to libzstd, which would - // then happily write past `max_length`. - let mut scratch: Vec = vec![0u8; chunk_size]; - let mut hit_max = false; - let mut iteration = 0usize; - - loop { - iteration += 1; - // Honor `max_length`: stop growing the output buffer once we have - // produced enough. Special-case the first iteration when the cap - // is zero so a zero-output frame (skippable frame, empty content - // frame) can still complete; we hand libzstd a 1-byte scratch - // and discard the byte if it ends up writing one. - let grow = match max_length { - Some(maxl) if output.len() >= maxl && iteration > 1 => { - hit_max = true; - break; - } - Some(maxl) if output.len() >= maxl => 1, - Some(maxl) => (maxl - output.len()).min(chunk_size), - None => chunk_size, - }; - let result; - let written; - { - let slot = &mut scratch[..grow]; - let mut out_buf = OutBuffer::around(slot as &mut [u8]); - result = state.dctx.decompress_stream(&mut out_buf, &mut input); - written = out_buf.pos(); - } - output.extend_from_slice(&scratch[..written]); - match result { - Ok(0) => { - // Frame fully decompressed; the decompressor is at EOF. - state.eof = true; - break; - } - Ok(_) => { - let output_was_full = written == grow; - let input_consumed = input.pos == input.src.len(); - - if let Some(maxl) = max_length - && output.len() >= maxl - && iteration > 1 - { + // Release the GIL for the streaming loop. Safety: see `do_compress`; + // the closure captures only Rust-owned buffers and `&mut state` (held + // under the decompressor's `PyMutex`), and surfaces error codes as + // `usize` so we can build the exception after re-attaching. + let loop_result: Result<(Vec, bool, usize), usize> = vm.allow_threads(|| { + let mut input = InBuffer::around(&work_data); + let mut output: Vec = Vec::new(); + // Reusable scratch buffer for each decompress_stream call. We need + // an exact-size output buffer because `Vec::reserve` may + // over-allocate; `OutBuffer` reports the full Vec capacity to + // libzstd, which would then happily write past `max_length`. + let mut scratch: Vec = vec![0u8; chunk_size]; + let mut hit_max = false; + let mut iteration = 0usize; + + let outcome = loop { + iteration += 1; + // Honor `max_length`: stop growing the output buffer once + // we have produced enough. Special-case the first iteration + // when the cap is zero so a zero-output frame (skippable + // frame, empty content frame) can still complete; we hand + // libzstd a 1-byte scratch and discard the byte if it ends + // up writing one. + let grow = match max_length { + Some(maxl) if output.len() >= maxl && iteration > 1 => { hit_max = true; - break; + break Ok(()); } - - // Input is gone and libzstd had room to write but did not, - // which means the frame is incomplete and the caller has - // to supply more input. - if input_consumed && !output_was_full { - break; + Some(maxl) if output.len() >= maxl => 1, + Some(maxl) => (maxl - output.len()).min(chunk_size), + None => chunk_size, + }; + let result; + let written; + { + let slot = &mut scratch[..grow]; + let mut out_buf = OutBuffer::around(slot as &mut [u8]); + result = state.dctx.decompress_stream(&mut out_buf, &mut input); + written = out_buf.pos(); + } + output.extend_from_slice(&scratch[..written]); + match result { + Ok(0) => { + // Frame fully decompressed; the decompressor is at EOF. + state.eof = true; + break Ok(()); } + Ok(_) => { + let output_was_full = written == grow; + let input_consumed = input.pos == input.src.len(); + + if let Some(maxl) = max_length + && output.len() >= maxl + && iteration > 1 + { + hit_max = true; + break Ok(()); + } + + // Input is gone and libzstd had room to write but did + // not, which means the frame is incomplete and the + // caller has to supply more input. + if input_consumed && !output_was_full { + break Ok(()); + } + } + Err(code) => break Err(code), } - Err(code) => return Err(catch_zstd_error(code, vm)), - } - } + }; + outcome.map(|()| (output, hit_max, input.pos)) + }); + + let (mut output, mut hit_max, consumed) = + loop_result.map_err(|c| catch_zstd_error(c, vm))?; // If `max_length == 0` opened a courtesy iteration that produced more // bytes than the caller asked for, truncate. Should not happen with @@ -1235,7 +1265,6 @@ mod _zstd { hit_max = true; } - let consumed = input.pos; let remaining = &work_data[consumed..]; if state.eof { @@ -1396,17 +1425,34 @@ mod _zstd { Ok(v as usize) } + /// Sum the per-sample sizes and check they exactly cover `expected_total`, + /// rejecting overflow. A safe-fn interface that calls into libzstd with + /// these sizes must not let a wrapping sum sneak past the equality check, + /// since libzstd would then read past the samples buffer. + fn check_sample_sizes_match( + sizes: &[usize], + expected_total: usize, + vm: &VirtualMachine, + ) -> PyResult<()> { + let mismatch = || -> PyBaseExceptionRef { + vm.new_value_error("The samples size tuple doesn't match the concatenation's size") + }; + let total = sizes + .iter() + .try_fold(0usize, |a, &b| a.checked_add(b)) + .ok_or_else(mismatch)?; + if total != expected_total { + return Err(mismatch()); + } + Ok(()) + } + #[pyfunction] fn train_dict(args: TrainDictArgs, vm: &VirtualMachine) -> PyResult> { let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; let samples_buffer = args.samples_bytes.as_bytes().to_vec(); let sizes = parse_sample_sizes(args.samples_sizes, vm)?; - let total: usize = sizes.iter().sum(); - if total != samples_buffer.len() { - return Err( - vm.new_value_error("The samples size tuple doesn't match the concatenation's size") - ); - } + check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; let mut dict_buffer: Vec = Vec::with_capacity(dict_size); zstd_safe::train_from_buffer(&mut dict_buffer, &samples_buffer, &sizes) .map_err(|c| catch_zstd_error(c, vm))?; @@ -1439,16 +1485,11 @@ mod _zstd { #[pyfunction] fn finalize_dict(args: FinalizeDictArgs, vm: &VirtualMachine) -> PyResult> { let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; - let compression_level = pyobj_to_i32(&args.compression_level, vm)?; + let compression_level: i32 = args.compression_level.try_to_value(vm)?; let custom_dict = args.custom_dict_bytes.as_bytes().to_vec(); let samples_buffer = args.samples_bytes.as_bytes().to_vec(); let sizes = parse_sample_sizes(args.samples_sizes, vm)?; - let total: usize = sizes.iter().sum(); - if total != samples_buffer.len() { - return Err( - vm.new_value_error("The samples size tuple doesn't match the concatenation's size") - ); - } + check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; let mut dict_buffer: Vec = vec![0u8; dict_size]; let params = zstd_sys::ZDICT_params_t { From acfb59b7c846c7b9d4e6a68ff58419164c3267ff Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 3 Jun 2026 14:48:21 +0100 Subject: [PATCH 10/25] Identify zstd parameter enums by type identity, not name - Have set_parameter_types stash the CompressionParameter/DecompressionParameter classes as private _zstd module attributes instead of validating and discarding them - check_wrong_param_kind now compares key classes by identity against the registered type, matching CPython's Py_TYPE check, and skips when unregistered - Error message names the actual key type as an attribute --- crates/stdlib/src/zstd.rs | 87 ++++++++++++++++++--------------------- 1 file changed, 41 insertions(+), 46 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 58f5f807996..c17edae6559 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -39,16 +39,6 @@ mod _zstd { use zstd_safe::zstd_sys; use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; - /// Class names of the `CompressionParameter` and `DecompressionParameter` - /// `IntEnum`s defined by the pure-Python wrapper at - /// `Lib/compression/zstd/__init__.py`. We identify these by name rather - /// than identity because storing `PyTypeRef` in a Rust `static` requires - /// `Sync`, which `PyMutex` does not provide on single-threaded targets - /// (Android, iOS, wasm32). Name comparison is sufficient since both names - /// originate from a module we ship and control. - const C_PARAMETER_TYPE_NAME: &str = "CompressionParameter"; - const D_PARAMETER_TYPE_NAME: &str = "DecompressionParameter"; - // ========================================================================= // Module-level constants // ========================================================================= @@ -177,29 +167,30 @@ mod _zstd { // Parameter helpers // ========================================================================= - /// If `key`'s class is the wrong kind of parameter enum for the caller's - /// context (a `CompressionParameter` passed to a decompressor, or vice - /// versa), return a `TypeError`. Used by both (de)compressor option-dict - /// loops so the resulting message names the actual type rather than the - /// generic "out of range" wording. + /// Reject an options-dict `key` whose class is the parameter enum that is + /// invalid for the caller's context (a `CompressionParameter` passed to a + /// decompressor, or vice versa) with a `TypeError` naming the type. /// - /// We identify the enums by class name (see [`C_PARAMETER_TYPE_NAME`]), - /// which avoids storing a `PyTypeRef` in a Rust `static` (impossible on - /// single-threaded builds whose `PyMutex` is not `Sync`). + /// `forbidden` is that invalid enum class, resolved once by the caller from + /// the types the pure-Python wrapper registers via [`set_parameter_types`]; + /// `None` (the wrapper never ran) skips the check, matching CPython's NULL + /// module-state pointers. The comparison is by identity, mirroring + /// CPython's `Py_TYPE(key) == ...` check. fn check_wrong_param_kind( key: &PyObjectRef, - is_compress: bool, + forbidden: Option<&PyObjectRef>, + kind: &str, vm: &VirtualMachine, ) -> PyResult<()> { - let class_name = key.class().name(); - let (forbidden, kind) = if is_compress { - (D_PARAMETER_TYPE_NAME, "compression") - } else { - (C_PARAMETER_TYPE_NAME, "decompression") + let Some(forbidden) = forbidden else { + return Ok(()); }; - if &*class_name == forbidden { + if key.class().is(forbidden) { + // `key`'s class is `forbidden` here, so name it directly (the same + // string CPython formats from `Py_TYPE(key)->tp_name`). + let name = key.class().name(); return Err(vm.new_type_error(format!( - "{kind} options dictionary key must not be a {forbidden}" + "{kind} options dictionary key must not be a {name} attribute" ))); } Ok(()) @@ -705,12 +696,21 @@ mod _zstd { let dict = options_obj .downcast::() .map_err(|_| vm.new_type_error("options must be a dict"))?; + // Resolve the parameter enum class that is invalid for this direction + // (registered by the pure-Python wrapper via set_parameter_types) once, + // so a key from the wrong family yields a clear TypeError naming the + // type. `None` when the wrapper never ran (e.g. `_zstd` used directly). + let (wrong_kind_attr, kind) = if is_compress { + ("_decompression_parameter_type", "compression") + } else { + ("_compression_parameter_type", "decompression") + }; + let wrong_kind = vm.get_attribute_opt(vm.import("_zstd", 0)?, wrong_kind_attr)?; for (k, v) in dict { - // Reject the wrong enum kind ("CompressionParameter" vs - // "DecompressionParameter") before doing any numeric coercion, - // so the error carries the documented type name rather than a + // Reject a key from the wrong parameter family before any numeric + // coercion, so the error names the type rather than giving a // generic out-of-range message. - check_wrong_param_kind(&k, is_compress, vm)?; + check_wrong_param_kind(&k, wrong_kind.as_ref(), kind, vm)?; let key_int: i32 = k.try_to_value(vm)?; let val_int: i32 = v.try_to_value(vm)?; // libzstd silently clamps out-of-range values for some @@ -1572,26 +1572,21 @@ mod _zstd { Ok((bounds.lowerBound, bounds.upperBound)) } - // The (de)compressor constructors identify the parameter enums by class - // *name* (see [`check_wrong_param_kind`]) rather than identity, so this - // function does not need to retain the type objects beyond the call. - // CPython retains them to produce nicer error messages naming the specific - // enum member; matching that would require a `Sync` cell-of-`PyTypeRef`, - // which is not available on single-threaded targets. The pure-Python - // wrapper at `Lib/compression/zstd/__init__.py:242` calls this exactly - // once at import time. + // Register the `CompressionParameter` / `DecompressionParameter` enum + // classes defined by the pure-Python wrapper so [`check_wrong_param_kind`] + // can reject a key from the wrong parameter family by identity. The types + // are stashed as private `_zstd` module attributes — the RustPython + // equivalent of the module state CPython keeps these in. The wrapper in + // `Lib/compression/zstd/__init__.py` calls this exactly once at import. #[pyfunction] fn set_parameter_types( - c_parameter_type: PyObjectRef, - d_parameter_type: PyObjectRef, + c_parameter_type: PyTypeRef, + d_parameter_type: PyTypeRef, vm: &VirtualMachine, ) -> PyResult<()> { - let _ = c_parameter_type - .downcast::() - .map_err(|_| vm.new_type_error("c_parameter_type must be a type object"))?; - let _ = d_parameter_type - .downcast::() - .map_err(|_| vm.new_type_error("d_parameter_type must be a type object"))?; + let module = vm.import("_zstd", 0)?; + module.set_attr("_compression_parameter_type", c_parameter_type, vm)?; + module.set_attr("_decompression_parameter_type", d_parameter_type, vm)?; Ok(()) } } From cca16a0276629f0db0578366932f4237ab8240ef Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 3 Jun 2026 15:24:11 +0100 Subject: [PATCH 11/25] Match CPython error behavior in _zstd ZstdCompressor - Raise TypeError (not RuntimeError), matching CPython, when both `level` and `options` are passed to ZstdCompressor - Parse set_pledged_input_size's argument before taking the state lock so a re-entrant __index__ can't deadlock, and fix the off-by-one in its bound message to match CPython Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/stdlib/src/zstd.rs | 51 ++++++++++++++++++++++----------------- 1 file changed, 29 insertions(+), 22 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index c17edae6559..e6f96b60a1a 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -638,7 +638,7 @@ mod _zstd { let dict_opt = args.zstd_dict.flatten(); if level_opt.is_some() && options_opt.is_some() { - return Err(vm.new_runtime_error("Only one of level or options should be used")); + return Err(vm.new_type_error("Only one of level or options should be used.")); } let mut cctx = CCtx::<'static>::create(); @@ -1000,37 +1000,44 @@ mod _zstd { #[pymethod] fn set_pledged_input_size(&self, size: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - let mut state = self.state.lock(); - if state.last_mode != COMP_MODE_FLUSH_FRAME { - return Err(vm.new_value_error( - "set_pledged_input_size() method must be called when last_mode == FLUSH_FRAME", - )); - } + // Parse the argument *before* taking the lock: `try_index` can run a + // Python `__index__`, and doing that while holding `self.state` would + // let a re-entrant call into this compressor deadlock. CPython + // likewise converts the argument before touching the compressor. + // // Python passes `None` to mean "unknown"; libzstd represents that // internally as `ZSTD_CONTENTSIZE_UNKNOWN` (`u64::MAX`), and - // `zstd_safe` translates `None` accordingly. Reject any concrete - // int outside the valid range up front so callers see the - // documented `ValueError`, not a libzstd-level error. The valid - // upper bound is `ZSTD_CONTENTSIZE_ERROR - 1`, i.e. `u64::MAX - 2`. + // `zstd_safe` translates `None` accordingly. libzstd also reserves + // `ZSTD_CONTENTSIZE_ERROR` (`u64::MAX - 1`), so a concrete size must + // be strictly less than that; reject anything else up front so + // callers see the documented `ValueError`, not a libzstd-level error. let pledged: Option = if vm.is_none(&size) { None } else { - const UPPER_BOUND: u64 = u64::MAX - 2; - let int_ref = size.try_index(vm)?; - // Try fitting into u64; OverflowError covers negatives and - // values above `u64::MAX`. - let v: u64 = int_ref.try_to_primitive(vm).map_err(|_| { + const LIMIT: u64 = u64::MAX - 1; + let err = || { vm.new_value_error(format!( - "size argument should be a positive int less than {UPPER_BOUND}" + "size argument should be a positive int less than {LIMIT}" )) - })?; - if v > UPPER_BOUND { - return Err(vm.new_value_error(format!( - "size argument should be a positive int less than {UPPER_BOUND}" - ))); + }; + // `try_to_primitive` fails (OverflowError) for negatives and for + // values above `u64::MAX`; the explicit check covers the rest of + // the reserved range. + let v: u64 = size + .try_index(vm)? + .try_to_primitive(vm) + .map_err(|_| err())?; + if v >= LIMIT { + return Err(err()); } Some(v) }; + let mut state = self.state.lock(); + if state.last_mode != COMP_MODE_FLUSH_FRAME { + return Err(vm.new_value_error( + "set_pledged_input_size() method must be called when last_mode == FLUSH_FRAME", + )); + } state .cctx .set_pledged_src_size(pledged) From 5b558cfb37851964f53b859a88fe4d2cd46a8fd1 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sun, 9 Aug 2026 16:31:50 +0100 Subject: [PATCH 12/25] Address review: drop decorative separators and unused zstd crate Remove the remaining `// ===...` section separators from zstd.rs per the project comment guidelines, drop the unused `zstd` crate dependency (only zstd-safe/zstd-sys are used) and declare zstd-safe's std feature explicitly instead of relying on the zstd crate's feature unification, and convert a doc comment on #[extend_class] to a regular comment. --- Cargo.toml | 3 +-- crates/stdlib/Cargo.toml | 1 - crates/stdlib/src/zstd.rs | 36 ++++-------------------------------- 3 files changed, 5 insertions(+), 35 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 5775e41b86c..e1b8c232e30 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -320,8 +320,7 @@ which = "8" x509-parser = "0.18" xml = "1.3" writeable = "0.6" -zstd = { version = "0.13", default-features = false, features = ["zdict_builder"] } -zstd-safe = { version = "7", default-features = false, features = ["zdict_builder"] } +zstd-safe = { version = "7", default-features = false, features = ["std", "zdict_builder"] } # Lints diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index 2586138ac79..47747c7a7f9 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -99,7 +99,6 @@ webpki-roots = { workspace = true, optional = true } libsqlite3-sys = { workspace = true, features = ["bundled"], optional = true } xz = { workspace = true } xz-sys = { workspace = true } -zstd = { workspace = true } zstd-safe = { workspace = true } [target.'cfg(windows)'.dependencies] diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index e6f96b60a1a..2c7e1558d02 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -39,10 +39,6 @@ mod _zstd { use zstd_safe::zstd_sys; use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; - // ========================================================================= - // Module-level constants - // ========================================================================= - #[pyattr] const ZSTD_CLEVEL_DEFAULT: i32 = zstd_sys::ZSTD_CLEVEL_DEFAULT as i32; @@ -138,10 +134,6 @@ mod _zstd { const DICT_TYPE_UNDIGESTED: i32 = 1; const DICT_TYPE_PREFIX: i32 = 2; - // ========================================================================= - // ZstdError exception - // ========================================================================= - #[pyattr(once, name = "ZstdError")] fn zstd_error(vm: &VirtualMachine) -> PyTypeRef { vm.ctx.new_exception_type( @@ -163,10 +155,6 @@ mod _zstd { new_zstd_error(zstd_safe::get_error_name(code).to_string(), vm) } - // ========================================================================= - // Parameter helpers - // ========================================================================= - /// Reject an options-dict `key` whose class is the parameter enum that is /// invalid for the caller's context (a `CompressionParameter` passed to a /// decompressor, or vice versa) with a `TypeError` naming the type. @@ -442,10 +430,6 @@ mod _zstd { Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")) } - // ========================================================================= - // ZstdDict - // ========================================================================= - #[derive(FromArgs)] pub(super) struct ZstdDictArgs { #[pyarg(positional)] @@ -558,10 +542,6 @@ mod _zstd { } } - // ========================================================================= - // ZstdCompressor - // ========================================================================= - // The three flush modes for `ZstdCompressor.compress()`, mirrored as // class attributes via `extend_class` below. Values are positional and // chosen to match what CPython exposes. @@ -1050,10 +1030,10 @@ mod _zstd { self.state.lock().last_mode } - /// Install class-level constants `CONTINUE`, `FLUSH_BLOCK`, and - /// `FLUSH_FRAME` so callers can reference them as - /// `ZstdCompressor.FLUSH_FRAME` (as the Python `ZstdFile` wrapper - /// does). + // Install class-level constants `CONTINUE`, `FLUSH_BLOCK`, and + // `FLUSH_FRAME` so callers can reference them as + // `ZstdCompressor.FLUSH_FRAME` (as the Python `ZstdFile` wrapper + // does). #[extend_class] fn extend_class(ctx: &Context, class: &'static Py) { class.set_attr( @@ -1085,10 +1065,6 @@ mod _zstd { mode: Option, } - // ========================================================================= - // ZstdDecompressor - // ========================================================================= - /// Internal state of a `ZstdDecompressor`. The CPython decompressor is /// single-frame: once we hit end-of-frame, additional bytes go into /// `unused_data` and further `decompress` calls raise `EOFError`. Field @@ -1339,10 +1315,6 @@ mod _zstd { } } - // ========================================================================= - // Module-level functions - // ========================================================================= - #[pyfunction] fn get_frame_size(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { let buf = frame_buffer.with_ref(|b| b.to_vec()); From 265e5914be28edf38ae9bdb8766bf5d8c3ad0f93 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sun, 9 Aug 2026 17:26:37 +0100 Subject: [PATCH 13/25] Switch _zstd backend from zstd-safe to pure-Rust libzstd-rs-sys Replace the C-library bindings (zstd-safe/zstd-sys) with Trifecta Tech Foundation's libzstd-rs-sys (pinned git rev), a c2rust translation of upstream libzstd. The module now calls its C-API-compatible functions through small owning RAII wrappers (CCtx/DCtx/CDict/DDict) with the same drop-order invariants as before. The new backend builds zstdmt_compress in, so nb_workers bounds are (0, 256) and test_zstd_multithread_compress passes for real; drop the now-obsolete expectedFailureIf marker. --- Cargo.lock | 10 + Cargo.toml | 2 +- Lib/test/test_zstd.py | 5 - crates/stdlib/Cargo.toml | 2 +- crates/stdlib/src/zstd.rs | 824 +++++++++++++++++++++++++------------- 5 files changed, 549 insertions(+), 294 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fb915315f66..1cf7ede336f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2321,6 +2321,15 @@ dependencies = [ "zlib-rs", ] +[[package]] +name = "libzstd-rs-sys" +version = "0.0.1-prerelease.2" +source = "git+https://github.com/trifectatechfoundation/libzstd-rs-sys?rev=99fb4c096a21a9b37feebe180333a2ff8645e23b#99fb4c096a21a9b37feebe180333a2ff8645e23b" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -4024,6 +4033,7 @@ dependencies = [ "itertools 0.15.0", "libc", "libsqlite3-sys", + "libzstd-rs-sys", "malachite-bigint", "memchr", "mt19937", diff --git a/Cargo.toml b/Cargo.toml index e1b8c232e30..eb904821767 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -241,6 +241,7 @@ xz-sys = "0.4.6-rc.0" libsqlite3-sys = "0.38" libz-rs-sys = "0.6" zlib-rs = { version = "=0.6.8", default-features = false, features = ["rust-allocator", "__internal-api"] } +libzstd-rs-sys = { git = "https://github.com/trifectatechfoundation/libzstd-rs-sys", rev = "99fb4c096a21a9b37feebe180333a2ff8645e23b" } lock_api = "0.4" log = "0.4.30" lz4_flex = "0.14" @@ -320,7 +321,6 @@ which = "8" x509-parser = "0.18" xml = "1.3" writeable = "0.6" -zstd-safe = { version = "7", default-features = false, features = ["std", "zdict_builder"] } # Lints diff --git a/Lib/test/test_zstd.py b/Lib/test/test_zstd.py index 1b4b819d2ef..44a5bbb4b19 100644 --- a/Lib/test/test_zstd.py +++ b/Lib/test/test_zstd.py @@ -321,11 +321,6 @@ def test_unknown_compression_parameter(self): with self.assertRaisesRegex(ValueError, pattern): ZstdCompressor(options=option) - # TODO: RUSTPYTHON - @unittest.expectedFailureIf( - not SUPPORT_MULTITHREADING, - "TODO: RUSTPYTHON; libzstd built without multi-threaded compression support", - ) def test_zstd_multithread_compress(self): size = 40*_1M b = THIS_FILE_BYTES * (size // len(THIS_FILE_BYTES)) diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index 47747c7a7f9..1493d821c3c 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -97,9 +97,9 @@ webpki-roots = { workspace = true, optional = true } [target.'cfg(not(any(target_os = "android", target_arch = "wasm32")))'.dependencies] libsqlite3-sys = { workspace = true, features = ["bundled"], optional = true } +libzstd-rs-sys = { workspace = true } xz = { workspace = true } xz-sys = { workspace = true } -zstd-safe = { workspace = true } [target.'cfg(windows)'.dependencies] paste = { workspace = true } diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 2c7e1558d02..8d9e194bfd6 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -3,7 +3,7 @@ // spell-checker:ignore btlazy dfast nbworkers windowlogmax windowlog overlap targetcblock // spell-checker:ignore srcsize zdict refprefix refcdict refddict pledgedsrcsize getframecontentsize // spell-checker:ignore Zstd Zstandard pylib RFC -// spell-checker:ignore CLEVEL zstdmodule cparameter dparameter maxl +// spell-checker:ignore CLEVEL zstdmodule cparameter dparameter maxl c2rust // spell-checker:ignore cctx dctx CCTX DCTX ldm cdict ddict windowlog hashlog chainlog searchlog CLEVEL //! The `_zstd` extension module. Backs the pure-Python `compression.zstd` @@ -12,10 +12,9 @@ //! `Lib/compression/zstd/__init__.py` imports from this module unconditionally, //! so the names and call signatures here must stay in sync with CPython. //! -//! Backend: the `zstd_safe` Rust crate (a thin safe wrapper over Facebook's -//! libzstd C library, which is what CPython links against). A handful of -//! routines that `zstd_safe` does not expose at a safe level (parameter -//! bounds, dictionary finalization) drop down to raw `zstd_sys` FFI calls. +//! Backend: the `libzstd-rs-sys` crate (Trifecta Tech Foundation's pure-Rust +//! c2rust translation of Facebook's libzstd, the C library CPython links +//! against), used directly through its `unsafe extern "C"` API surface. pub(crate) use _zstd::module_def; @@ -26,7 +25,23 @@ pub(crate) use _zstd::module_def; #[allow(non_upper_case_globals)] #[pymodule] mod _zstd { - use core::ffi::c_int; + use core::ffi::{CStr, c_int}; + use libzstd_rs_sys::lib::compress::zstd_compress::{ZSTD_e_continue, ZSTD_e_end, ZSTD_e_flush}; + use libzstd_rs_sys::lib::zdict::ZDICT_finalizeDictionary; + use libzstd_rs_sys::{ + ZDICT_getErrorName, ZDICT_isError, ZDICT_params_t, ZDICT_trainFromBuffer, ZSTD_CCtx, + ZSTD_CCtx_loadDictionary, ZSTD_CCtx_refCDict, ZSTD_CCtx_refPrefix, ZSTD_CCtx_setParameter, + ZSTD_CCtx_setPledgedSrcSize, ZSTD_CDict, ZSTD_CONTENTSIZE_ERROR, ZSTD_CONTENTSIZE_UNKNOWN, + ZSTD_CStreamOutSize, ZSTD_DCtx, ZSTD_DCtx_loadDictionary, ZSTD_DCtx_refDDict, + ZSTD_DCtx_refPrefix, ZSTD_DCtx_setParameter, ZSTD_DDict, ZSTD_DStreamOutSize, + ZSTD_EndDirective, ZSTD_cParam_getBounds, ZSTD_cParameter, ZSTD_compressStream2, + ZSTD_createCCtx, ZSTD_createCDict, ZSTD_createDCtx, ZSTD_createDDict, + ZSTD_dParam_getBounds, ZSTD_dParameter, ZSTD_decompressStream, + ZSTD_findFrameCompressedSize, ZSTD_freeCCtx, ZSTD_freeCDict, ZSTD_freeDCtx, ZSTD_freeDDict, + ZSTD_getDictID_fromDict, ZSTD_getDictID_fromFrame, ZSTD_getErrorName, + ZSTD_getFrameContentSize, ZSTD_inBuffer, ZSTD_isError, ZSTD_outBuffer, ZSTD_strategy, + ZSTD_versionNumber, ZSTD_versionString, + }; use rustpython_common::lock::PyMutex; use rustpython_vm::builtins::{ PyBaseExceptionRef, PyBytesRef, PyDict, PyTupleRef, PyType, PyTypeRef, @@ -36,94 +51,98 @@ mod _zstd { use rustpython_vm::{ AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, }; - use zstd_safe::zstd_sys; - use zstd_safe::{CCtx, CParameter, DCtx, DParameter, InBuffer, OutBuffer}; #[pyattr] - const ZSTD_CLEVEL_DEFAULT: i32 = zstd_sys::ZSTD_CLEVEL_DEFAULT as i32; + const ZSTD_CLEVEL_DEFAULT: i32 = libzstd_rs_sys::ZSTD_CLEVEL_DEFAULT; // Compression parameter identifiers. Values match the `ZSTD_cParameter` // enum in libzstd, which is what the public `CompressionParameter` IntEnum // in `Lib/compression/zstd/__init__.py` derives its members from. + // libzstd-rs-sys models the enum as a newtype whose inner value is + // crate-private, so the ids — frozen C ABI values from zstd.h — are + // spelled out as literals here. #[pyattr] - const ZSTD_c_compressionLevel: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_compressionLevel as i32; + const ZSTD_c_compressionLevel: i32 = 100; #[pyattr] - const ZSTD_c_windowLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_windowLog as i32; + const ZSTD_c_windowLog: i32 = 101; #[pyattr] - const ZSTD_c_hashLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_hashLog as i32; + const ZSTD_c_hashLog: i32 = 102; #[pyattr] - const ZSTD_c_chainLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_chainLog as i32; + const ZSTD_c_chainLog: i32 = 103; #[pyattr] - const ZSTD_c_searchLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_searchLog as i32; + const ZSTD_c_searchLog: i32 = 104; #[pyattr] - const ZSTD_c_minMatch: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_minMatch as i32; + const ZSTD_c_minMatch: i32 = 105; #[pyattr] - const ZSTD_c_targetLength: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_targetLength as i32; + const ZSTD_c_targetLength: i32 = 106; #[pyattr] - const ZSTD_c_strategy: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_strategy as i32; + const ZSTD_c_strategy: i32 = 107; #[pyattr] - const ZSTD_c_enableLongDistanceMatching: i32 = - zstd_sys::ZSTD_cParameter::ZSTD_c_enableLongDistanceMatching as i32; + const ZSTD_c_enableLongDistanceMatching: i32 = 160; #[pyattr] - const ZSTD_c_ldmHashLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmHashLog as i32; + const ZSTD_c_ldmHashLog: i32 = 161; #[pyattr] - const ZSTD_c_ldmMinMatch: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmMinMatch as i32; + const ZSTD_c_ldmMinMatch: i32 = 162; #[pyattr] - const ZSTD_c_ldmBucketSizeLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmBucketSizeLog as i32; + const ZSTD_c_ldmBucketSizeLog: i32 = 163; #[pyattr] - const ZSTD_c_ldmHashRateLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_ldmHashRateLog as i32; + const ZSTD_c_ldmHashRateLog: i32 = 164; #[pyattr] - const ZSTD_c_contentSizeFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_contentSizeFlag as i32; + const ZSTD_c_contentSizeFlag: i32 = 200; #[pyattr] - const ZSTD_c_checksumFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_checksumFlag as i32; + const ZSTD_c_checksumFlag: i32 = 201; #[pyattr] - const ZSTD_c_dictIDFlag: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_dictIDFlag as i32; + const ZSTD_c_dictIDFlag: i32 = 202; #[pyattr] - const ZSTD_c_nbWorkers: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_nbWorkers as i32; + const ZSTD_c_nbWorkers: i32 = 400; #[pyattr] - const ZSTD_c_jobSize: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_jobSize as i32; + const ZSTD_c_jobSize: i32 = 401; #[pyattr] - const ZSTD_c_overlapLog: i32 = zstd_sys::ZSTD_cParameter::ZSTD_c_overlapLog as i32; + const ZSTD_c_overlapLog: i32 = 402; // Decompression parameter identifiers. libzstd only exposes one non- // experimental decompression parameter. #[pyattr] - const ZSTD_d_windowLogMax: i32 = zstd_sys::ZSTD_dParameter::ZSTD_d_windowLogMax as i32; + const ZSTD_d_windowLogMax: i32 = 100; // Strategy enum members ordered from fastest to strongest. These power // the `Strategy` IntEnum in `Lib/compression/zstd/__init__.py`. #[pyattr] - const ZSTD_fast: i32 = zstd_sys::ZSTD_strategy::ZSTD_fast as i32; + const ZSTD_fast: i32 = libzstd_rs_sys::lib::zstd::ZSTD_fast as i32; #[pyattr] - const ZSTD_dfast: i32 = zstd_sys::ZSTD_strategy::ZSTD_dfast as i32; + const ZSTD_dfast: i32 = libzstd_rs_sys::lib::zstd::ZSTD_dfast as i32; #[pyattr] - const ZSTD_greedy: i32 = zstd_sys::ZSTD_strategy::ZSTD_greedy as i32; + const ZSTD_greedy: i32 = libzstd_rs_sys::lib::zstd::ZSTD_greedy as i32; #[pyattr] - const ZSTD_lazy: i32 = zstd_sys::ZSTD_strategy::ZSTD_lazy as i32; + const ZSTD_lazy: i32 = libzstd_rs_sys::lib::zstd::ZSTD_lazy as i32; #[pyattr] - const ZSTD_lazy2: i32 = zstd_sys::ZSTD_strategy::ZSTD_lazy2 as i32; + const ZSTD_lazy2: i32 = libzstd_rs_sys::lib::zstd::ZSTD_lazy2 as i32; #[pyattr] - const ZSTD_btlazy2: i32 = zstd_sys::ZSTD_strategy::ZSTD_btlazy2 as i32; + const ZSTD_btlazy2: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btlazy2 as i32; #[pyattr] - const ZSTD_btopt: i32 = zstd_sys::ZSTD_strategy::ZSTD_btopt as i32; + const ZSTD_btopt: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btopt as i32; #[pyattr] - const ZSTD_btultra: i32 = zstd_sys::ZSTD_strategy::ZSTD_btultra as i32; + const ZSTD_btultra: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btultra as i32; #[pyattr] - const ZSTD_btultra2: i32 = zstd_sys::ZSTD_strategy::ZSTD_btultra2 as i32; + const ZSTD_btultra2: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btultra2 as i32; #[pyattr(once, name = "zstd_version")] fn zstd_version(_vm: &VirtualMachine) -> String { - zstd_safe::version_string().to_string() + // SAFETY: `ZSTD_versionString` returns a pointer to libzstd's static, + // NUL-terminated version string. + unsafe { CStr::from_ptr(ZSTD_versionString()) } + .to_string_lossy() + .into_owned() } #[pyattr(once, name = "zstd_version_number")] fn zstd_version_number(_vm: &VirtualMachine) -> u32 { - zstd_safe::version_number() + ZSTD_versionNumber() } #[pyattr(once, name = "ZSTD_DStreamOutSize")] fn zstd_dstream_out_size(_vm: &VirtualMachine) -> usize { - DCtx::out_size() + ZSTD_DStreamOutSize() } // Dictionary load type markers. The `ZstdDict.as_*` properties wrap the @@ -152,7 +171,10 @@ mod _zstd { /// functions when `ZSTD_isError(code)` is non-zero) into a `ZstdError` /// carrying the human-readable message from `ZSTD_getErrorName`. fn catch_zstd_error(code: usize, vm: &VirtualMachine) -> PyBaseExceptionRef { - new_zstd_error(zstd_safe::get_error_name(code).to_string(), vm) + // SAFETY: `ZSTD_getErrorName` returns a pointer to a static, + // NUL-terminated error string from libzstd's error table. + let name = unsafe { CStr::from_ptr(ZSTD_getErrorName(code)) }; + new_zstd_error(name.to_string_lossy(), vm) } /// Reject an options-dict `key` whose class is the parameter enum that is @@ -184,92 +206,80 @@ mod _zstd { Ok(()) } - /// Map a compression parameter id to its raw libzstd `ZSTD_cParameter` - /// enum value. Returns `None` for unknown ids so callers can surface a + /// Map a compression parameter id to its libzstd `ZSTD_cParameter` + /// constant. Returns `None` for unknown ids so callers can surface a /// targeted `ValueError`. Done with an explicit match rather than /// `mem::transmute` so passing junk like `ZSTD_cParameter(42)` cannot /// be triggered from Python. - fn c_param_enum(param: i32) -> Option { - use zstd_sys::ZSTD_cParameter as P; + fn c_param_enum(param: i32) -> Option { Some(match param { - ZSTD_c_compressionLevel => P::ZSTD_c_compressionLevel, - ZSTD_c_windowLog => P::ZSTD_c_windowLog, - ZSTD_c_hashLog => P::ZSTD_c_hashLog, - ZSTD_c_chainLog => P::ZSTD_c_chainLog, - ZSTD_c_searchLog => P::ZSTD_c_searchLog, - ZSTD_c_minMatch => P::ZSTD_c_minMatch, - ZSTD_c_targetLength => P::ZSTD_c_targetLength, - ZSTD_c_strategy => P::ZSTD_c_strategy, - ZSTD_c_enableLongDistanceMatching => P::ZSTD_c_enableLongDistanceMatching, - ZSTD_c_ldmHashLog => P::ZSTD_c_ldmHashLog, - ZSTD_c_ldmMinMatch => P::ZSTD_c_ldmMinMatch, - ZSTD_c_ldmBucketSizeLog => P::ZSTD_c_ldmBucketSizeLog, - ZSTD_c_ldmHashRateLog => P::ZSTD_c_ldmHashRateLog, - ZSTD_c_contentSizeFlag => P::ZSTD_c_contentSizeFlag, - ZSTD_c_checksumFlag => P::ZSTD_c_checksumFlag, - ZSTD_c_dictIDFlag => P::ZSTD_c_dictIDFlag, - ZSTD_c_nbWorkers => P::ZSTD_c_nbWorkers, - ZSTD_c_jobSize => P::ZSTD_c_jobSize, - ZSTD_c_overlapLog => P::ZSTD_c_overlapLog, + ZSTD_c_compressionLevel => ZSTD_cParameter::ZSTD_c_compressionLevel, + ZSTD_c_windowLog => ZSTD_cParameter::ZSTD_c_windowLog, + ZSTD_c_hashLog => ZSTD_cParameter::ZSTD_c_hashLog, + ZSTD_c_chainLog => ZSTD_cParameter::ZSTD_c_chainLog, + ZSTD_c_searchLog => ZSTD_cParameter::ZSTD_c_searchLog, + ZSTD_c_minMatch => ZSTD_cParameter::ZSTD_c_minMatch, + ZSTD_c_targetLength => ZSTD_cParameter::ZSTD_c_targetLength, + ZSTD_c_strategy => ZSTD_cParameter::ZSTD_c_strategy, + ZSTD_c_enableLongDistanceMatching => ZSTD_cParameter::ZSTD_c_enableLongDistanceMatching, + ZSTD_c_ldmHashLog => ZSTD_cParameter::ZSTD_c_ldmHashLog, + ZSTD_c_ldmMinMatch => ZSTD_cParameter::ZSTD_c_ldmMinMatch, + ZSTD_c_ldmBucketSizeLog => ZSTD_cParameter::ZSTD_c_ldmBucketSizeLog, + ZSTD_c_ldmHashRateLog => ZSTD_cParameter::ZSTD_c_ldmHashRateLog, + ZSTD_c_contentSizeFlag => ZSTD_cParameter::ZSTD_c_contentSizeFlag, + ZSTD_c_checksumFlag => ZSTD_cParameter::ZSTD_c_checksumFlag, + ZSTD_c_dictIDFlag => ZSTD_cParameter::ZSTD_c_dictIDFlag, + ZSTD_c_nbWorkers => ZSTD_cParameter::ZSTD_c_nbWorkers, + ZSTD_c_jobSize => ZSTD_cParameter::ZSTD_c_jobSize, + ZSTD_c_overlapLog => ZSTD_cParameter::ZSTD_c_overlapLog, _ => return None, }) } - /// Map a decompression parameter id to its raw libzstd `ZSTD_dParameter` - /// enum value. See [`c_param_enum`] for rationale. - fn d_param_enum(param: i32) -> Option { - use zstd_sys::ZSTD_dParameter as P; + /// Map a decompression parameter id to its libzstd `ZSTD_dParameter` + /// constant. See [`c_param_enum`] for rationale. + fn d_param_enum(param: i32) -> Option { match param { - ZSTD_d_windowLogMax => Some(P::ZSTD_d_windowLogMax), + ZSTD_d_windowLogMax => Some(ZSTD_dParameter::ZSTD_d_windowLogMax), _ => None, } } - /// Map a compression-parameter id and integer value to a `CParameter` - /// variant. Used by the compressor's `options=` constructor argument. - fn cparameter_from_int(param: i32, value: i32, vm: &VirtualMachine) -> PyResult { - let p = match param { - ZSTD_c_compressionLevel => CParameter::CompressionLevel(value), - ZSTD_c_windowLog => CParameter::WindowLog(value as u32), - ZSTD_c_hashLog => CParameter::HashLog(value as u32), - ZSTD_c_chainLog => CParameter::ChainLog(value as u32), - ZSTD_c_searchLog => CParameter::SearchLog(value as u32), - ZSTD_c_minMatch => CParameter::MinMatch(value as u32), - ZSTD_c_targetLength => CParameter::TargetLength(value as u32), - ZSTD_c_strategy => { - CParameter::Strategy(strategy_from_int(value).ok_or_else(|| { - new_zstd_error(format!("invalid strategy value: {value}"), vm) - })?) - } - ZSTD_c_enableLongDistanceMatching => CParameter::EnableLongDistanceMatching(value != 0), - ZSTD_c_ldmHashLog => CParameter::LdmHashLog(value as u32), - ZSTD_c_ldmMinMatch => CParameter::LdmMinMatch(value as u32), - ZSTD_c_ldmBucketSizeLog => CParameter::LdmBucketSizeLog(value as u32), - ZSTD_c_ldmHashRateLog => CParameter::LdmHashRateLog(value as u32), - ZSTD_c_contentSizeFlag => CParameter::ContentSizeFlag(value != 0), - ZSTD_c_checksumFlag => CParameter::ChecksumFlag(value != 0), - ZSTD_c_dictIDFlag => CParameter::DictIdFlag(value != 0), - ZSTD_c_nbWorkers => CParameter::NbWorkers(value as u32), - ZSTD_c_jobSize => CParameter::JobSize(value as u32), - ZSTD_c_overlapLog => CParameter::OverlapSizeLog(value as u32), - _ => { - return Err(vm.new_value_error(format!( - "invalid compression parameter 'unknown parameter (key {param})'" - ))); - } - }; + /// Validate a compression-parameter id and value pair, returning the + /// parameter constant on success. Used by the compressor's `options=` + /// constructor argument. The value is only validated here for the + /// strategy parameter (which libzstd does not bounds-check on its own in + /// a way CPython can rely on); everything else is guarded by the bounds + /// pre-validation in `apply_options` and passed through to libzstd + /// untouched, like CPython does. + fn cparameter_from_int( + param: i32, + value: i32, + vm: &VirtualMachine, + ) -> PyResult { + let p = c_param_enum(param).ok_or_else(|| { + vm.new_value_error(format!( + "invalid compression parameter 'unknown parameter (key {param})'" + )) + })?; + if param == ZSTD_c_strategy && strategy_from_int(value).is_none() { + return Err(new_zstd_error( + format!("invalid strategy value: {value}"), + vm, + )); + } Ok(p) } - /// Map a decompression-parameter id and integer value to a `DParameter`. - /// Used by the decompressor's `options=` constructor argument. - fn dparameter_from_int(param: i32, value: i32, vm: &VirtualMachine) -> PyResult { - match param { - ZSTD_d_windowLogMax => Ok(DParameter::WindowLogMax(value as u32)), - _ => Err(vm.new_value_error(format!( + /// Validate a decompression-parameter id, returning the parameter + /// constant on success. Used by the decompressor's `options=` constructor + /// argument. + fn dparameter_from_int(param: i32, vm: &VirtualMachine) -> PyResult { + d_param_enum(param).ok_or_else(|| { + vm.new_value_error(format!( "invalid decompression parameter 'unknown parameter (key {param})'" - ))), - } + )) + }) } /// Build the CPython-compatible " parameter 'name' received an @@ -312,21 +322,23 @@ mod _zstd { /// parameter id. Returns `None` if the id is not recognized (callers /// validate the id separately). fn lookup_param_bounds(param: i32, is_compress: bool) -> Option<(i32, i32)> { - // SAFETY: `ZSTD_*Param_getBounds` reads no memory beyond the enum - // discriminant; the helpers above validated that `param` maps to a - // real enum variant. - let bounds = if is_compress { + // The helpers above validated that `param` maps to a real parameter + // constant before `ZSTD_*Param_getBounds` inspects it. The two + // functions return distinct (but layout-identical) `ZSTD_bounds` + // types, so destructure into a tuple right away. + let (error, lo, hi) = if is_compress { let p = c_param_enum(param)?; - unsafe { zstd_sys::ZSTD_cParam_getBounds(p) } + let b = ZSTD_cParam_getBounds(p); + (b.error, b.lowerBound, b.upperBound) } else { let p = d_param_enum(param)?; - unsafe { zstd_sys::ZSTD_dParam_getBounds(p) } + let b = ZSTD_dParam_getBounds(p); + (b.error, b.lowerBound, b.upperBound) }; - // SAFETY: ZSTD_isError just inspects the bounds.error integer. - if unsafe { zstd_sys::ZSTD_isError(bounds.error) } != 0 { + if ZSTD_isError(error) != 0 { return None; } - Some((bounds.lowerBound, bounds.upperBound)) + Some((lo, hi)) } /// Map a parameter integer id back to the Python-visible enum member @@ -365,21 +377,21 @@ mod _zstd { } /// Map a strategy integer (as exposed via the `Strategy` IntEnum) back - /// to the underlying `ZSTD_strategy` C enum value. Done via an explicit + /// to the underlying `ZSTD_strategy` C value. Done via an explicit /// match for the same reason as [`c_param_enum`]: an untrusted int - /// might not correspond to any real enum variant. - fn strategy_from_int(v: i32) -> Option { - use zstd_sys::ZSTD_strategy as S; + /// might not correspond to any real strategy value. + fn strategy_from_int(v: i32) -> Option { + use libzstd_rs_sys::lib::zstd as s; Some(match v { - ZSTD_fast => S::ZSTD_fast, - ZSTD_dfast => S::ZSTD_dfast, - ZSTD_greedy => S::ZSTD_greedy, - ZSTD_lazy => S::ZSTD_lazy, - ZSTD_lazy2 => S::ZSTD_lazy2, - ZSTD_btlazy2 => S::ZSTD_btlazy2, - ZSTD_btopt => S::ZSTD_btopt, - ZSTD_btultra => S::ZSTD_btultra, - ZSTD_btultra2 => S::ZSTD_btultra2, + ZSTD_fast => s::ZSTD_fast, + ZSTD_dfast => s::ZSTD_dfast, + ZSTD_greedy => s::ZSTD_greedy, + ZSTD_lazy => s::ZSTD_lazy, + ZSTD_lazy2 => s::ZSTD_lazy2, + ZSTD_btlazy2 => s::ZSTD_btlazy2, + ZSTD_btopt => s::ZSTD_btopt, + ZSTD_btultra => s::ZSTD_btultra, + ZSTD_btultra2 => s::ZSTD_btultra2, _ => return None, }) } @@ -462,12 +474,16 @@ mod _zstd { fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { let dict_content = args.dict_content.with_ref(|b| b.to_vec()); - // libzstd's `get_dict_id_from_dict` returns `None` either when + // libzstd's `ZSTD_getDictID_fromDict` returns 0 either when // the content is too small to contain a valid header or when it // does not carry the dictionary magic. Both are runtime errors // when `is_raw=False`, matching CPython's behavior of raising // `ValueError` on a non-conformant dictionary. - let parsed_id = zstd_safe::get_dict_id_from_dict(&dict_content).map_or(0, |n| n.get()); + // SAFETY: `dict_content` is a live slice for the duration of the + // call. + let parsed_id = unsafe { + ZSTD_getDictID_fromDict(dict_content.as_ptr().cast(), dict_content.len()) + }; if !args.is_raw && parsed_id == 0 { return Err(vm.new_value_error( "ZSTD_DICT_MAGIC_NUMBER not found, dict_content cannot be a 'raw content' \ @@ -549,6 +565,150 @@ mod _zstd { const COMP_MODE_FLUSH_BLOCK: i32 = 1; const COMP_MODE_FLUSH_FRAME: i32 = 2; + /// Owning wrapper around a raw libzstd compression context. Frees the + /// context with `ZSTD_freeCCtx` on drop. + struct CCtx(*mut ZSTD_CCtx); + + // SAFETY: the wrapped context is only ever accessed while holding the + // owning pyclass's `PyMutex` (`CompressorState`), so no concurrent + // access can occur. + unsafe impl Send for CCtx {} + + impl CCtx { + fn create() -> Self { + // SAFETY: `ZSTD_createCCtx` has no preconditions; it returns NULL + // only on allocation failure. + let ptr = unsafe { ZSTD_createCCtx() }; + assert!(!ptr.is_null(), "ZSTD_createCCtx failed"); + Self(ptr) + } + + /// Forward a `(parameter, value)` pair to libzstd, passing the value + /// straight through like CPython does. The error code is returned on + /// failure so the caller can pick the right error mapping. + fn set_parameter(&mut self, param: ZSTD_cParameter, value: c_int) -> Result<(), usize> { + // SAFETY: `self.0` is a live context owned by this wrapper. + let code = unsafe { ZSTD_CCtx_setParameter(self.0, param, value) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } + } + + fn set_pledged_src_size(&mut self, pledged: Option) -> Result<(), usize> { + // libzstd represents "unknown" as `ZSTD_CONTENTSIZE_UNKNOWN`. + let pledged = pledged.map_or(ZSTD_CONTENTSIZE_UNKNOWN, |v| v); + // SAFETY: `self.0` is a live context owned by this wrapper. + let code = unsafe { ZSTD_CCtx_setPledgedSrcSize(self.0, pledged) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } + } + } + + impl Drop for CCtx { + fn drop(&mut self) { + // SAFETY: `self.0` was created by `ZSTD_createCCtx` and is freed + // exactly once, here. + unsafe { ZSTD_freeCCtx(self.0) }; + } + } + + /// Owning wrapper around a raw libzstd decompression context. Frees the + /// context with `ZSTD_freeDCtx` on drop. + struct DCtx(*mut ZSTD_DCtx); + + // SAFETY: same reasoning as `CCtx`: access is serialized by the owning + // pyclass's `PyMutex` (`DecompressorState`). + unsafe impl Send for DCtx {} + + impl DCtx { + fn create() -> Self { + // SAFETY: `ZSTD_createDCtx` has no preconditions; it returns NULL + // only on allocation failure. + let ptr = unsafe { ZSTD_createDCtx() }; + assert!(!ptr.is_null(), "ZSTD_createDCtx failed"); + Self(ptr) + } + + /// See [`CCtx::set_parameter`]. + fn set_parameter(&mut self, param: ZSTD_dParameter, value: c_int) -> Result<(), usize> { + // SAFETY: `self.0` is a live context owned by this wrapper. + let code = unsafe { ZSTD_DCtx_setParameter(self.0, param, value) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } + } + } + + impl Drop for DCtx { + fn drop(&mut self) { + // SAFETY: `self.0` was created by `ZSTD_createDCtx` and is freed + // exactly once, here. + unsafe { ZSTD_freeDCtx(self.0) }; + } + } + + /// Owning wrapper around a digested compression dictionary. Frees the + /// dictionary with `ZSTD_freeCDict` on drop. + struct CDict(*mut ZSTD_CDict); + + // SAFETY: the dictionary is only referenced by the `CCtx` stored + // alongside it in `CompressorState`, which the pyclass's `PyMutex` + // serializes access to. + unsafe impl Send for CDict {} + + impl CDict { + /// Build a digested dictionary; libzstd copies the bytes into its own + /// storage. Returns `None` when libzstd rejects the content (NULL), + /// e.g. for a corrupted dictionary. + fn try_create(bytes: &[u8], level: c_int) -> Option { + // SAFETY: `bytes` is a live slice for the duration of the call, + // and `ZSTD_createCDict` copies what it needs. + let ptr = unsafe { ZSTD_createCDict(bytes.as_ptr().cast(), bytes.len(), level) }; + if ptr.is_null() { None } else { Some(Self(ptr)) } + } + } + + impl Drop for CDict { + fn drop(&mut self) { + // SAFETY: `self.0` was created by `ZSTD_createCDict` and is freed + // exactly once, here. + unsafe { ZSTD_freeCDict(self.0) }; + } + } + + /// Owning wrapper around a digested decompression dictionary. Frees the + /// dictionary with `ZSTD_freeDDict` on drop. + struct DDict(*mut ZSTD_DDict); + + // SAFETY: same reasoning as `CDict` (guarded by the decompressor's + // `PyMutex` via `DecompressorState`). + unsafe impl Send for DDict {} + + impl DDict { + /// See [`CDict::try_create`]. + fn try_create(bytes: &[u8]) -> Option { + // SAFETY: `bytes` is a live slice for the duration of the call, + // and `ZSTD_createDDict` copies what it needs. + let ptr = unsafe { ZSTD_createDDict(bytes.as_ptr().cast(), bytes.len()) }; + if ptr.is_null() { None } else { Some(Self(ptr)) } + } + } + + impl Drop for DDict { + fn drop(&mut self) { + // SAFETY: `self.0` was created by `ZSTD_createDDict` and is freed + // exactly once, here. + unsafe { ZSTD_freeDDict(self.0) }; + } + } + /// Internal state of a `ZstdCompressor`. Holds the libzstd context, the /// last mode used (for `last_mode` and `set_pledged_input_size` /// validation), and the dictionary handles that the context may reference @@ -557,12 +717,12 @@ mod _zstd { /// `PyRef` go away afterwards, which is the safe order for /// teardown. struct CompressorState { - cctx: CCtx<'static>, + cctx: CCtx, /// Cached digested dictionary. The CCtx references this via - /// `ref_cdict`, so it must outlive the CCtx (handled by Rust's - /// field drop order: `cctx` drops first). - _cdict: Option>, - /// Keeps the ZstdDict's bytes alive for `ref_prefix` mode. + /// `ZSTD_CCtx_refCDict`, so it must outlive the CCtx (handled by + /// Rust's field drop order: `cctx` drops first). + _cdict: Option, + /// Keeps the ZstdDict's bytes alive for `ZSTD_CCtx_refPrefix` mode. _dict: Option>, last_mode: i32, } @@ -592,14 +752,11 @@ mod _zstd { /// Translate the public `mode` int to the libzstd `ZSTD_EndDirective` /// the streaming API takes. - fn end_directive_from_mode( - mode: i32, - vm: &VirtualMachine, - ) -> PyResult { + fn end_directive_from_mode(mode: i32, vm: &VirtualMachine) -> PyResult { match mode { - COMP_MODE_CONTINUE => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_continue), - COMP_MODE_FLUSH_BLOCK => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_flush), - COMP_MODE_FLUSH_FRAME => Ok(zstd_sys::ZSTD_EndDirective::ZSTD_e_end), + COMP_MODE_CONTINUE => Ok(ZSTD_e_continue), + COMP_MODE_FLUSH_BLOCK => Ok(ZSTD_e_flush), + COMP_MODE_FLUSH_FRAME => Ok(ZSTD_e_end), _ => Err(vm.new_value_error(format!( "mode argument wrong value, it should be one of \ ZstdCompressor.CONTINUE ({COMP_MODE_CONTINUE}), \ @@ -621,11 +778,11 @@ mod _zstd { return Err(vm.new_type_error("Only one of level or options should be used.")); } - let mut cctx = CCtx::<'static>::create(); + let mut cctx = CCtx::create(); if let Some(level_obj) = level_opt { let level = parse_compression_level(&level_obj, vm)?; - cctx.set_parameter(CParameter::CompressionLevel(level)) + cctx.set_parameter(ZSTD_cParameter::ZSTD_c_compressionLevel, level) .map_err(|_| param_value_error_for(ZSTD_c_compressionLevel, level, true, vm))?; } @@ -708,25 +865,25 @@ mod _zstd { /// Trait wrapper over `CCtx::set_parameter` and `DCtx::set_parameter` so /// `apply_options` can drive either context without duplicated code. - /// `set_parameter` translates the (id, value) pair to the appropriate - /// `CParameter` / `DParameter` enum variant, then forwards to libzstd. + /// `apply` validates the (id, value) pair via `cparameter_from_int` / + /// `dparameter_from_int`, then forwards to libzstd. trait ParamSetter { fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()>; } - impl ParamSetter for CCtx<'static> { + impl ParamSetter for CCtx { fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()> { let p = cparameter_from_int(param, value, vm)?; - self.set_parameter(p) + self.set_parameter(p, value) .map_err(|_| param_value_error_for(param, value, true, vm))?; Ok(()) } } - impl ParamSetter for DCtx<'static> { + impl ParamSetter for DCtx { fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()> { - let p = dparameter_from_int(param, value, vm)?; - self.set_parameter(p) + let p = dparameter_from_int(param, vm)?; + self.set_parameter(p, value) .map_err(|_| param_value_error_for(param, value, false, vm))?; Ok(()) } @@ -736,46 +893,95 @@ mod _zstd { /// and decompressor consume a dictionary: the name of the type that /// appears in error messages, the eager-validation constructor for the /// digested variant, and the three ways of attaching it to the context. - trait DictLoader<'a> { + /// The `Err` payload of each method is a libzstd error code. + trait DictLoader { type Digested; const KIND_NAME: &'static str; fn try_create_digested(bytes: &[u8]) -> Option; - fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult; - fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult; - fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult; + fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize>; + fn load_undigested(&mut self, bytes: &[u8]) -> Result<(), usize>; + fn ref_prefix(&mut self, bytes: &[u8]) -> Result<(), usize>; } - impl DictLoader<'static> for CCtx<'static> { - type Digested = zstd_safe::CDict<'static>; + impl DictLoader for CCtx { + type Digested = CDict; const KIND_NAME: &'static str = "ZSTD_CDict"; fn try_create_digested(bytes: &[u8]) -> Option { - zstd_safe::CDict::try_create(bytes, ZSTD_CLEVEL_DEFAULT) + CDict::try_create(bytes, libzstd_rs_sys::ZSTD_CLEVEL_DEFAULT) } - fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult { - self.ref_cdict(dict) + fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize> { + // SAFETY: both handles are live; `load_dict`'s safety contract + // keeps `dict` alive at least as long as `self`. + let code = unsafe { ZSTD_CCtx_refCDict(self.0, dict.0.cast_const()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } - fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult { - self.load_dictionary(bytes) + fn load_undigested(&mut self, bytes: &[u8]) -> Result<(), usize> { + // SAFETY: `bytes` is a live slice for the duration of the call; + // its contents are copied into the context. + let code = + unsafe { ZSTD_CCtx_loadDictionary(self.0, bytes.as_ptr().cast(), bytes.len()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } - fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult { - self.ref_prefix(bytes) + fn ref_prefix(&mut self, bytes: &[u8]) -> Result<(), usize> { + // SAFETY: `bytes` is a live slice for the duration of the call. + // libzstd keeps a raw pointer to the bytes afterwards; + // `load_dict`'s safety contract keeps them alive at least as + // long as `self`. + let code = unsafe { ZSTD_CCtx_refPrefix(self.0, bytes.as_ptr().cast(), bytes.len()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } } - impl DictLoader<'static> for DCtx<'static> { - type Digested = zstd_safe::DDict<'static>; + impl DictLoader for DCtx { + type Digested = DDict; const KIND_NAME: &'static str = "ZSTD_DDict"; fn try_create_digested(bytes: &[u8]) -> Option { - zstd_safe::DDict::try_create(bytes) + DDict::try_create(bytes) } - fn ref_digested(&mut self, dict: &Self::Digested) -> zstd_safe::SafeResult { - self.ref_ddict(dict) + fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize> { + // SAFETY: both handles are live; `load_dict`'s safety contract + // keeps `dict` alive at least as long as `self`. + let code = unsafe { ZSTD_DCtx_refDDict(self.0, dict.0.cast_const()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } - fn load_undigested(&mut self, bytes: &[u8]) -> zstd_safe::SafeResult { - self.load_dictionary(bytes) + fn load_undigested(&mut self, bytes: &[u8]) -> Result<(), usize> { + // SAFETY: `bytes` is a live slice for the duration of the call; + // its contents are copied into the context. + let code = + unsafe { ZSTD_DCtx_loadDictionary(self.0, bytes.as_ptr().cast(), bytes.len()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } - fn ref_prefix_static(&mut self, bytes: &'static [u8]) -> zstd_safe::SafeResult { - self.ref_prefix(bytes) + fn ref_prefix(&mut self, bytes: &[u8]) -> Result<(), usize> { + // SAFETY: `bytes` is a live slice for the duration of the call. + // libzstd keeps a raw pointer to the bytes afterwards; + // `load_dict`'s safety contract keeps them alive at least as + // long as `self`. + let code = unsafe { ZSTD_DCtx_refPrefix(self.0, bytes.as_ptr().cast(), bytes.len()) }; + if ZSTD_isError(code) != 0 { + Err(code) + } else { + Ok(()) + } } } @@ -803,7 +1009,7 @@ mod _zstd { /// In `undigested` mode the bytes are copied into `ctx`, so neither /// return value carries a safety obligation — but the caller cannot tell /// the modes apart, so it must keep both alive regardless. - unsafe fn load_dict>( + unsafe fn load_dict( ctx: &mut L, dict_obj: Option, vm: &VirtualMachine, @@ -825,17 +1031,7 @@ mod _zstd { let mut digested = None; match marker { DICT_TYPE_PREFIX => { - // SAFETY: we extend `dict_bytes`' lifetime to `'static` only - // to thread it through `ref_prefix_static`'s signature. - // `load_dict`'s own safety contract requires the caller to - // keep the returned `PyRef` (which owns - // `dict_bytes`) alive at least as long as `ctx`, so the - // raw pointer libzstd stores stays valid for the required - // window. - let static_bytes: &'static [u8] = - unsafe { core::slice::from_raw_parts(dict_bytes.as_ptr(), dict_bytes.len()) }; - ctx.ref_prefix_static(static_bytes) - .map_err(|_| bad_dict_err())?; + ctx.ref_prefix(dict_bytes).map_err(|_| bad_dict_err())?; } DICT_TYPE_DIGESTED => { // Build the digested dict eagerly so a corrupted dictionary @@ -861,7 +1057,7 @@ mod _zstd { /// here, both invariants `load_dict` documents become structural and a /// safe-Rust caller cannot split the pieces apart. fn build_compressor_state( - mut cctx: CCtx<'static>, + mut cctx: CCtx, dict_obj: Option, last_mode: i32, vm: &VirtualMachine, @@ -874,7 +1070,7 @@ mod _zstd { // and `_dict` (prefix mode) are freed. `CompressorState` is private // to this module and is never destructured, so no safe caller can // reorder the drops. - let (cdict, dict) = unsafe { load_dict::>(&mut cctx, dict_obj, vm) }?; + let (cdict, dict) = unsafe { load_dict::(&mut cctx, dict_obj, vm) }?; Ok(CompressorState { cctx, _cdict: cdict, @@ -887,12 +1083,12 @@ mod _zstd { /// [`build_compressor_state`] for the safety reasoning; /// `DecompressorState`'s field order plays the same role here. fn build_decompressor_state( - mut dctx: DCtx<'static>, + mut dctx: DCtx, dict_obj: Option, vm: &VirtualMachine, ) -> PyResult { // SAFETY: see [`build_compressor_state`]. - let (ddict, dict) = unsafe { load_dict::>(&mut dctx, dict_obj, vm) }?; + let (ddict, dict) = unsafe { load_dict::(&mut dctx, dict_obj, vm) }?; Ok(DecompressorState { dctx, _ddict: ddict, @@ -904,13 +1100,14 @@ mod _zstd { }) } - /// Drive `compress_stream2` until the input is fully consumed and, for - /// flush/end directives, the internal buffers report zero remaining bytes. - /// Grows the output `Vec` by `CStreamOutSize` chunks as needed. + /// Drive `ZSTD_compressStream2` until the input is fully consumed and, + /// for flush/end directives, the internal buffers report zero remaining + /// bytes. Grows the output `Vec` by `ZSTD_CStreamOutSize` chunks as + /// needed. fn do_compress( state: &mut CompressorState, data: &[u8], - end_op: zstd_sys::ZSTD_EndDirective, + end_op: ZSTD_EndDirective, vm: &VirtualMachine, ) -> PyResult> { // Release the GIL for the duration of the compression loop. Safety: @@ -921,21 +1118,41 @@ mod _zstd { // and the output `Vec` is local to this function. No Python object // access happens inside the closure — error codes are surfaced as // `usize` and converted into exceptions after re-attaching. - let is_end = end_op != zstd_sys::ZSTD_EndDirective::ZSTD_e_continue; - let chunk_size = CCtx::out_size().max(1); + let is_end = end_op != ZSTD_e_continue; + let chunk_size = ZSTD_CStreamOutSize().max(1); let result: Result, usize> = vm.allow_threads(|| { - let mut output = Vec::new(); - let mut input = InBuffer::around(data); + let mut output: Vec = Vec::new(); + let mut input = ZSTD_inBuffer { + src: data.as_ptr().cast(), + size: data.len(), + pos: 0, + }; loop { let prev_len = output.len(); output.reserve(chunk_size); let remaining = { - let mut out_buf = OutBuffer::around_pos(&mut output, prev_len); - state - .cctx - .compress_stream2(&mut out_buf, &mut input, end_op) - }?; - let consumed_all = input.pos == input.src.len(); + let spare = output.spare_capacity_mut(); + let mut out_buf = ZSTD_outBuffer { + dst: spare.as_mut_ptr().cast(), + size: spare.len(), + pos: 0, + }; + // SAFETY: `state.cctx` is a live context; `out_buf` + // covers the Vec's spare capacity, which is valid for + // `spare.len()` bytes; `input` points to `data`, which + // outlives the call. + let code = unsafe { + ZSTD_compressStream2(state.cctx.0, &mut out_buf, &mut input, end_op) + }; + if ZSTD_isError(code) != 0 { + return Err(code); + } + // SAFETY: libzstd wrote `out_buf.pos` bytes + // (`out_buf.pos <= spare.len()`) into the spare capacity. + unsafe { output.set_len(prev_len + out_buf.pos) }; + code + }; + let consumed_all = input.pos == input.size; // Stop when input is fully consumed and, for flush/end // directives, libzstd reports that all internal buffers have // been drained (remaining == 0). Otherwise loop; the next @@ -986,8 +1203,8 @@ mod _zstd { // likewise converts the argument before touching the compressor. // // Python passes `None` to mean "unknown"; libzstd represents that - // internally as `ZSTD_CONTENTSIZE_UNKNOWN` (`u64::MAX`), and - // `zstd_safe` translates `None` accordingly. libzstd also reserves + // internally as `ZSTD_CONTENTSIZE_UNKNOWN` (`u64::MAX`), which + // `set_pledged_src_size` maps `None` to. libzstd also reserves // `ZSTD_CONTENTSIZE_ERROR` (`u64::MAX - 1`), so a concrete size must // be strictly less than that; reject anything else up front so // callers see the documented `ValueError`, not a libzstd-level error. @@ -1072,9 +1289,10 @@ mod _zstd { /// the `dctx` is freed first and must give up its internal pointers /// before any referenced `DDict`/`PyRef` is dropped. struct DecompressorState { - dctx: DCtx<'static>, - /// Cached decompression dictionary referenced by the DCtx. - _ddict: Option>, + dctx: DCtx, + /// Cached decompression dictionary referenced by the DCtx via + /// `ZSTD_DCtx_refDDict`. + _ddict: Option, _dict: Option>, eof: bool, needs_input: bool, @@ -1113,7 +1331,7 @@ mod _zstd { let dict_opt = args.zstd_dict.flatten(); let options_opt = args.options.flatten(); - let mut dctx = DCtx::<'static>::create(); + let mut dctx = DCtx::create(); if let Some(options_obj) = options_opt { apply_options(&mut dctx, options_obj, false, vm)?; @@ -1162,18 +1380,22 @@ mod _zstd { alloc::borrow::Cow::Owned(combined) }; - let chunk_size = DCtx::out_size().max(1); + let chunk_size = ZSTD_DStreamOutSize().max(1); // Release the GIL for the streaming loop. Safety: see `do_compress`; // the closure captures only Rust-owned buffers and `&mut state` (held // under the decompressor's `PyMutex`), and surfaces error codes as // `usize` so we can build the exception after re-attaching. let loop_result: Result<(Vec, bool, usize), usize> = vm.allow_threads(|| { - let mut input = InBuffer::around(&work_data); + let mut input = ZSTD_inBuffer { + src: work_data.as_ptr().cast(), + size: work_data.len(), + pos: 0, + }; let mut output: Vec = Vec::new(); // Reusable scratch buffer for each decompress_stream call. We need // an exact-size output buffer because `Vec::reserve` may - // over-allocate; `OutBuffer` reports the full Vec capacity to - // libzstd, which would then happily write past `max_length`. + // over-allocate; reporting the full Vec capacity to libzstd would + // let it write past `max_length`. let mut scratch: Vec = vec![0u8; chunk_size]; let mut hit_max = false; let mut iteration = 0usize; @@ -1195,41 +1417,46 @@ mod _zstd { Some(maxl) => (maxl - output.len()).min(chunk_size), None => chunk_size, }; - let result; + let code; let written; { let slot = &mut scratch[..grow]; - let mut out_buf = OutBuffer::around(slot as &mut [u8]); - result = state.dctx.decompress_stream(&mut out_buf, &mut input); - written = out_buf.pos(); + let mut out_buf = ZSTD_outBuffer { + dst: slot.as_mut_ptr().cast(), + size: slot.len(), + pos: 0, + }; + // SAFETY: `state.dctx` is a live context; `out_buf` + // covers `slot`; `input` points into `work_data`, which + // outlives the call. + code = unsafe { ZSTD_decompressStream(state.dctx.0, &mut out_buf, &mut input) }; + written = out_buf.pos; } output.extend_from_slice(&scratch[..written]); - match result { - Ok(0) => { - // Frame fully decompressed; the decompressor is at EOF. - state.eof = true; - break Ok(()); - } - Ok(_) => { - let output_was_full = written == grow; - let input_consumed = input.pos == input.src.len(); - - if let Some(maxl) = max_length - && output.len() >= maxl - && iteration > 1 - { - hit_max = true; - break Ok(()); - } - - // Input is gone and libzstd had room to write but did - // not, which means the frame is incomplete and the - // caller has to supply more input. - if input_consumed && !output_was_full { - break Ok(()); - } - } - Err(code) => break Err(code), + if ZSTD_isError(code) != 0 { + break Err(code); + } + if code == 0 { + // Frame fully decompressed; the decompressor is at EOF. + state.eof = true; + break Ok(()); + } + let output_was_full = written == grow; + let input_consumed = input.pos == input.size; + + if let Some(maxl) = max_length + && output.len() >= maxl + && iteration > 1 + { + hit_max = true; + break Ok(()); + } + + // Input is gone and libzstd had room to write but did + // not, which means the frame is incomplete and the + // caller has to supply more input. + if input_consumed && !output_was_full { + break Ok(()); } }; outcome.map(|()| (output, hit_max, input.pos)) @@ -1318,14 +1545,17 @@ mod _zstd { #[pyfunction] fn get_frame_size(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { let buf = frame_buffer.with_ref(|b| b.to_vec()); - zstd_safe::find_frame_compressed_size(&buf).map_err(|_| { - new_zstd_error( + // SAFETY: `buf` is a live slice for the duration of the call. + let size = unsafe { ZSTD_findFrameCompressedSize(buf.as_ptr().cast(), buf.len()) }; + if ZSTD_isError(size) != 0 { + return Err(new_zstd_error( "Error when finding the compressed size of a Zstandard frame. \ Ensure the frame_buffer argument starts from the beginning of a frame, \ and its length not less than this complete frame.", vm, - ) - }) + )); + } + Ok(size) } #[pyfunction] @@ -1334,19 +1564,23 @@ mod _zstd { vm: &VirtualMachine, ) -> PyResult<(PyObjectRef, u32)> { let buf = frame_buffer.with_ref(|b| b.to_vec()); - let content_size = zstd_safe::get_frame_content_size(&buf).map_err(|_| { - new_zstd_error( + // SAFETY: `buf` is a live slice for the duration of the call. + let content_size = unsafe { ZSTD_getFrameContentSize(buf.as_ptr().cast(), buf.len()) }; + if content_size == ZSTD_CONTENTSIZE_ERROR { + return Err(new_zstd_error( "Error when getting information from the header of a Zstandard frame. \ Ensure the frame_buffer argument starts from the beginning of a frame, \ and its length not less than the frame header (6~18 bytes).", vm, - ) - })?; - let content_size_obj: PyObjectRef = match content_size { - Some(n) => vm.ctx.new_int(n).into(), - None => vm.ctx.none(), + )); + } + let content_size_obj: PyObjectRef = if content_size == ZSTD_CONTENTSIZE_UNKNOWN { + vm.ctx.none() + } else { + vm.ctx.new_int(content_size).into() }; - let dict_id = zstd_safe::get_dict_id_from_frame(&buf).map_or(0, |n| n.get()); + // SAFETY: `buf` is a live slice for the duration of the call. + let dict_id = unsafe { ZSTD_getDictID_fromFrame(buf.as_ptr().cast(), buf.len()) }; Ok((content_size_obj, dict_id)) } @@ -1405,9 +1639,9 @@ mod _zstd { } /// Sum the per-sample sizes and check they exactly cover `expected_total`, - /// rejecting overflow. A safe-fn interface that calls into libzstd with - /// these sizes must not let a wrapping sum sneak past the equality check, - /// since libzstd would then read past the samples buffer. + /// rejecting overflow. Any call into libzstd with these sizes must not + /// let a wrapping sum sneak past the equality check, since libzstd would + /// then read past the samples buffer. fn check_sample_sizes_match( sizes: &[usize], expected_total: usize, @@ -1433,8 +1667,27 @@ mod _zstd { let sizes = parse_sample_sizes(args.samples_sizes, vm)?; check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; let mut dict_buffer: Vec = Vec::with_capacity(dict_size); - zstd_safe::train_from_buffer(&mut dict_buffer, &samples_buffer, &sizes) - .map_err(|c| catch_zstd_error(c, vm))?; + // SAFETY: `dict_buffer`'s spare capacity is valid for writes of up to + // `dict_size` bytes; the samples and sizes buffers are live slices + // for the duration of the call. + let written = unsafe { + ZDICT_trainFromBuffer( + dict_buffer.as_mut_ptr().cast(), + dict_size, + samples_buffer.as_ptr().cast(), + sizes.as_ptr(), + sizes.len() as u32, + ) + }; + if ZDICT_isError(written) != 0 { + // SAFETY: `ZDICT_getErrorName` returns a pointer to a static, + // NUL-terminated error string from libzstd's error table. + let name = unsafe { CStr::from_ptr(ZDICT_getErrorName(written)) }; + return Err(new_zstd_error(name.to_string_lossy(), vm)); + } + // SAFETY: `ZDICT_trainFromBuffer` wrote `written` (<= dict_size) + // bytes into the buffer's spare capacity. + unsafe { dict_buffer.set_len(written) }; Ok(dict_buffer) } @@ -1459,8 +1712,6 @@ mod _zstd { compression_level: PyObjectRef, } - // `zstd_safe` does not wrap `ZDICT_finalizeDictionary`, so we drop down to - // raw `zstd_sys` FFI for it. #[pyfunction] fn finalize_dict(args: FinalizeDictArgs, vm: &VirtualMachine) -> PyResult> { let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; @@ -1471,7 +1722,7 @@ mod _zstd { check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; let mut dict_buffer: Vec = vec![0u8; dict_size]; - let params = zstd_sys::ZDICT_params_t { + let params = ZDICT_params_t { compressionLevel: compression_level, notificationLevel: 0, dictID: 0, @@ -1481,7 +1732,7 @@ mod _zstd { // that outlive the FFI call. ZDICT_finalizeDictionary just reads from // the sample/dict buffers and writes into `dict_buffer`. let written = unsafe { - zstd_sys::ZDICT_finalizeDictionary( + ZDICT_finalizeDictionary( dict_buffer.as_mut_ptr() as *mut _, dict_buffer.len(), custom_dict.as_ptr() as *const _, @@ -1492,15 +1743,14 @@ mod _zstd { params, ) }; - // SAFETY: ZDICT_isError just inspects the integer return code. - if unsafe { zstd_sys::ZDICT_isError(written) } != 0 { - // SAFETY: ZDICT_getErrorName returns a static NUL-terminated - // C string from libzstd's internal error table. - let err_ptr = unsafe { zstd_sys::ZDICT_getErrorName(written) }; + if ZDICT_isError(written) != 0 { + let err_ptr = ZDICT_getErrorName(written); let msg = if err_ptr.is_null() { "zstd dictionary finalization failed".to_string() } else { - unsafe { core::ffi::CStr::from_ptr(err_ptr) } + // SAFETY: `ZDICT_getErrorName` returns a pointer to a static, + // NUL-terminated error string from libzstd's error table. + unsafe { CStr::from_ptr(err_ptr) } .to_string_lossy() .into_owned() }; @@ -1533,22 +1783,22 @@ mod _zstd { }; // Validate the id via the same safe enum-lookup helpers used in // `lookup_param_bounds`, then call libzstd directly so we can - // distinguish a libzstd-reported error from our own "unknown". - let bounds = if args.is_compress { + // distinguish a libzstd-reported error from our own "unknown". The + // two `getBounds` functions return distinct (but layout-identical) + // `ZSTD_bounds` types, so destructure into a tuple right away. + let (error, lo, hi) = if args.is_compress { let p = c_param_enum(args.parameter).ok_or_else(unknown)?; - // SAFETY: `c_param_enum` returned `Some`, so `p` is a real - // `ZSTD_cParameter` discriminant. - unsafe { zstd_sys::ZSTD_cParam_getBounds(p) } + let b = ZSTD_cParam_getBounds(p); + (b.error, b.lowerBound, b.upperBound) } else { let p = d_param_enum(args.parameter).ok_or_else(unknown)?; - // SAFETY: same as above. - unsafe { zstd_sys::ZSTD_dParam_getBounds(p) } + let b = ZSTD_dParam_getBounds(p); + (b.error, b.lowerBound, b.upperBound) }; - // SAFETY: ZSTD_isError just inspects the integer error code. - if unsafe { zstd_sys::ZSTD_isError(bounds.error) } != 0 { - return Err(catch_zstd_error(bounds.error, vm)); + if ZSTD_isError(error) != 0 { + return Err(catch_zstd_error(error, vm)); } - Ok((bounds.lowerBound, bounds.upperBound)) + Ok((lo, hi)) } // Register the `CompressionParameter` / `DecompressionParameter` enum From e84dabab673a66a6bb95396b571e5fdc036ffdf7 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sun, 9 Aug 2026 17:52:07 +0100 Subject: [PATCH 14/25] Fix data loss in ZstdDecompressor.decompress(max_length=0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The zero-cap probe iteration handed libzstd a 1-byte buffer and the post-loop truncation discarded whatever it emitted, while the input counter had already advanced past the bytes that produced it — every decompress(..., max_length=0) call with pending output silently dropped one byte. The suite never caught it because max_length=0 was only exercised against zero-output (skippable) frames. Match CPython's mechanism instead: probe with a zero-size output buffer, so libzstd consumes input without emitting and never reports frame completion with un-emitted content. Add a regression snippet covering lossless zero-cap probes (content, truncated, and skippable frames), verified byte-exact against CPython. --- crates/stdlib/src/zstd.rs | 31 +++++++----------- extra_tests/snippets/stdlib_zstd.py | 51 +++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 19 deletions(-) create mode 100644 extra_tests/snippets/stdlib_zstd.py diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 8d9e194bfd6..e7d3613fef0 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -1403,17 +1403,18 @@ mod _zstd { let outcome = loop { iteration += 1; // Honor `max_length`: stop growing the output buffer once - // we have produced enough. Special-case the first iteration - // when the cap is zero so a zero-output frame (skippable - // frame, empty content frame) can still complete; we hand - // libzstd a 1-byte scratch and discard the byte if it ends - // up writing one. + // we have produced enough. When the cap is zero, hand + // libzstd a zero-size buffer instead: it then consumes + // input without emitting, so zero-output frames (skippable + // frame, empty content frame) still complete while real + // output stays inside libzstd until a later call has room. + // CPython's decompressor uses the same zero-size mechanism. let grow = match max_length { Some(maxl) if output.len() >= maxl && iteration > 1 => { hit_max = true; break Ok(()); } - Some(maxl) if output.len() >= maxl => 1, + Some(maxl) if output.len() >= maxl => 0, Some(maxl) => (maxl - output.len()).min(chunk_size), None => chunk_size, }; @@ -1441,7 +1442,10 @@ mod _zstd { state.eof = true; break Ok(()); } - let output_was_full = written == grow; + // Only meaningful with a non-zero buffer: libzstd had more + // to emit but ran out of room. A zero-size call (the + // max_length == 0 probe) is not "full". + let output_was_full = grow > 0 && written == grow; let input_consumed = input.pos == input.size; if let Some(maxl) = max_length @@ -1462,18 +1466,7 @@ mod _zstd { outcome.map(|()| (output, hit_max, input.pos)) }); - let (mut output, mut hit_max, consumed) = - loop_result.map_err(|c| catch_zstd_error(c, vm))?; - - // If `max_length == 0` opened a courtesy iteration that produced more - // bytes than the caller asked for, truncate. Should not happen with - // the scratch slicing above, but keep the safety net. - if let Some(maxl) = max_length - && output.len() > maxl - { - output.truncate(maxl); - hit_max = true; - } + let (output, hit_max, consumed) = loop_result.map_err(|c| catch_zstd_error(c, vm))?; let remaining = &work_data[consumed..]; diff --git a/extra_tests/snippets/stdlib_zstd.py b/extra_tests/snippets/stdlib_zstd.py new file mode 100644 index 00000000000..58fa742a3d7 --- /dev/null +++ b/extra_tests/snippets/stdlib_zstd.py @@ -0,0 +1,51 @@ +from compression.zstd import ZstdCompressor, ZstdDecompressor + +# basic roundtrip +payload = b"hello world" * 100 +c = ZstdCompressor() +frame = c.compress(payload) + c.flush() +d = ZstdDecompressor() +out = d.decompress(frame) +assert out == payload +assert d.eof +assert not d.needs_input + +# max_length=0 must not drop output bytes: probing with a zero cap has to +# leave every byte available to later drain calls (regression: the probe +# used to emit one byte per call into the void). +d = ZstdDecompressor() +out = d.decompress(frame, 0) +assert out == b"" +assert not d.eof +assert not d.needs_input +rest = d.decompress(b"", len(payload)) +assert rest == payload +assert d.eof + +# two zero-cap probes in a row, then drain; still lossless +d = ZstdDecompressor() +assert d.decompress(frame, 0) == b"" +assert d.decompress(b"", 0) == b"" +assert d.decompress(b"", len(payload)) == payload +assert d.eof + +# zero-cap probe on a truncated frame: no output, not at end +d = ZstdDecompressor() +out = d.decompress(frame[:-5], 0) +assert out == b"" +assert not d.eof +assert not d.needs_input + +# zero-cap probe on a skippable frame completes it +skippable = b"\x50\x2a\x4d\x18\x04\x00\x00\x00abcd" +d = ZstdDecompressor() +assert d.decompress(skippable, 0) == b"" +assert d.eof +assert not d.needs_input +assert d.unused_data == b"" + +# skippable frame followed by trailing bytes lands in unused_data +d = ZstdDecompressor() +assert d.decompress(skippable + b"trail", 0) == b"" +assert d.eof +assert d.unused_data == b"trail" From 2466964a650dade0c317d6783617e7817e524339 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Sun, 9 Aug 2026 18:03:50 +0100 Subject: [PATCH 15/25] Drop duplicate spell-checker line in zstd.rs --- crates/stdlib/src/zstd.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index e7d3613fef0..69e952e73ce 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -4,7 +4,6 @@ // spell-checker:ignore srcsize zdict refprefix refcdict refddict pledgedsrcsize getframecontentsize // spell-checker:ignore Zstd Zstandard pylib RFC // spell-checker:ignore CLEVEL zstdmodule cparameter dparameter maxl c2rust -// spell-checker:ignore cctx dctx CCTX DCTX ldm cdict ddict windowlog hashlog chainlog searchlog CLEVEL //! The `_zstd` extension module. Backs the pure-Python `compression.zstd` //! package by exposing the same classes, functions and constants that From 3372949d8fbd1398253da95f1b4a9a49ba0e07aa Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 19 Aug 2026 09:58:51 +0100 Subject: [PATCH 16/25] Address review on _zstd constants, version attrs and dict markers - Derive the compression/decompression parameter ids from libzstd's own `ZSTD_cParameter`/`ZSTD_dParameter` constants instead of restating them as literals. The crate declares both as `#[repr(transparent)]` newtypes over a private `u32` with no accessor, `From` impl or `Deref`, so the id is read back through that guaranteed layout. - Drop the redundant `name = ` on `#[pyattr(once)] fn zstd_version` and `fn zstd_version_number`; the function names already match. - Use `CStr::to_str` for the version string. libzstd returns a fixed ASCII `MAJOR.MINOR.RELEASE` literal, so there is nothing for a lossy conversion to repair. - Replace the `DICT_TYPE_*` int constants with a `DictType` enum carrying the pinned discriminants, decoded at the Python boundary by `DictType::from_marker`. - Explain the non-trivial assertions: why `level_bounds`' `expect` cannot fire, and why a NULL context allocation panics rather than raising `MemoryError`. Co-Authored-By: Claude Opus 5 --- crates/stdlib/src/zstd.rs | 184 ++++++++++++++++++++++++++------------ 1 file changed, 127 insertions(+), 57 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 69e952e73ce..833e4da5cea 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -54,55 +54,79 @@ mod _zstd { #[pyattr] const ZSTD_CLEVEL_DEFAULT: i32 = libzstd_rs_sys::ZSTD_CLEVEL_DEFAULT; - // Compression parameter identifiers. Values match the `ZSTD_cParameter` - // enum in libzstd, which is what the public `CompressionParameter` IntEnum - // in `Lib/compression/zstd/__init__.py` derives its members from. - // libzstd-rs-sys models the enum as a newtype whose inner value is - // crate-private, so the ids — frozen C ABI values from zstd.h — are - // spelled out as literals here. + /// Read the numeric id back out of a libzstd compression-parameter + /// constant. + /// + /// The ids are needed as plain ints because the module constants below are + /// what `Lib/compression/zstd/__init__.py` turns into the public + /// `CompressionParameter` IntEnum. libzstd-rs-sys declares + /// `ZSTD_cParameter` as `#[repr(transparent)] pub struct + /// ZSTD_cParameter(pub(crate) u32)` and gives it no accessor, `From` impl + /// or `Deref`, so the id is unreachable through its public API. Reading + /// the wrapped `u32` back keeps libzstd the single source of truth for + /// these values instead of restating them here as literals. + const fn c_param_id(param: ZSTD_cParameter) -> i32 { + // SAFETY: `#[repr(transparent)]` guarantees `ZSTD_cParameter` has the + // same size, alignment and validity invariants as the `u32` it wraps, + // so reading one back as a `u32` is always well-defined. + unsafe { core::mem::transmute::(param) as i32 } + } + + /// `ZSTD_dParameter` counterpart of [`c_param_id`]; that type is likewise + /// a `#[repr(transparent)]` newtype over a private `u32`. + const fn d_param_id(param: ZSTD_dParameter) -> i32 { + // SAFETY: as in `c_param_id`. + unsafe { core::mem::transmute::(param) as i32 } + } + + // Compression parameter identifiers, taken straight from libzstd's + // `ZSTD_cParameter` constants. These are the ids the public + // `CompressionParameter` IntEnum in `Lib/compression/zstd/__init__.py` + // derives its members from. #[pyattr] - const ZSTD_c_compressionLevel: i32 = 100; + const ZSTD_c_compressionLevel: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_compressionLevel); #[pyattr] - const ZSTD_c_windowLog: i32 = 101; + const ZSTD_c_windowLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_windowLog); #[pyattr] - const ZSTD_c_hashLog: i32 = 102; + const ZSTD_c_hashLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_hashLog); #[pyattr] - const ZSTD_c_chainLog: i32 = 103; + const ZSTD_c_chainLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_chainLog); #[pyattr] - const ZSTD_c_searchLog: i32 = 104; + const ZSTD_c_searchLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_searchLog); #[pyattr] - const ZSTD_c_minMatch: i32 = 105; + const ZSTD_c_minMatch: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_minMatch); #[pyattr] - const ZSTD_c_targetLength: i32 = 106; + const ZSTD_c_targetLength: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_targetLength); #[pyattr] - const ZSTD_c_strategy: i32 = 107; + const ZSTD_c_strategy: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_strategy); #[pyattr] - const ZSTD_c_enableLongDistanceMatching: i32 = 160; + const ZSTD_c_enableLongDistanceMatching: i32 = + c_param_id(ZSTD_cParameter::ZSTD_c_enableLongDistanceMatching); #[pyattr] - const ZSTD_c_ldmHashLog: i32 = 161; + const ZSTD_c_ldmHashLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmHashLog); #[pyattr] - const ZSTD_c_ldmMinMatch: i32 = 162; + const ZSTD_c_ldmMinMatch: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmMinMatch); #[pyattr] - const ZSTD_c_ldmBucketSizeLog: i32 = 163; + const ZSTD_c_ldmBucketSizeLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmBucketSizeLog); #[pyattr] - const ZSTD_c_ldmHashRateLog: i32 = 164; + const ZSTD_c_ldmHashRateLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmHashRateLog); #[pyattr] - const ZSTD_c_contentSizeFlag: i32 = 200; + const ZSTD_c_contentSizeFlag: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_contentSizeFlag); #[pyattr] - const ZSTD_c_checksumFlag: i32 = 201; + const ZSTD_c_checksumFlag: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_checksumFlag); #[pyattr] - const ZSTD_c_dictIDFlag: i32 = 202; + const ZSTD_c_dictIDFlag: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_dictIDFlag); #[pyattr] - const ZSTD_c_nbWorkers: i32 = 400; + const ZSTD_c_nbWorkers: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_nbWorkers); #[pyattr] - const ZSTD_c_jobSize: i32 = 401; + const ZSTD_c_jobSize: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_jobSize); #[pyattr] - const ZSTD_c_overlapLog: i32 = 402; + const ZSTD_c_overlapLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_overlapLog); // Decompression parameter identifiers. libzstd only exposes one non- // experimental decompression parameter. #[pyattr] - const ZSTD_d_windowLogMax: i32 = 100; + const ZSTD_d_windowLogMax: i32 = d_param_id(ZSTD_dParameter::ZSTD_d_windowLogMax); // Strategy enum members ordered from fastest to strongest. These power // the `Strategy` IntEnum in `Lib/compression/zstd/__init__.py`. @@ -125,16 +149,21 @@ mod _zstd { #[pyattr] const ZSTD_btultra2: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btultra2 as i32; - #[pyattr(once, name = "zstd_version")] + #[pyattr(once)] fn zstd_version(_vm: &VirtualMachine) -> String { // SAFETY: `ZSTD_versionString` returns a pointer to libzstd's static, // NUL-terminated version string. - unsafe { CStr::from_ptr(ZSTD_versionString()) } - .to_string_lossy() - .into_owned() + let version = unsafe { CStr::from_ptr(ZSTD_versionString()) }; + // The string is a fixed `MAJOR.MINOR.RELEASE` ASCII literal baked into + // libzstd, so it is always valid UTF-8 and there is nothing for a + // lossy conversion to repair. + version + .to_str() + .expect("libzstd version string is ASCII") + .to_owned() } - #[pyattr(once, name = "zstd_version_number")] + #[pyattr(once)] fn zstd_version_number(_vm: &VirtualMachine) -> u32 { ZSTD_versionNumber() } @@ -144,13 +173,35 @@ mod _zstd { ZSTD_DStreamOutSize() } - // Dictionary load type markers. The `ZstdDict.as_*` properties wrap the - // dictionary in a `(zdict, marker)` tuple so the compressor or decompressor - // constructor knows which load mode to apply. Numbering matches CPython's - // `Modules/_zstd/_zstdmodule.h::DictType`. - const DICT_TYPE_DIGESTED: i32 = 0; - const DICT_TYPE_UNDIGESTED: i32 = 1; - const DICT_TYPE_PREFIX: i32 = 2; + /// Which of the three load modes a `zstd_dict=` argument asks for. The + /// `ZstdDict.as_*` properties pair the dictionary with the matching + /// discriminant in a `(zdict, marker)` tuple, and the compressor and + /// decompressor constructors turn that marker back into this enum. + /// + /// The discriminants are pinned because they cross the Python boundary as + /// plain ints: they are the values CPython's + /// `Modules/_zstd/_zstdmodule.h::DictType` uses, and `test_zstd` builds + /// the tuples by hand with those literals. + #[derive(Copy, Clone)] + enum DictType { + Digested = 0, + Undigested = 1, + Prefix = 2, + } + + impl DictType { + /// Decode the int half of a `(zdict, marker)` tuple, which may come + /// straight from user code. `None` for anything outside the three + /// documented markers. + fn from_marker(marker: i32) -> Option { + Some(match marker { + m if m == Self::Digested as i32 => Self::Digested, + m if m == Self::Undigested as i32 => Self::Undigested, + m if m == Self::Prefix as i32 => Self::Prefix, + _ => return None, + }) + } + } #[pyattr(once, name = "ZstdError")] fn zstd_error(vm: &VirtualMachine) -> PyTypeRef { @@ -313,6 +364,11 @@ mod _zstd { /// libzstd silently clamps out-of-range values rather than surfacing /// them as errors. fn level_bounds() -> (i32, i32) { + // `lookup_param_bounds` only returns `None` for an unrecognized id or + // when libzstd flags the bounds as an error. `ZSTD_c_compressionLevel` + // is recognized, and `ZSTD_cParam_getBounds` answers it from + // `ZSTD_minCLevel()`/`ZSTD_maxCLevel()` without ever setting the error + // field, so `None` is unreachable here. lookup_param_bounds(ZSTD_c_compressionLevel, true) .expect("compressionLevel always has valid bounds") } @@ -401,12 +457,12 @@ mod _zstd { fn parse_zstd_dict_arg( obj: PyObjectRef, vm: &VirtualMachine, - ) -> PyResult<(PyRef, i32)> { + ) -> PyResult<(PyRef, DictType)> { // The first downcast clones `obj` because we fall through to the // tuple branch if it fails. The second downcast (the tuple one) is // the last use of `obj`, so we let it move directly. if let Ok(d) = obj.clone().downcast::() { - return Ok((d, DICT_TYPE_DIGESTED)); + return Ok((d, DictType::Digested)); } if let Ok(tuple) = obj.downcast::() { let items = tuple.as_slice(); @@ -433,10 +489,10 @@ mod _zstd { vm.new_type_error("zstd_dict argument should be a ZstdDict object") } })?; - if !(DICT_TYPE_DIGESTED..=DICT_TYPE_PREFIX).contains(&marker) { + let Some(dict_type) = DictType::from_marker(marker) else { return Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")); - } - return Ok((d, marker)); + }; + return Ok((d, dict_type)); } Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")) } @@ -538,22 +594,26 @@ mod _zstd { #[pygetset] fn as_digested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { - vm.ctx - .new_tuple(vec![zelf.into(), vm.ctx.new_int(DICT_TYPE_DIGESTED).into()]) + vm.ctx.new_tuple(vec![ + zelf.into(), + vm.ctx.new_int(DictType::Digested as i32).into(), + ]) } #[pygetset] fn as_undigested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { vm.ctx.new_tuple(vec![ zelf.into(), - vm.ctx.new_int(DICT_TYPE_UNDIGESTED).into(), + vm.ctx.new_int(DictType::Undigested as i32).into(), ]) } #[pygetset] fn as_prefix(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { - vm.ctx - .new_tuple(vec![zelf.into(), vm.ctx.new_int(DICT_TYPE_PREFIX).into()]) + vm.ctx.new_tuple(vec![ + zelf.into(), + vm.ctx.new_int(DictType::Prefix as i32).into(), + ]) } } @@ -576,8 +636,13 @@ mod _zstd { impl CCtx { fn create() -> Self { // SAFETY: `ZSTD_createCCtx` has no preconditions; it returns NULL - // only on allocation failure. + // only when its `malloc` of the context struct fails. let ptr = unsafe { ZSTD_createCCtx() }; + // A NULL return therefore means the allocator is out of memory, + // which is the one case Rust already handles by aborting, so + // panicking matches the rest of the VM rather than threading a + // `vm` reference through every construction site to raise + // `MemoryError` the way CPython's `_zstd` does. assert!(!ptr.is_null(), "ZSTD_createCCtx failed"); Self(ptr) } @@ -627,8 +692,13 @@ mod _zstd { impl DCtx { fn create() -> Self { // SAFETY: `ZSTD_createDCtx` has no preconditions; it returns NULL - // only on allocation failure. + // only when its `malloc` of the context struct fails. let ptr = unsafe { ZSTD_createDCtx() }; + // A NULL return therefore means the allocator is out of memory, + // which is the one case Rust already handles by aborting, so + // panicking matches the rest of the VM rather than threading a + // `vm` reference through every construction site to raise + // `MemoryError` the way CPython's `_zstd` does. assert!(!ptr.is_null(), "ZSTD_createDCtx failed"); Self(ptr) } @@ -1016,7 +1086,7 @@ mod _zstd { let Some(dict_obj) = dict_obj else { return Ok((None, None)); }; - let (zdict, marker) = parse_zstd_dict_arg(dict_obj, vm)?; + let (zdict, dict_type) = parse_zstd_dict_arg(dict_obj, vm)?; let bad_dict_err = || -> PyBaseExceptionRef { new_zstd_error( format!( @@ -1028,11 +1098,11 @@ mod _zstd { }; let dict_bytes = zdict.dict_content.as_bytes(); let mut digested = None; - match marker { - DICT_TYPE_PREFIX => { + match dict_type { + DictType::Prefix => { ctx.ref_prefix(dict_bytes).map_err(|_| bad_dict_err())?; } - DICT_TYPE_DIGESTED => { + DictType::Digested => { // Build the digested dict eagerly so a corrupted dictionary // surfaces as a `ZstdError` at construction time, not when // the first compress/decompress call runs. @@ -1040,9 +1110,9 @@ mod _zstd { ctx.ref_digested(&d).map_err(|_| bad_dict_err())?; digested = Some(d); } - _ => { - // Undigested: copy the bytes into the context. Validation - // happens lazily at the first stream call in this mode. + DictType::Undigested => { + // Copy the bytes into the context. Validation happens lazily + // at the first stream call in this mode. ctx.load_undigested(dict_bytes) .map_err(|_| bad_dict_err())?; } From a7f161c6c2771e00ca41ee21e6ae0e515e79678f Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 19 Aug 2026 10:06:33 +0100 Subject: [PATCH 17/25] Replace the _zstd flush-mode constants with a CompressMode enum `COMP_MODE_CONTINUE`/`FLUSH_BLOCK`/`FLUSH_FRAME` become a `CompressMode` enum with pinned discriminants, decoded at the Python boundary by `CompressMode::from_int` and mapped to libzstd by `CompressMode::end_directive`. `CompressorState::last_mode` now holds the enum, so the states `set_pledged_input_size` accepts are checked by the type system rather than by comparing ints. `flush()` keeps rejecting `CONTINUE` and out-of-range modes with its own message rather than the one `compress()` uses, so the Python-visible errors are unchanged. Co-Authored-By: Claude Opus 5 --- crates/stdlib/src/zstd.rs | 121 ++++++++++++++++++++++++++------------ 1 file changed, 83 insertions(+), 38 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 833e4da5cea..63624ba6d2f 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -617,12 +617,42 @@ mod _zstd { } } - // The three flush modes for `ZstdCompressor.compress()`, mirrored as - // class attributes via `extend_class` below. Values are positional and - // chosen to match what CPython exposes. - const COMP_MODE_CONTINUE: i32 = 0; - const COMP_MODE_FLUSH_BLOCK: i32 = 1; - const COMP_MODE_FLUSH_FRAME: i32 = 2; + /// The flush modes `ZstdCompressor.compress()` and + /// `ZstdCompressor.flush()` accept, mirrored as the class attributes + /// `CONTINUE`, `FLUSH_BLOCK` and `FLUSH_FRAME` by `extend_class` below. + /// + /// The discriminants are pinned because they cross the Python boundary as + /// plain ints, both as those class attributes and as the value the + /// `last_mode` getset hands back. They match what CPython's `_zstd` + /// exposes. + #[derive(Copy, Clone, PartialEq, Eq)] + enum CompressMode { + Continue = 0, + FlushBlock = 1, + FlushFrame = 2, + } + + impl CompressMode { + /// Decode a user-supplied `mode=` argument. `None` for anything + /// outside the three documented modes. + fn from_int(mode: i32) -> Option { + Some(match mode { + m if m == Self::Continue as i32 => Self::Continue, + m if m == Self::FlushBlock as i32 => Self::FlushBlock, + m if m == Self::FlushFrame as i32 => Self::FlushFrame, + _ => return None, + }) + } + + /// The libzstd end directive this mode drives the streaming API with. + fn end_directive(self) -> ZSTD_EndDirective { + match self { + Self::Continue => ZSTD_e_continue, + Self::FlushBlock => ZSTD_e_flush, + Self::FlushFrame => ZSTD_e_end, + } + } + } /// Owning wrapper around a raw libzstd compression context. Frees the /// context with `ZSTD_freeCCtx` on drop. @@ -793,7 +823,7 @@ mod _zstd { _cdict: Option, /// Keeps the ZstdDict's bytes alive for `ZSTD_CCtx_refPrefix` mode. _dict: Option>, - last_mode: i32, + last_mode: CompressMode, } #[pyattr] @@ -819,20 +849,22 @@ mod _zstd { zstd_dict: OptionalOption, } - /// Translate the public `mode` int to the libzstd `ZSTD_EndDirective` - /// the streaming API takes. - fn end_directive_from_mode(mode: i32, vm: &VirtualMachine) -> PyResult { - match mode { - COMP_MODE_CONTINUE => Ok(ZSTD_e_continue), - COMP_MODE_FLUSH_BLOCK => Ok(ZSTD_e_flush), - COMP_MODE_FLUSH_FRAME => Ok(ZSTD_e_end), - _ => Err(vm.new_value_error(format!( + /// Decode the public `mode` int accepted by `ZstdCompressor.compress()`, + /// which takes all three modes. + fn compress_mode_from_int(mode: i32, vm: &VirtualMachine) -> PyResult { + CompressMode::from_int(mode).ok_or_else(|| { + let (cont, block, frame) = ( + CompressMode::Continue as i32, + CompressMode::FlushBlock as i32, + CompressMode::FlushFrame as i32, + ); + vm.new_value_error(format!( "mode argument wrong value, it should be one of \ - ZstdCompressor.CONTINUE ({COMP_MODE_CONTINUE}), \ - ZstdCompressor.FLUSH_BLOCK ({COMP_MODE_FLUSH_BLOCK}), or \ - ZstdCompressor.FLUSH_FRAME ({COMP_MODE_FLUSH_FRAME})" - ))), - } + ZstdCompressor.CONTINUE ({cont}), \ + ZstdCompressor.FLUSH_BLOCK ({block}), or \ + ZstdCompressor.FLUSH_FRAME ({frame})" + )) + }) } impl Constructor for ZstdCompressor { @@ -859,7 +891,7 @@ mod _zstd { apply_options(&mut cctx, options_obj, true, vm)?; } - let state = build_compressor_state(cctx, dict_opt, COMP_MODE_FLUSH_FRAME, vm)?; + let state = build_compressor_state(cctx, dict_opt, CompressMode::FlushFrame, vm)?; Ok(Self { state: PyMutex::new(state), }) @@ -1128,7 +1160,7 @@ mod _zstd { fn build_compressor_state( mut cctx: CCtx, dict_obj: Option, - last_mode: i32, + last_mode: CompressMode, vm: &VirtualMachine, ) -> PyResult { // SAFETY: `load_dict` requires its two return values to outlive `ctx`. @@ -1238,8 +1270,11 @@ mod _zstd { impl ZstdCompressor { #[pymethod] fn compress(&self, args: CompressMethodArgs, vm: &VirtualMachine) -> PyResult> { - let mode = args.mode.unwrap_or(COMP_MODE_CONTINUE); - let end_op = end_directive_from_mode(mode, vm)?; + let mode = match args.mode { + Some(mode) => compress_mode_from_int(mode, vm)?, + None => CompressMode::Continue, + }; + let end_op = mode.end_directive(); let data = args.data.with_ref(|b| b.to_vec()); let mut state = self.state.lock(); let out = do_compress(&mut state, &data, end_op, vm)?; @@ -1249,15 +1284,25 @@ mod _zstd { #[pymethod] fn flush(&self, args: FlushMethodArgs, vm: &VirtualMachine) -> PyResult> { - let mode = args.mode.unwrap_or(COMP_MODE_FLUSH_FRAME); - if mode != COMP_MODE_FLUSH_BLOCK && mode != COMP_MODE_FLUSH_FRAME { - return Err(vm.new_value_error(format!( - "mode argument wrong value, it should be \ - ZstdCompressor.FLUSH_FRAME ({COMP_MODE_FLUSH_FRAME}) or \ - ZstdCompressor.FLUSH_BLOCK ({COMP_MODE_FLUSH_BLOCK})" - ))); - } - let end_op = end_directive_from_mode(mode, vm)?; + // `flush()` takes only the two flushing modes, and rejects + // everything else — `CONTINUE` included — with its own message + // rather than the one `compress()` uses. + let mode = match args.mode.map(CompressMode::from_int) { + None => CompressMode::FlushFrame, + Some(Some(mode @ (CompressMode::FlushBlock | CompressMode::FlushFrame))) => mode, + Some(_) => { + let (block, frame) = ( + CompressMode::FlushBlock as i32, + CompressMode::FlushFrame as i32, + ); + return Err(vm.new_value_error(format!( + "mode argument wrong value, it should be \ + ZstdCompressor.FLUSH_FRAME ({frame}) or \ + ZstdCompressor.FLUSH_BLOCK ({block})" + ))); + } + }; + let end_op = mode.end_directive(); let mut state = self.state.lock(); let out = do_compress(&mut state, &[], end_op, vm)?; state.last_mode = mode; @@ -1299,7 +1344,7 @@ mod _zstd { Some(v) }; let mut state = self.state.lock(); - if state.last_mode != COMP_MODE_FLUSH_FRAME { + if state.last_mode != CompressMode::FlushFrame { return Err(vm.new_value_error( "set_pledged_input_size() method must be called when last_mode == FLUSH_FRAME", )); @@ -1313,7 +1358,7 @@ mod _zstd { #[pygetset] fn last_mode(&self) -> i32 { - self.state.lock().last_mode + self.state.lock().last_mode as i32 } // Install class-level constants `CONTINUE`, `FLUSH_BLOCK`, and @@ -1324,15 +1369,15 @@ mod _zstd { fn extend_class(ctx: &Context, class: &'static Py) { class.set_attr( ctx.intern_str("CONTINUE"), - ctx.new_int(COMP_MODE_CONTINUE).into(), + ctx.new_int(CompressMode::Continue as i32).into(), ); class.set_attr( ctx.intern_str("FLUSH_BLOCK"), - ctx.new_int(COMP_MODE_FLUSH_BLOCK).into(), + ctx.new_int(CompressMode::FlushBlock as i32).into(), ); class.set_attr( ctx.intern_str("FLUSH_FRAME"), - ctx.new_int(COMP_MODE_FLUSH_FRAME).into(), + ctx.new_int(CompressMode::FlushFrame as i32).into(), ); } } From 04d560cf8988c71613395c10fac509323e9471ea Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 19 Aug 2026 10:54:25 +0100 Subject: [PATCH 18/25] Correct the _zstd explanatory comments and raise MemoryError on OOM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Verified each claim the previous two commits added, against the pinned libzstd-rs-sys rev, CPython's `Modules/_zstd/` and `Lib/test/test_zstd.py`, and corrected the ones that did not hold: - The `DictType` discriminants come from CPython's `dictionary_type` enum, not a type named `DictType`, and `test_zstd` never hand-builds a valid `(zdict, marker)` tuple — it pins the accepted range from the outside with `(zd, -1)` and `(zd, 3)` rejection cases. - `#[repr(transparent)]` guarantees layout, not validity invariants. The transmute is sound because every bit pattern is a valid `u32`, which holds only in the enum-to-int direction; say so, and note that this is why `c_param_enum` decodes untrusted ids with an explicit match. - The id is unreachable through the crate's public API only in the sense that no accessor exists; the derived `Debug` does print it. - Undigested dictionary loading is not uniformly lazy: `ZSTD_DCtx_loadDictionary` digests eagerly and rejects corrupted content at construction time, while `ZSTD_CCtx_loadDictionary` accepts it. - `CompressMode`'s discriminants are libzstd's `ZSTD_e_continue`/ `ZSTD_e_flush`/`ZSTD_e_end` values. `CCtx::create`/`DCtx::create` now return `PyResult` and raise `MemoryError` like CPython's `_zstd` does, rather than panicking. The old justification for the panic claimed a `vm` reference would have to be threaded through every construction site; there are two, and both already had one. Co-Authored-By: Claude Opus 5 --- crates/stdlib/src/zstd.rs | 75 +++++++++++++++++++++------------------ 1 file changed, 40 insertions(+), 35 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 63624ba6d2f..59247cb2fed 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -62,13 +62,17 @@ mod _zstd { /// `CompressionParameter` IntEnum. libzstd-rs-sys declares /// `ZSTD_cParameter` as `#[repr(transparent)] pub struct /// ZSTD_cParameter(pub(crate) u32)` and gives it no accessor, `From` impl - /// or `Deref`, so the id is unreachable through its public API. Reading - /// the wrapped `u32` back keeps libzstd the single source of truth for - /// these values instead of restating them here as literals. + /// or `Deref`, so short of parsing its derived `Debug` output there is no + /// way to ask it for its id. Reading the wrapped `u32` back keeps libzstd + /// the single source of truth for these values instead of restating them + /// here as literals. const fn c_param_id(param: ZSTD_cParameter) -> i32 { - // SAFETY: `#[repr(transparent)]` guarantees `ZSTD_cParameter` has the - // same size, alignment and validity invariants as the `u32` it wraps, - // so reading one back as a `u32` is always well-defined. + // SAFETY: `#[repr(transparent)]` gives `ZSTD_cParameter` the layout of + // the single `u32` it wraps, and every bit pattern is a valid `u32`, + // so reading one back out is well-defined. Note this holds only in + // this direction: `u32 -> ZSTD_cParameter` would be a widening into a + // type whose valid values libzstd defines, which is why `c_param_enum` + // below decodes untrusted ids with an explicit match instead. unsafe { core::mem::transmute::(param) as i32 } } @@ -179,9 +183,10 @@ mod _zstd { /// decompressor constructors turn that marker back into this enum. /// /// The discriminants are pinned because they cross the Python boundary as - /// plain ints: they are the values CPython's - /// `Modules/_zstd/_zstdmodule.h::DictType` uses, and `test_zstd` builds - /// the tuples by hand with those literals. + /// plain ints: they are the values of CPython's `dictionary_type` enum in + /// `Modules/_zstd/_zstdmodule.h`. `test_zstd` pins the accepted range from + /// the outside, hand-building `(zd, -1)` and `(zd, 3)` tuples and + /// requiring both to raise. #[derive(Copy, Clone)] enum DictType { Digested = 0, @@ -621,10 +626,11 @@ mod _zstd { /// `ZstdCompressor.flush()` accept, mirrored as the class attributes /// `CONTINUE`, `FLUSH_BLOCK` and `FLUSH_FRAME` by `extend_class` below. /// - /// The discriminants are pinned because they cross the Python boundary as - /// plain ints, both as those class attributes and as the value the - /// `last_mode` getset hands back. They match what CPython's `_zstd` - /// exposes. + /// The discriminants are libzstd's own `ZSTD_e_continue`, `ZSTD_e_flush` + /// and `ZSTD_e_end` values, which is what CPython's `_zstd` exposes under + /// these attribute names. They are pinned here because they cross the + /// Python boundary as plain ints, both as those class attributes and as + /// the value the `last_mode` getset hands back. #[derive(Copy, Clone, PartialEq, Eq)] enum CompressMode { Continue = 0, @@ -664,17 +670,15 @@ mod _zstd { unsafe impl Send for CCtx {} impl CCtx { - fn create() -> Self { + fn create(vm: &VirtualMachine) -> PyResult { // SAFETY: `ZSTD_createCCtx` has no preconditions; it returns NULL - // only when its `malloc` of the context struct fails. + // only when its allocation of the context struct fails, so that is + // the sole case to handle. let ptr = unsafe { ZSTD_createCCtx() }; - // A NULL return therefore means the allocator is out of memory, - // which is the one case Rust already handles by aborting, so - // panicking matches the rest of the VM rather than threading a - // `vm` reference through every construction site to raise - // `MemoryError` the way CPython's `_zstd` does. - assert!(!ptr.is_null(), "ZSTD_createCCtx failed"); - Self(ptr) + if ptr.is_null() { + return Err(vm.new_memory_error("failed to allocate a zstd context")); + } + Ok(Self(ptr)) } /// Forward a `(parameter, value)` pair to libzstd, passing the value @@ -720,17 +724,15 @@ mod _zstd { unsafe impl Send for DCtx {} impl DCtx { - fn create() -> Self { + fn create(vm: &VirtualMachine) -> PyResult { // SAFETY: `ZSTD_createDCtx` has no preconditions; it returns NULL - // only when its `malloc` of the context struct fails. + // only when its allocation of the context struct fails, so that is + // the sole case to handle. let ptr = unsafe { ZSTD_createDCtx() }; - // A NULL return therefore means the allocator is out of memory, - // which is the one case Rust already handles by aborting, so - // panicking matches the rest of the VM rather than threading a - // `vm` reference through every construction site to raise - // `MemoryError` the way CPython's `_zstd` does. - assert!(!ptr.is_null(), "ZSTD_createDCtx failed"); - Self(ptr) + if ptr.is_null() { + return Err(vm.new_memory_error("failed to allocate a zstd context")); + } + Ok(Self(ptr)) } /// See [`CCtx::set_parameter`]. @@ -879,7 +881,7 @@ mod _zstd { return Err(vm.new_type_error("Only one of level or options should be used.")); } - let mut cctx = CCtx::create(); + let mut cctx = CCtx::create(vm)?; if let Some(level_obj) = level_opt { let level = parse_compression_level(&level_obj, vm)?; @@ -1143,8 +1145,11 @@ mod _zstd { digested = Some(d); } DictType::Undigested => { - // Copy the bytes into the context. Validation happens lazily - // at the first stream call in this mode. + // Copy the bytes into the context. Whether a corrupted + // dictionary is caught here is up to libzstd and differs by + // direction: `ZSTD_DCtx_loadDictionary` digests the content + // eagerly and fails now, while `ZSTD_CCtx_loadDictionary` + // defers that work and accepts it. ctx.load_undigested(dict_bytes) .map_err(|_| bad_dict_err())?; } @@ -1445,7 +1450,7 @@ mod _zstd { let dict_opt = args.zstd_dict.flatten(); let options_opt = args.options.flatten(); - let mut dctx = DCtx::create(); + let mut dctx = DCtx::create(vm)?; if let Some(options_obj) = options_opt { apply_options(&mut dctx, options_obj, false, vm)?; From 97b6f48e010b179f973944b7d4bf00fc0f0887d4 Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 19 Aug 2026 12:44:11 +0100 Subject: [PATCH 19/25] Fix CPython divergences and a crash in the _zstd module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Found by reviewing the module against CPython's `Modules/_zstd/` and the pinned libzstd-rs-sys sources; each item below was reproduced before the fix and re-checked after. - `train_dict`/`finalize_dict` sized their output buffer with an infallible `Vec::with_capacity`/`vec![]` from a caller-supplied `dict_size`, so `train_dict([b'x'], 2**62)` aborted the interpreter instead of raising. Reserve fallibly and raise `MemoryError`, as CPython does. - A digested dictionary was always built at `ZSTD_CLEVEL_DEFAULT`. Since `ZSTD_CCtx_refCDict` takes its parameters from the CDict, that silently discarded the requested level: `level=19` with `as_digested_dict` compressed a test corpus to 39817 bytes instead of 31951. Build the CDict at the compressor's effective level, tracked through `level=` and the `options` dict like CPython's `self->compression_level`. - A bare `ZstdDict` was loaded as digested in both directions; CPython compresses with an undigested dictionary by default and only decompresses with a digested one. Give `DictLoader` a per-direction default. - Neither compressor nor decompressor reset its session after an error, leaving the object permanently broken — every later call failed with "Operation not authorized at current processing stage" or re-reported the original corruption on valid input. Reset as CPython's error paths do, including restoring `last_mode` to `FLUSH_FRAME`. - `apply_options` range-checked every parameter, rejecting values libzstd and CPython accept: 0 ("use the default") for `window_log`, `strategy`, `window_log_max` and friends, any non-zero value for the boolean flags, and the clamped `nb_workers`/`job_size`/`overlap_log`. Only `compression_level` needs the check — libzstd clamps that one silently, which is why `test_compress_parameters` requires a `ValueError` — so everything else now goes to libzstd, whose error code already maps to the same message. - `ZstdDict` accepted content shorter than 8 bytes; CPython rejects it up front regardless of `is_raw`. - `compress`/`decompress` refused `data` as a keyword and `get_param_bounds` required `is_compress` as one, both contrary to CPython's clinic signatures. - The decompressor allocated a full 128 KiB scratch buffer even when `max_length` capped output far below that. `test_zstd` (119 tests) and the snippet test pass; clippy is clean. Co-Authored-By: Claude Opus 5 --- crates/stdlib/src/zstd.rs | 230 ++++++++++++++++++++++++++++++-------- 1 file changed, 182 insertions(+), 48 deletions(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 59247cb2fed..c873fda3cbd 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -29,17 +29,17 @@ mod _zstd { use libzstd_rs_sys::lib::zdict::ZDICT_finalizeDictionary; use libzstd_rs_sys::{ ZDICT_getErrorName, ZDICT_isError, ZDICT_params_t, ZDICT_trainFromBuffer, ZSTD_CCtx, - ZSTD_CCtx_loadDictionary, ZSTD_CCtx_refCDict, ZSTD_CCtx_refPrefix, ZSTD_CCtx_setParameter, - ZSTD_CCtx_setPledgedSrcSize, ZSTD_CDict, ZSTD_CONTENTSIZE_ERROR, ZSTD_CONTENTSIZE_UNKNOWN, - ZSTD_CStreamOutSize, ZSTD_DCtx, ZSTD_DCtx_loadDictionary, ZSTD_DCtx_refDDict, - ZSTD_DCtx_refPrefix, ZSTD_DCtx_setParameter, ZSTD_DDict, ZSTD_DStreamOutSize, - ZSTD_EndDirective, ZSTD_cParam_getBounds, ZSTD_cParameter, ZSTD_compressStream2, - ZSTD_createCCtx, ZSTD_createCDict, ZSTD_createDCtx, ZSTD_createDDict, - ZSTD_dParam_getBounds, ZSTD_dParameter, ZSTD_decompressStream, - ZSTD_findFrameCompressedSize, ZSTD_freeCCtx, ZSTD_freeCDict, ZSTD_freeDCtx, ZSTD_freeDDict, - ZSTD_getDictID_fromDict, ZSTD_getDictID_fromFrame, ZSTD_getErrorName, - ZSTD_getFrameContentSize, ZSTD_inBuffer, ZSTD_isError, ZSTD_outBuffer, ZSTD_strategy, - ZSTD_versionNumber, ZSTD_versionString, + ZSTD_CCtx_loadDictionary, ZSTD_CCtx_refCDict, ZSTD_CCtx_refPrefix, ZSTD_CCtx_reset, + ZSTD_CCtx_setParameter, ZSTD_CCtx_setPledgedSrcSize, ZSTD_CDict, ZSTD_CONTENTSIZE_ERROR, + ZSTD_CONTENTSIZE_UNKNOWN, ZSTD_CStreamOutSize, ZSTD_DCtx, ZSTD_DCtx_loadDictionary, + ZSTD_DCtx_refDDict, ZSTD_DCtx_refPrefix, ZSTD_DCtx_reset, ZSTD_DCtx_setParameter, + ZSTD_DDict, ZSTD_DStreamOutSize, ZSTD_EndDirective, ZSTD_ResetDirective, + ZSTD_cParam_getBounds, ZSTD_cParameter, ZSTD_compressStream2, ZSTD_createCCtx, + ZSTD_createCDict, ZSTD_createDCtx, ZSTD_createDDict, ZSTD_dParam_getBounds, + ZSTD_dParameter, ZSTD_decompressStream, ZSTD_findFrameCompressedSize, ZSTD_freeCCtx, + ZSTD_freeCDict, ZSTD_freeDCtx, ZSTD_freeDDict, ZSTD_getDictID_fromDict, + ZSTD_getDictID_fromFrame, ZSTD_getErrorName, ZSTD_getFrameContentSize, ZSTD_inBuffer, + ZSTD_isError, ZSTD_outBuffer, ZSTD_strategy, ZSTD_versionNumber, ZSTD_versionString, }; use rustpython_common::lock::PyMutex; use rustpython_vm::builtins::{ @@ -317,7 +317,9 @@ mod _zstd { "invalid compression parameter 'unknown parameter (key {param})'" )) })?; - if param == ZSTD_c_strategy && strategy_from_int(value).is_none() { + // Strategy 0 is libzstd's "use the default" sentinel, like the other + // parameters; only a non-zero value has to name a real strategy. + if param == ZSTD_c_strategy && value != 0 && strategy_from_int(value).is_none() { return Err(new_zstd_error( format!("invalid strategy value: {value}"), vm, @@ -457,17 +459,18 @@ mod _zstd { } /// Decode the `zstd_dict=` constructor argument. Accepts either a - /// `ZstdDict` instance (treated as the default digested form) or a - /// `(ZstdDict, marker)` tuple produced by one of `ZstdDict.as_*`. + /// `ZstdDict` instance, which takes the caller's `default` load mode, or + /// a `(ZstdDict, marker)` tuple produced by one of `ZstdDict.as_*`. fn parse_zstd_dict_arg( obj: PyObjectRef, + default: DictType, vm: &VirtualMachine, ) -> PyResult<(PyRef, DictType)> { // The first downcast clones `obj` because we fall through to the // tuple branch if it fails. The second downcast (the tuple one) is // the last use of `obj`, so we let it move directly. if let Ok(d) = obj.clone().downcast::() { - return Ok((d, DictType::Digested)); + return Ok((d, default)); } if let Ok(tuple) = obj.downcast::() { let items = tuple.as_slice(); @@ -534,6 +537,15 @@ mod _zstd { fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { let dict_content = args.dict_content.with_ref(|b| b.to_vec()); + // CPython checks this first, before and independently of the + // `is_raw` handling below, so a too-short buffer never reaches + // libzstd. + if dict_content.len() < 8 { + return Err(vm.new_value_error( + "Zstandard dictionary content too short \ + (must have at least eight bytes)", + )); + } // libzstd's `ZSTD_getDictID_fromDict` returns 0 either when // the content is too small to contain a valid header or when it // does not carry the dictionary magic. Both are runtime errors @@ -694,6 +706,14 @@ mod _zstd { } } + /// Return the context to a frame boundary, discarding a partially + /// written frame but keeping parameters and dictionary. libzstd + /// documents a session-only reset as infallible. + fn reset_session(&mut self) { + // SAFETY: `self.0` is a live context owned by this wrapper. + unsafe { ZSTD_CCtx_reset(self.0, ZSTD_ResetDirective::ZSTD_reset_session_only) }; + } + fn set_pledged_src_size(&mut self, pledged: Option) -> Result<(), usize> { // libzstd represents "unknown" as `ZSTD_CONTENTSIZE_UNKNOWN`. let pledged = pledged.map_or(ZSTD_CONTENTSIZE_UNKNOWN, |v| v); @@ -735,6 +755,12 @@ mod _zstd { Ok(Self(ptr)) } + /// See [`CCtx::reset_session`]. + fn reset_session(&mut self) { + // SAFETY: `self.0` is a live context owned by this wrapper. + unsafe { ZSTD_DCtx_reset(self.0, ZSTD_ResetDirective::ZSTD_reset_session_only) }; + } + /// See [`CCtx::set_parameter`]. fn set_parameter(&mut self, param: ZSTD_dParameter, value: c_int) -> Result<(), usize> { // SAFETY: `self.0` is a live context owned by this wrapper. @@ -828,6 +854,17 @@ mod _zstd { last_mode: CompressMode, } + impl CompressorState { + /// Recover from a failed `compress()`/`flush()`, as CPython does: + /// drop the half-written frame and report `last_mode == FLUSH_FRAME`. + /// Without this the context stays mid-frame and every later call + /// fails with "Operation not authorized at current processing stage". + fn reset_session(&mut self) { + self.last_mode = CompressMode::FlushFrame; + self.cctx.reset_session(); + } + } + #[pyattr] #[pyclass(name = "ZstdCompressor")] #[derive(PyPayload)] @@ -883,17 +920,25 @@ mod _zstd { let mut cctx = CCtx::create(vm)?; + // Mirror CPython's `self->compression_level`: the level libzstd + // will actually compress at, which a digested dictionary has to + // be built for. + let mut level = ZSTD_CLEVEL_DEFAULT; + if let Some(level_obj) = level_opt { - let level = parse_compression_level(&level_obj, vm)?; + level = parse_compression_level(&level_obj, vm)?; cctx.set_parameter(ZSTD_cParameter::ZSTD_c_compressionLevel, level) .map_err(|_| param_value_error_for(ZSTD_c_compressionLevel, level, true, vm))?; } - if let Some(options_obj) = options_opt { - apply_options(&mut cctx, options_obj, true, vm)?; + if let Some(options_obj) = options_opt + && let Some(from_options) = apply_options(&mut cctx, options_obj, true, vm)? + { + level = from_options; } - let state = build_compressor_state(cctx, dict_opt, CompressMode::FlushFrame, vm)?; + let state = + build_compressor_state(cctx, dict_opt, level, CompressMode::FlushFrame, vm)?; Ok(Self { state: PyMutex::new(state), }) @@ -932,7 +977,8 @@ mod _zstd { options_obj: PyObjectRef, is_compress: bool, vm: &VirtualMachine, - ) -> PyResult<()> { + ) -> PyResult> { + let mut level = None; let dict = options_obj .downcast::() .map_err(|_| vm.new_type_error("options must be a dict"))?; @@ -953,17 +999,34 @@ mod _zstd { check_wrong_param_kind(&k, wrong_kind.as_ref(), kind, vm)?; let key_int: i32 = k.try_to_value(vm)?; let val_int: i32 = v.try_to_value(vm)?; - // libzstd silently clamps out-of-range values for some - // parameters (notably compression_level) rather than rejecting - // them, so validate against the documented bounds upfront. - if let Some((lo, hi)) = lookup_param_bounds(key_int, is_compress) + // Only the compression level is range-checked here. libzstd + // silently clamps it (`ZSTD_cParam_clampBounds`) instead of + // reporting an error, so without this check the `ValueError` + // that `test_compress_parameters` requires would never be + // raised. + // + // Every other parameter is left to libzstd, which does its own + // bounds checking and whose error code becomes the same + // `ValueError` via `ParamSetter::apply`. A blanket range test + // here would reject values CPython accepts: 0 means "use the + // default" for most parameters, the boolean flags take any + // non-zero value, and `nb_workers`/`job_size`/`overlap_log` are + // clamped rather than rejected. + if is_compress + && key_int == ZSTD_c_compressionLevel + && let Some((lo, hi)) = lookup_param_bounds(key_int, is_compress) && (val_int < lo || val_int > hi) { return Err(param_value_error_for(key_int, val_int, is_compress, vm)); } ctx.apply(key_int, val_int, vm)?; + // Track the level the compressor ends up with: a digested + // dictionary has to be built at that level (see `load_dict`). + if is_compress && key_int == ZSTD_c_compressionLevel { + level = Some(val_int); + } } - Ok(()) + Ok(level) } /// Trait wrapper over `CCtx::set_parameter` and `DCtx::set_parameter` so @@ -1000,7 +1063,13 @@ mod _zstd { trait DictLoader { type Digested; const KIND_NAME: &'static str; - fn try_create_digested(bytes: &[u8]) -> Option; + /// Load mode applied when the caller passes a bare `ZstdDict` instead + /// of one of the `as_*` tuples. CPython compresses with an undigested + /// dictionary by default and decompresses with a digested one. + const DEFAULT_DICT_TYPE: DictType; + /// `level` is the compression level the digested dictionary must be + /// built for; the decompressor has no level and ignores it. + fn try_create_digested(bytes: &[u8], level: c_int) -> Option; fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize>; fn load_undigested(&mut self, bytes: &[u8]) -> Result<(), usize>; fn ref_prefix(&mut self, bytes: &[u8]) -> Result<(), usize>; @@ -1009,8 +1078,12 @@ mod _zstd { impl DictLoader for CCtx { type Digested = CDict; const KIND_NAME: &'static str = "ZSTD_CDict"; - fn try_create_digested(bytes: &[u8]) -> Option { - CDict::try_create(bytes, libzstd_rs_sys::ZSTD_CLEVEL_DEFAULT) + const DEFAULT_DICT_TYPE: DictType = DictType::Undigested; + fn try_create_digested(bytes: &[u8], level: c_int) -> Option { + // The level matters: `ZSTD_CCtx_refCDict` takes its compression + // parameters from the CDict, so a CDict built at the default + // level would override whatever the caller asked for. + CDict::try_create(bytes, level) } fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize> { // SAFETY: both handles are live; `load_dict`'s safety contract @@ -1050,7 +1123,8 @@ mod _zstd { impl DictLoader for DCtx { type Digested = DDict; const KIND_NAME: &'static str = "ZSTD_DDict"; - fn try_create_digested(bytes: &[u8]) -> Option { + const DEFAULT_DICT_TYPE: DictType = DictType::Digested; + fn try_create_digested(bytes: &[u8], _level: c_int) -> Option { DDict::try_create(bytes) } fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize> { @@ -1115,12 +1189,13 @@ mod _zstd { unsafe fn load_dict( ctx: &mut L, dict_obj: Option, + level: c_int, vm: &VirtualMachine, ) -> DictLoadResult { let Some(dict_obj) = dict_obj else { return Ok((None, None)); }; - let (zdict, dict_type) = parse_zstd_dict_arg(dict_obj, vm)?; + let (zdict, dict_type) = parse_zstd_dict_arg(dict_obj, L::DEFAULT_DICT_TYPE, vm)?; let bad_dict_err = || -> PyBaseExceptionRef { new_zstd_error( format!( @@ -1140,7 +1215,7 @@ mod _zstd { // Build the digested dict eagerly so a corrupted dictionary // surfaces as a `ZstdError` at construction time, not when // the first compress/decompress call runs. - let d = L::try_create_digested(dict_bytes).ok_or_else(bad_dict_err)?; + let d = L::try_create_digested(dict_bytes, level).ok_or_else(bad_dict_err)?; ctx.ref_digested(&d).map_err(|_| bad_dict_err())?; digested = Some(d); } @@ -1165,6 +1240,7 @@ mod _zstd { fn build_compressor_state( mut cctx: CCtx, dict_obj: Option, + level: c_int, last_mode: CompressMode, vm: &VirtualMachine, ) -> PyResult { @@ -1176,7 +1252,7 @@ mod _zstd { // and `_dict` (prefix mode) are freed. `CompressorState` is private // to this module and is never destructured, so no safe caller can // reorder the drops. - let (cdict, dict) = unsafe { load_dict::(&mut cctx, dict_obj, vm) }?; + let (cdict, dict) = unsafe { load_dict::(&mut cctx, dict_obj, level, vm) }?; Ok(CompressorState { cctx, _cdict: cdict, @@ -1194,7 +1270,9 @@ mod _zstd { vm: &VirtualMachine, ) -> PyResult { // SAFETY: see [`build_compressor_state`]. - let (ddict, dict) = unsafe { load_dict::(&mut dctx, dict_obj, vm) }?; + // The decompressor has no compression level; `DCtx` ignores it. + let (ddict, dict) = + unsafe { load_dict::(&mut dctx, dict_obj, ZSTD_CLEVEL_DEFAULT, vm) }?; Ok(DecompressorState { dctx, _ddict: ddict, @@ -1282,9 +1360,16 @@ mod _zstd { let end_op = mode.end_directive(); let data = args.data.with_ref(|b| b.to_vec()); let mut state = self.state.lock(); - let out = do_compress(&mut state, &data, end_op, vm)?; - state.last_mode = mode; - Ok(out) + match do_compress(&mut state, &data, end_op, vm) { + Ok(out) => { + state.last_mode = mode; + Ok(out) + } + Err(e) => { + state.reset_session(); + Err(e) + } + } } #[pymethod] @@ -1309,9 +1394,16 @@ mod _zstd { }; let end_op = mode.end_directive(); let mut state = self.state.lock(); - let out = do_compress(&mut state, &[], end_op, vm)?; - state.last_mode = mode; - Ok(out) + match do_compress(&mut state, &[], end_op, vm) { + Ok(out) => { + state.last_mode = mode; + Ok(out) + } + Err(e) => { + state.reset_session(); + Err(e) + } + } } #[pymethod] @@ -1389,7 +1481,9 @@ mod _zstd { #[derive(FromArgs)] pub(super) struct CompressMethodArgs { - #[pyarg(positional)] + /// Positional-or-keyword, as in CPython's clinic signature for + /// `_zstd.ZstdCompressor.compress` (no `/` marker). + #[pyarg(any)] data: ArgBytesLike, #[pyarg(any, optional)] mode: Option, @@ -1422,6 +1516,20 @@ mod _zstd { input_buffer: Vec, } + impl DecompressorState { + /// Recover from a failed `decompress()`, mirroring CPython's + /// `decompressor_reset_session_lock_held`. Without it the DCtx stays + /// wedged on the corrupt frame and every later call re-reports the + /// same error, even for valid input. + fn reset_session(&mut self, vm: &VirtualMachine) { + self.input_buffer.clear(); + self.unused_data = vm.ctx.empty_bytes.clone(); + self.needs_input = true; + self.eof = false; + self.dctx.reset_session(); + } + } + #[pyattr] #[pyclass(name = "ZstdDecompressor")] #[derive(PyPayload)] @@ -1465,7 +1573,9 @@ mod _zstd { #[derive(FromArgs)] pub(super) struct DecompressMethodArgs { - #[pyarg(positional)] + /// Positional-or-keyword, as in CPython's clinic signature for + /// `_zstd.ZstdDecompressor.decompress`. + #[pyarg(any)] data: ArgBytesLike, #[pyarg(any, default = -1)] max_length: isize, @@ -1514,8 +1624,13 @@ mod _zstd { // Reusable scratch buffer for each decompress_stream call. We need // an exact-size output buffer because `Vec::reserve` may // over-allocate; reporting the full Vec capacity to libzstd would - // let it write past `max_length`. - let mut scratch: Vec = vec![0u8; chunk_size]; + // let it write past `max_length`. A cap below one chunk bounds + // how much of it can ever be used, so do not allocate more. + let scratch_size = match max_length { + Some(maxl) => maxl.min(chunk_size), + None => chunk_size, + }; + let mut scratch: Vec = vec![0u8; scratch_size]; let mut hit_max = false; let mut iteration = 0usize; @@ -1585,7 +1700,13 @@ mod _zstd { outcome.map(|()| (output, hit_max, input.pos)) }); - let (output, hit_max, consumed) = loop_result.map_err(|c| catch_zstd_error(c, vm))?; + let (output, hit_max, consumed) = match loop_result { + Ok(v) => v, + Err(code) => { + state.reset_session(vm); + return Err(catch_zstd_error(code, vm)); + } + }; let remaining = &work_data[consumed..]; @@ -1778,7 +1899,13 @@ mod _zstd { let samples_buffer = args.samples_bytes.as_bytes().to_vec(); let sizes = parse_sample_sizes(args.samples_sizes, vm)?; check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; - let mut dict_buffer: Vec = Vec::with_capacity(dict_size); + // `dict_size` is caller-controlled, so reserve fallibly: an + // infallible `Vec::with_capacity` aborts the whole interpreter on a + // failed allocation, where CPython raises `MemoryError`. + let mut dict_buffer: Vec = Vec::new(); + dict_buffer + .try_reserve_exact(dict_size) + .map_err(|_| vm.new_memory_error("failed to allocate the dictionary buffer"))?; // SAFETY: `dict_buffer`'s spare capacity is valid for writes of up to // `dict_size` bytes; the samples and sizes buffers are live slices // for the duration of the call. @@ -1833,7 +1960,12 @@ mod _zstd { let sizes = parse_sample_sizes(args.samples_sizes, vm)?; check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; - let mut dict_buffer: Vec = vec![0u8; dict_size]; + // Reserve fallibly before filling: see `train_dict`. + let mut dict_buffer: Vec = Vec::new(); + dict_buffer + .try_reserve_exact(dict_size) + .map_err(|_| vm.new_memory_error("failed to allocate the dictionary buffer"))?; + dict_buffer.resize(dict_size, 0); let params = ZDICT_params_t { compressionLevel: compression_level, notificationLevel: 0, @@ -1874,9 +2006,11 @@ mod _zstd { #[derive(FromArgs)] pub(super) struct ParamBoundsArgs { - #[pyarg(positional)] + #[pyarg(any)] parameter: i32, - #[pyarg(named)] + /// Positional-or-keyword, matching CPython's clinic signature + /// `get_param_bounds(parameter, is_compress)`. + #[pyarg(any)] is_compress: bool, } From acd733c9aa533b248409f71b0f214047b630d23c Mon Sep 17 00:00:00 2001 From: James David Clarke Date: Wed, 19 Aug 2026 13:40:51 +0100 Subject: [PATCH 20/25] Restore Lib/test/test_zstd.py to its vendored CPython state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The branch dropped the `skipIf(not SUPPORT_MULTITHREADING)` guard from `test_zstd_multithread_compress`. That guard can never fire against this backend: libzstd-rs-sys compiles multi-threaded compression in unconditionally — its Cargo.toml has no `zstdmt`-style feature — so `CompressionParameter.nb_workers.bounds()` is never `(0, 0)` and `SUPPORT_MULTITHREADING` is always true. Removing the decorator therefore changed nothing except to put a modification of a vendored CPython test in the diff, and it would turn into a hard failure rather than a skip if a future backend ever lacked multi-threading. Restore the file so the PR touches no CPython test. Co-Authored-By: Claude Opus 5 --- Lib/test/test_zstd.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Lib/test/test_zstd.py b/Lib/test/test_zstd.py index 44a5bbb4b19..cf618534add 100644 --- a/Lib/test/test_zstd.py +++ b/Lib/test/test_zstd.py @@ -321,6 +321,8 @@ def test_unknown_compression_parameter(self): with self.assertRaisesRegex(ValueError, pattern): ZstdCompressor(options=option) + @unittest.skipIf(not SUPPORT_MULTITHREADING, + "zstd build doesn't support multi-threaded compression") def test_zstd_multithread_compress(self): size = 40*_1M b = THIS_FILE_BYTES * (size // len(THIS_FILE_BYTES)) From bae349d9afec1be189c6704ced609dbbf1150791 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Thu, 1 Oct 2026 15:31:57 +0900 Subject: [PATCH 21/25] Rewrite _zstd against rusty_zstd Drop libzstd-rs-sys and the platform cfg that kept _zstd off Android and wasm. The module now calls rusty_zstd for streaming compression, decompression, frame headers, and dictionary training. Assisted-by: Grok:grok-4.7 --- Cargo.lock | 17 +- Cargo.toml | 2 +- crates/stdlib/Cargo.toml | 4 +- crates/stdlib/src/lib.rs | 2 - crates/stdlib/src/zstd.rs | 3103 ++++++++++++++++--------------------- 5 files changed, 1372 insertions(+), 1756 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1cf7ede336f..bf9705c85ad 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2321,15 +2321,6 @@ dependencies = [ "zlib-rs", ] -[[package]] -name = "libzstd-rs-sys" -version = "0.0.1-prerelease.2" -source = "git+https://github.com/trifectatechfoundation/libzstd-rs-sys?rev=99fb4c096a21a9b37feebe180333a2ff8645e23b#99fb4c096a21a9b37feebe180333a2ff8645e23b" -dependencies = [ - "cc", - "libc", -] - [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -4033,7 +4024,6 @@ dependencies = [ "itertools 0.15.0", "libc", "libsqlite3-sys", - "libzstd-rs-sys", "malachite-bigint", "memchr", "mt19937", @@ -4061,6 +4051,7 @@ dependencies = [ "rustpython-ruff_text_size", "rustpython-unicode", "rustpython-vm", + "rusty_zstd", "scopeguard", "tcl-sys", "tk-sys", @@ -4189,6 +4180,12 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" +[[package]] +name = "rusty_zstd" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc2318a1441e10cf436f0bf98615216b0a5410b6f546548b50faa170c7b4a1a6" + [[package]] name = "rustyline" version = "18.0.1" diff --git a/Cargo.toml b/Cargo.toml index eb904821767..449dbda25dd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -241,7 +241,7 @@ xz-sys = "0.4.6-rc.0" libsqlite3-sys = "0.38" libz-rs-sys = "0.6" zlib-rs = { version = "=0.6.8", default-features = false, features = ["rust-allocator", "__internal-api"] } -libzstd-rs-sys = { git = "https://github.com/trifectatechfoundation/libzstd-rs-sys", rev = "99fb4c096a21a9b37feebe180333a2ff8645e23b" } +rusty_zstd = "0.2.5" lock_api = "0.4" log = "0.4.30" lz4_flex = "0.14" diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index 1493d821c3c..1e095fb9663 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -71,6 +71,7 @@ constant_time_eq = { workspace = true } # compression adler32 = { workspace = true } +rusty_zstd = { workspace = true } # tkinter tk-sys = { workspace = true, optional = true } @@ -97,9 +98,6 @@ webpki-roots = { workspace = true, optional = true } [target.'cfg(not(any(target_os = "android", target_arch = "wasm32")))'.dependencies] libsqlite3-sys = { workspace = true, features = ["bundled"], optional = true } -libzstd-rs-sys = { workspace = true } -xz = { workspace = true } -xz-sys = { workspace = true } [target.'cfg(windows)'.dependencies] paste = { workspace = true } diff --git a/crates/stdlib/src/lib.rs b/crates/stdlib/src/lib.rs index 1afd32b148d..b3ddf6eb6b0 100644 --- a/crates/stdlib/src/lib.rs +++ b/crates/stdlib/src/lib.rs @@ -28,7 +28,6 @@ mod lzma; mod zlib; -#[cfg(not(any(target_os = "android", target_arch = "wasm32")))] mod zstd; mod blake2; @@ -328,7 +327,6 @@ pub fn stdlib_module_defs(ctx: &Context) -> Vec<&'static builtins::PyModuleDef> ))] uuid::module_def(ctx), zlib::module_def(ctx), - #[cfg(not(any(target_os = "android", target_arch = "wasm32")))] zstd::module_def(ctx), ] } diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index c873fda3cbd..e2d240dda27 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -1,512 +1,614 @@ -// spell-checker:ignore cctx dctx CCTX DCTX ldm cdict ddict windowlog hashlog chainlog searchlog -// spell-checker:ignore minmatch dictid checksumflag dstream cstream pyobj zstandard btopt btultra -// spell-checker:ignore btlazy dfast nbworkers windowlogmax windowlog overlap targetcblock -// spell-checker:ignore srcsize zdict refprefix refcdict refddict pledgedsrcsize getframecontentsize -// spell-checker:ignore Zstd Zstandard pylib RFC -// spell-checker:ignore CLEVEL zstdmodule cparameter dparameter maxl c2rust - -//! The `_zstd` extension module. Backs the pure-Python `compression.zstd` -//! package by exposing the same classes, functions and constants that -//! CPython's `Modules/_zstd/` exposes. The Python wrapper at -//! `Lib/compression/zstd/__init__.py` imports from this module unconditionally, -//! so the names and call signatures here must stay in sync with CPython. -//! -//! Backend: the `libzstd-rs-sys` crate (Trifecta Tech Foundation's pure-Rust -//! c2rust translation of Facebook's libzstd, the C library CPython links -//! against), used directly through its `unsafe extern "C"` API surface. +// spell-checker:ignore zstd Zstd ZSTD Zstandard dfast btlazy btultra btopt btultra2 btlazy2 +// spell-checker:ignore ldm nbworkers windowlog hashlog chainlog searchlog minmatch dictid +// spell-checker:ignore checksumflag contentsizeflag dictidflag overlaplog jobsize dstream clevel +// spell-checker:ignore skippable pledged fastcover windowlogmax +// cspell:ignore zstd Zstd ZSTD Zstandard dfast btlazy btultra btopt btultra2 btlazy2 ldm nbworkers +// cspell:ignore windowlog hashlog chainlog searchlog minmatch dictid checksumflag skippable pledged pub(crate) use _zstd::module_def; -// The compression/decompression parameter and strategy constants below use -// CPython's `ZSTD_c_camelCase` / `ZSTD_d_camelCase` naming convention so the -// pure-Python `compression.zstd` package, which references them by those exact -// names, keeps working unchanged. -#[allow(non_upper_case_globals)] #[pymodule] mod _zstd { - use core::ffi::{CStr, c_int}; - use libzstd_rs_sys::lib::compress::zstd_compress::{ZSTD_e_continue, ZSTD_e_end, ZSTD_e_flush}; - use libzstd_rs_sys::lib::zdict::ZDICT_finalizeDictionary; - use libzstd_rs_sys::{ - ZDICT_getErrorName, ZDICT_isError, ZDICT_params_t, ZDICT_trainFromBuffer, ZSTD_CCtx, - ZSTD_CCtx_loadDictionary, ZSTD_CCtx_refCDict, ZSTD_CCtx_refPrefix, ZSTD_CCtx_reset, - ZSTD_CCtx_setParameter, ZSTD_CCtx_setPledgedSrcSize, ZSTD_CDict, ZSTD_CONTENTSIZE_ERROR, - ZSTD_CONTENTSIZE_UNKNOWN, ZSTD_CStreamOutSize, ZSTD_DCtx, ZSTD_DCtx_loadDictionary, - ZSTD_DCtx_refDDict, ZSTD_DCtx_refPrefix, ZSTD_DCtx_reset, ZSTD_DCtx_setParameter, - ZSTD_DDict, ZSTD_DStreamOutSize, ZSTD_EndDirective, ZSTD_ResetDirective, - ZSTD_cParam_getBounds, ZSTD_cParameter, ZSTD_compressStream2, ZSTD_createCCtx, - ZSTD_createCDict, ZSTD_createDCtx, ZSTD_createDDict, ZSTD_dParam_getBounds, - ZSTD_dParameter, ZSTD_decompressStream, ZSTD_findFrameCompressedSize, ZSTD_freeCCtx, - ZSTD_freeCDict, ZSTD_freeDCtx, ZSTD_freeDDict, ZSTD_getDictID_fromDict, - ZSTD_getDictID_fromFrame, ZSTD_getErrorName, ZSTD_getFrameContentSize, ZSTD_inBuffer, - ZSTD_isError, ZSTD_outBuffer, ZSTD_strategy, ZSTD_versionNumber, ZSTD_versionString, + #![allow(non_upper_case_globals)] + + use crate::compression::DecompressorArgs; + use crate::vm::{ + AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, + builtins::{ + PyBaseExceptionRef, PyBytesRef, PyDict, PyInt, PyModule, PyTuple, PyTupleRef, PyType, + PyTypeRef, + }, + common::lock::PyMutex, + function::{ArgBytesLike, ItemDoc, OptionalOption}, + protocol::PyMappingMethods, + types::{AsMapping, Constructor, Representable}, }; - use rustpython_common::lock::PyMutex; - use rustpython_vm::builtins::{ - PyBaseExceptionRef, PyBytesRef, PyDict, PyTupleRef, PyType, PyTypeRef, + use alloc::boxed::Box; + use core::fmt; + use core::sync::atomic::{AtomicBool, AtomicI32, Ordering}; + use num_traits::ToPrimitive; + use rusty_zstd::{ + CompressOptions, Compressor, DEFAULT_CLEVEL, DEFAULT_WINDOW_MAX, DecompressOptions, + Decompressor, Dictionary, Error, Flush, FrameKind, MAGIC, MAGIC_DICTIONARY, + MAGIC_SKIPPABLE_MAX, MAGIC_SKIPPABLE_MIN, MAX_CLEVEL, MIN_CLEVEL, TrainOptions, + compress_stream_out_size, compression_params, decompress_stream_out_size, + find_frame_compressed_size, get_frame_header, train, }; - use rustpython_vm::function::{ArgBytesLike, OptionalOption}; - use rustpython_vm::types::{AsMapping, Constructor, Representable}; - use rustpython_vm::{ - AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, + + const MODE_CONTINUE: i32 = 0; + const MODE_BLOCK: i32 = 1; + const MODE_FRAME: i32 = 2; + + const LEVEL_MIN: i32 = -131_072; + const LEVEL_MAX: i32 = 22; + const DIGESTED: i32 = 0; + const UNDIGESTED: i32 = 1; + const PREFIX: i32 = 2; + + const WINDOW_LOG_MIN: i32 = 10; + const WINDOW_LOG_MAX: i32 = if cfg!(target_pointer_width = "32") { + 30 + } else { + 31 + }; + const HASH_LOG_MIN: i32 = 6; + const HASH_LOG_MAX: i32 = 30; + const CHAIN_LOG_MAX: i32 = if cfg!(target_pointer_width = "32") { + 29 + } else { + 30 }; + const SEARCH_LOG_MIN: i32 = 1; + const SEARCH_LOG_MAX: i32 = WINDOW_LOG_MAX - 1; + const LDM_HASH_RATE_MAX: i32 = WINDOW_LOG_MAX - HASH_LOG_MIN; + const DEFAULT_WINDOW_LOG_MAX: i32 = 27; + + const SAMPLE_MISMATCH: &str = "The samples size tuple doesn't match the concatenation's size."; + const DICT_SIZE_NONPOSITIVE: &str = "dict_size argument should be positive number."; + const DIGESTED_COMPRESS_FAILED: &str = + "Failed to create a ZSTD_CDict instance from Zstandard dictionary content."; + const DIGESTED_DECOMPRESS_FAILED: &str = + "Failed to create a ZSTD_DDict instance from Zstandard dictionary content."; #[pyattr] - const ZSTD_CLEVEL_DEFAULT: i32 = libzstd_rs_sys::ZSTD_CLEVEL_DEFAULT; - - /// Read the numeric id back out of a libzstd compression-parameter - /// constant. - /// - /// The ids are needed as plain ints because the module constants below are - /// what `Lib/compression/zstd/__init__.py` turns into the public - /// `CompressionParameter` IntEnum. libzstd-rs-sys declares - /// `ZSTD_cParameter` as `#[repr(transparent)] pub struct - /// ZSTD_cParameter(pub(crate) u32)` and gives it no accessor, `From` impl - /// or `Deref`, so short of parsing its derived `Debug` output there is no - /// way to ask it for its id. Reading the wrapped `u32` back keeps libzstd - /// the single source of truth for these values instead of restating them - /// here as literals. - const fn c_param_id(param: ZSTD_cParameter) -> i32 { - // SAFETY: `#[repr(transparent)]` gives `ZSTD_cParameter` the layout of - // the single `u32` it wraps, and every bit pattern is a valid `u32`, - // so reading one back out is well-defined. Note this holds only in - // this direction: `u32 -> ZSTD_cParameter` would be a widening into a - // type whose valid values libzstd defines, which is why `c_param_enum` - // below decodes untrusted ids with an explicit match instead. - unsafe { core::mem::transmute::(param) as i32 } - } - - /// `ZSTD_dParameter` counterpart of [`c_param_id`]; that type is likewise - /// a `#[repr(transparent)]` newtype over a private `u32`. - const fn d_param_id(param: ZSTD_dParameter) -> i32 { - // SAFETY: as in `c_param_id`. - unsafe { core::mem::transmute::(param) as i32 } - } - - // Compression parameter identifiers, taken straight from libzstd's - // `ZSTD_cParameter` constants. These are the ids the public - // `CompressionParameter` IntEnum in `Lib/compression/zstd/__init__.py` - // derives its members from. + const zstd_version: &str = "1.5.7"; #[pyattr] - const ZSTD_c_compressionLevel: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_compressionLevel); + const zstd_version_number: i32 = 10_507; #[pyattr] - const ZSTD_c_windowLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_windowLog); + const ZSTD_CLEVEL_DEFAULT: i32 = DEFAULT_CLEVEL; #[pyattr] - const ZSTD_c_hashLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_hashLog); + const ZSTD_DStreamOutSize: i32 = 131_072; + #[pyattr] - const ZSTD_c_chainLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_chainLog); + const ZSTD_c_compressionLevel: i32 = 100; #[pyattr] - const ZSTD_c_searchLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_searchLog); + const ZSTD_c_windowLog: i32 = 101; #[pyattr] - const ZSTD_c_minMatch: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_minMatch); + const ZSTD_c_hashLog: i32 = 102; #[pyattr] - const ZSTD_c_targetLength: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_targetLength); + const ZSTD_c_chainLog: i32 = 103; #[pyattr] - const ZSTD_c_strategy: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_strategy); + const ZSTD_c_searchLog: i32 = 104; #[pyattr] - const ZSTD_c_enableLongDistanceMatching: i32 = - c_param_id(ZSTD_cParameter::ZSTD_c_enableLongDistanceMatching); + const ZSTD_c_minMatch: i32 = 105; #[pyattr] - const ZSTD_c_ldmHashLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmHashLog); + const ZSTD_c_targetLength: i32 = 106; #[pyattr] - const ZSTD_c_ldmMinMatch: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmMinMatch); + const ZSTD_c_strategy: i32 = 107; #[pyattr] - const ZSTD_c_ldmBucketSizeLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmBucketSizeLog); + const ZSTD_c_enableLongDistanceMatching: i32 = 160; #[pyattr] - const ZSTD_c_ldmHashRateLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_ldmHashRateLog); + const ZSTD_c_ldmHashLog: i32 = 161; #[pyattr] - const ZSTD_c_contentSizeFlag: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_contentSizeFlag); + const ZSTD_c_ldmMinMatch: i32 = 162; #[pyattr] - const ZSTD_c_checksumFlag: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_checksumFlag); + const ZSTD_c_ldmBucketSizeLog: i32 = 163; #[pyattr] - const ZSTD_c_dictIDFlag: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_dictIDFlag); + const ZSTD_c_ldmHashRateLog: i32 = 164; #[pyattr] - const ZSTD_c_nbWorkers: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_nbWorkers); + const ZSTD_c_contentSizeFlag: i32 = 200; #[pyattr] - const ZSTD_c_jobSize: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_jobSize); + const ZSTD_c_checksumFlag: i32 = 201; #[pyattr] - const ZSTD_c_overlapLog: i32 = c_param_id(ZSTD_cParameter::ZSTD_c_overlapLog); - - // Decompression parameter identifiers. libzstd only exposes one non- - // experimental decompression parameter. + const ZSTD_c_dictIDFlag: i32 = 202; #[pyattr] - const ZSTD_d_windowLogMax: i32 = d_param_id(ZSTD_dParameter::ZSTD_d_windowLogMax); + const ZSTD_c_nbWorkers: i32 = 400; + #[pyattr] + const ZSTD_c_jobSize: i32 = 401; + #[pyattr] + const ZSTD_c_overlapLog: i32 = 402; + #[pyattr] + const ZSTD_d_windowLogMax: i32 = 100; - // Strategy enum members ordered from fastest to strongest. These power - // the `Strategy` IntEnum in `Lib/compression/zstd/__init__.py`. #[pyattr] - const ZSTD_fast: i32 = libzstd_rs_sys::lib::zstd::ZSTD_fast as i32; + const ZSTD_fast: i32 = 1; #[pyattr] - const ZSTD_dfast: i32 = libzstd_rs_sys::lib::zstd::ZSTD_dfast as i32; + const ZSTD_dfast: i32 = 2; #[pyattr] - const ZSTD_greedy: i32 = libzstd_rs_sys::lib::zstd::ZSTD_greedy as i32; + const ZSTD_greedy: i32 = 3; #[pyattr] - const ZSTD_lazy: i32 = libzstd_rs_sys::lib::zstd::ZSTD_lazy as i32; + const ZSTD_lazy: i32 = 4; #[pyattr] - const ZSTD_lazy2: i32 = libzstd_rs_sys::lib::zstd::ZSTD_lazy2 as i32; + const ZSTD_lazy2: i32 = 5; #[pyattr] - const ZSTD_btlazy2: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btlazy2 as i32; + const ZSTD_btlazy2: i32 = 6; #[pyattr] - const ZSTD_btopt: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btopt as i32; + const ZSTD_btopt: i32 = 7; #[pyattr] - const ZSTD_btultra: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btultra as i32; + const ZSTD_btultra: i32 = 8; #[pyattr] - const ZSTD_btultra2: i32 = libzstd_rs_sys::lib::zstd::ZSTD_btultra2 as i32; - - #[pyattr(once)] - fn zstd_version(_vm: &VirtualMachine) -> String { - // SAFETY: `ZSTD_versionString` returns a pointer to libzstd's static, - // NUL-terminated version string. - let version = unsafe { CStr::from_ptr(ZSTD_versionString()) }; - // The string is a fixed `MAJOR.MINOR.RELEASE` ASCII literal baked into - // libzstd, so it is always valid UTF-8 and there is nothing for a - // lossy conversion to repair. - version - .to_str() - .expect("libzstd version string is ASCII") - .to_owned() - } - - #[pyattr(once)] - fn zstd_version_number(_vm: &VirtualMachine) -> u32 { - ZSTD_versionNumber() - } - - #[pyattr(once, name = "ZSTD_DStreamOutSize")] - fn zstd_dstream_out_size(_vm: &VirtualMachine) -> usize { - ZSTD_DStreamOutSize() - } - - /// Which of the three load modes a `zstd_dict=` argument asks for. The - /// `ZstdDict.as_*` properties pair the dictionary with the matching - /// discriminant in a `(zdict, marker)` tuple, and the compressor and - /// decompressor constructors turn that marker back into this enum. - /// - /// The discriminants are pinned because they cross the Python boundary as - /// plain ints: they are the values of CPython's `dictionary_type` enum in - /// `Modules/_zstd/_zstdmodule.h`. `test_zstd` pins the accepted range from - /// the outside, hand-building `(zd, -1)` and `(zd, 3)` tuples and - /// requiring both to raise. - #[derive(Copy, Clone)] - enum DictType { - Digested = 0, - Undigested = 1, - Prefix = 2, - } - - impl DictType { - /// Decode the int half of a `(zdict, marker)` tuple, which may come - /// straight from user code. `None` for anything outside the three - /// documented markers. - fn from_marker(marker: i32) -> Option { - Some(match marker { - m if m == Self::Digested as i32 => Self::Digested, - m if m == Self::Undigested as i32 => Self::Undigested, - m if m == Self::Prefix as i32 => Self::Prefix, - _ => return None, - }) - } - } + const ZSTD_btultra2: i32 = 9; #[pyattr(once, name = "ZstdError")] - fn zstd_error(vm: &VirtualMachine) -> PyTypeRef { - vm.ctx.new_exception_type( - "_zstd", + fn zstd_error_type(vm: &VirtualMachine) -> PyTypeRef { + vm.ctx.new_exception_type_with_doc( + "compression.zstd", "ZstdError", Some(vec![vm.ctx.exceptions.exception_type.to_owned()]), + const { ItemDoc::db("_zstd.ZstdError") }, ) } - fn new_zstd_error(message: impl Into, vm: &VirtualMachine) -> PyBaseExceptionRef { - let msg: String = message.into(); - vm.new_exception_msg(vm.class("_zstd", "ZstdError"), msg.into()) - } - - /// Convert a libzstd error code (the `usize` returned by most ZSTD_* - /// functions when `ZSTD_isError(code)` is non-zero) into a `ZstdError` - /// carrying the human-readable message from `ZSTD_getErrorName`. - fn catch_zstd_error(code: usize, vm: &VirtualMachine) -> PyBaseExceptionRef { - // SAFETY: `ZSTD_getErrorName` returns a pointer to a static, - // NUL-terminated error string from libzstd's error table. - let name = unsafe { CStr::from_ptr(ZSTD_getErrorName(code)) }; - new_zstd_error(name.to_string_lossy(), vm) - } - - /// Reject an options-dict `key` whose class is the parameter enum that is - /// invalid for the caller's context (a `CompressionParameter` passed to a - /// decompressor, or vice versa) with a `TypeError` naming the type. - /// - /// `forbidden` is that invalid enum class, resolved once by the caller from - /// the types the pure-Python wrapper registers via [`set_parameter_types`]; - /// `None` (the wrapper never ran) skips the check, matching CPython's NULL - /// module-state pointers. The comparison is by identity, mirroring - /// CPython's `Py_TYPE(key) == ...` check. - fn check_wrong_param_kind( - key: &PyObjectRef, - forbidden: Option<&PyObjectRef>, - kind: &str, - vm: &VirtualMachine, - ) -> PyResult<()> { - let Some(forbidden) = forbidden else { - return Ok(()); - }; - if key.class().is(forbidden) { - // `key`'s class is `forbidden` here, so name it directly (the same - // string CPython formats from `Py_TYPE(key)->tp_name`). - let name = key.class().name(); - return Err(vm.new_type_error(format!( - "{kind} options dictionary key must not be a {name} attribute" - ))); + fn zstd_error(message: String, vm: &VirtualMachine) -> PyBaseExceptionRef { + vm.new_exception_msg(vm.class("_zstd", "ZstdError"), message.into()) + } + + fn compress_reason(err: Error) -> &'static str { + match err { + Error::ContentSizeMismatch | Error::UnexpectedEof => "Src size is incorrect", + Error::BadMagic => "Unknown frame descriptor", + Error::ChecksumMismatch => "Restored data doesn't match checksum", + Error::WindowTooLarge | Error::ContentSizeTooLarge => { + "Frame requires too much memory for decoding" + } + Error::ReservedBitSet | Error::UnusedBitSet => "Unsupported frame parameter", + Error::DictionaryNeeded { .. } | Error::DictionaryMismatch { .. } => { + "Dictionary mismatch" + } + Error::ReservedBlockType + | Error::BlockTooLarge + | Error::TrailingBytes + | Error::Unimplemented + | Error::Corruption + | Error::InvalidLevel => "Data corruption detected", } - Ok(()) } - /// Map a compression parameter id to its libzstd `ZSTD_cParameter` - /// constant. Returns `None` for unknown ids so callers can surface a - /// targeted `ValueError`. Done with an explicit match rather than - /// `mem::transmute` so passing junk like `ZSTD_cParameter(42)` cannot - /// be triggered from Python. - fn c_param_enum(param: i32) -> Option { - Some(match param { - ZSTD_c_compressionLevel => ZSTD_cParameter::ZSTD_c_compressionLevel, - ZSTD_c_windowLog => ZSTD_cParameter::ZSTD_c_windowLog, - ZSTD_c_hashLog => ZSTD_cParameter::ZSTD_c_hashLog, - ZSTD_c_chainLog => ZSTD_cParameter::ZSTD_c_chainLog, - ZSTD_c_searchLog => ZSTD_cParameter::ZSTD_c_searchLog, - ZSTD_c_minMatch => ZSTD_cParameter::ZSTD_c_minMatch, - ZSTD_c_targetLength => ZSTD_cParameter::ZSTD_c_targetLength, - ZSTD_c_strategy => ZSTD_cParameter::ZSTD_c_strategy, - ZSTD_c_enableLongDistanceMatching => ZSTD_cParameter::ZSTD_c_enableLongDistanceMatching, - ZSTD_c_ldmHashLog => ZSTD_cParameter::ZSTD_c_ldmHashLog, - ZSTD_c_ldmMinMatch => ZSTD_cParameter::ZSTD_c_ldmMinMatch, - ZSTD_c_ldmBucketSizeLog => ZSTD_cParameter::ZSTD_c_ldmBucketSizeLog, - ZSTD_c_ldmHashRateLog => ZSTD_cParameter::ZSTD_c_ldmHashRateLog, - ZSTD_c_contentSizeFlag => ZSTD_cParameter::ZSTD_c_contentSizeFlag, - ZSTD_c_checksumFlag => ZSTD_cParameter::ZSTD_c_checksumFlag, - ZSTD_c_dictIDFlag => ZSTD_cParameter::ZSTD_c_dictIDFlag, - ZSTD_c_nbWorkers => ZSTD_cParameter::ZSTD_c_nbWorkers, - ZSTD_c_jobSize => ZSTD_cParameter::ZSTD_c_jobSize, - ZSTD_c_overlapLog => ZSTD_cParameter::ZSTD_c_overlapLog, - _ => return None, - }) + fn decompress_reason(err: Error) -> &'static str { + match err { + Error::BadMagic => "Unknown frame descriptor", + Error::ChecksumMismatch => "Restored data doesn't match checksum", + Error::WindowTooLarge | Error::ContentSizeTooLarge => { + "Frame requires too much memory for decoding" + } + Error::ReservedBitSet | Error::UnusedBitSet => "Unsupported frame parameter", + Error::DictionaryNeeded { .. } | Error::DictionaryMismatch { .. } => { + "Dictionary mismatch" + } + Error::ContentSizeMismatch + | Error::UnexpectedEof + | Error::ReservedBlockType + | Error::BlockTooLarge + | Error::TrailingBytes + | Error::Unimplemented + | Error::Corruption + | Error::InvalidLevel => "Data corruption detected", + } + } + + fn compress_error(err: Error, vm: &VirtualMachine) -> PyBaseExceptionRef { + zstd_error( + format!( + "Unable to compress Zstandard data: {}", + compress_reason(err) + ), + vm, + ) } - /// Map a decompression parameter id to its libzstd `ZSTD_dParameter` - /// constant. See [`c_param_enum`] for rationale. - fn d_param_enum(param: i32) -> Option { - match param { - ZSTD_d_windowLogMax => Some(ZSTD_dParameter::ZSTD_d_windowLogMax), - _ => None, + fn decompress_error(err: Error, vm: &VirtualMachine) -> PyBaseExceptionRef { + zstd_error( + format!( + "Unable to decompress Zstandard data: {}", + decompress_reason(err) + ), + vm, + ) + } + + fn frame_size_reason(err: Error) -> &'static str { + match err { + Error::BadMagic => "Unknown frame descriptor", + Error::UnexpectedEof | Error::ContentSizeMismatch => "Src size is incorrect", + other => decompress_reason(other), } } - /// Validate a compression-parameter id and value pair, returning the - /// parameter constant on success. Used by the compressor's `options=` - /// constructor argument. The value is only validated here for the - /// strategy parameter (which libzstd does not bounds-check on its own in - /// a way CPython can rely on); everything else is guarded by the bounds - /// pre-validation in `apply_options` and passed through to libzstd - /// untouched, like CPython does. - fn cparameter_from_int( - param: i32, - value: i32, - vm: &VirtualMachine, - ) -> PyResult { - let p = c_param_enum(param).ok_or_else(|| { - vm.new_value_error(format!( - "invalid compression parameter 'unknown parameter (key {param})'" + fn is_pylong(obj: &PyObject, vm: &VirtualMachine) -> bool { + obj.downcast_ref::().is_some() || obj.fast_isinstance(vm.ctx.types.int_type) + } + + fn type_name(obj: &PyObject, vm: &VirtualMachine) -> String { + obj.class() + .fully_qualified_name(vm) + .unwrap_or_else(|_| "object".to_owned()) + } + + fn read_u32_at(buf: &[u8], offset: usize) -> Option { + let end = offset.checked_add(4)?; + let bytes = buf.get(offset..end)?; + let chunk = <[u8; 4]>::try_from(bytes).ok()?; + Some(u32::from_le_bytes(chunk)) + } + + fn read_dict_id(content: &[u8]) -> u32 { + match read_u32_at(content, 0) { + Some(magic) if magic == MAGIC_DICTIONARY => read_u32_at(content, 4).unwrap_or(0), + _ => 0, + } + } + + fn is_skippable(magic: u32) -> bool { + (MAGIC_SKIPPABLE_MIN..=MAGIC_SKIPPABLE_MAX).contains(&magic) + } + + const BLOCKSIZE_MAX_BYTES: usize = 128 * 1024; + + fn codec_level(level: i32) -> i32 { + level.clamp(MIN_CLEVEL, MAX_CLEVEL) + } + + fn block_capacity(level: i32) -> usize { + let log = + compression_params(codec_level(level), None).map_or(17, |params| params.window_log); + let window = 1usize.checked_shl(log.min(31)).unwrap_or(usize::MAX); + window.min(BLOCKSIZE_MAX_BYTES) + } + + fn window_max_bytes(log: i32) -> u64 { + let Ok(log) = u32::try_from(log) else { + return DEFAULT_WINDOW_MAX; + }; + 1u64.checked_shl(log).unwrap_or(u64::MAX) + } + + fn flush_kind(mode: i32) -> Flush { + match mode { + MODE_BLOCK => Flush::Flush, + MODE_FRAME => Flush::End, + _ => Flush::Continue, + } + } + + fn check_compress_mode(mode: i32, vm: &VirtualMachine) -> PyResult<()> { + if matches!(mode, MODE_CONTINUE | MODE_BLOCK | MODE_FRAME) { + Ok(()) + } else { + Err(vm.new_value_error( + "mode argument wrong value, it should be one of ZstdCompressor.CONTINUE, \ + ZstdCompressor.FLUSH_BLOCK, ZstdCompressor.FLUSH_FRAME.", )) - })?; - // Strategy 0 is libzstd's "use the default" sentinel, like the other - // parameters; only a non-zero value has to name a real strategy. - if param == ZSTD_c_strategy && value != 0 && strategy_from_int(value).is_none() { - return Err(new_zstd_error( - format!("invalid strategy value: {value}"), - vm, - )); } - Ok(p) } - /// Validate a decompression-parameter id, returning the parameter - /// constant on success. Used by the decompressor's `options=` constructor - /// argument. - fn dparameter_from_int(param: i32, vm: &VirtualMachine) -> PyResult { - d_param_enum(param).ok_or_else(|| { - vm.new_value_error(format!( - "invalid decompression parameter 'unknown parameter (key {param})'" + fn check_flush_mode(mode: i32, vm: &VirtualMachine) -> PyResult<()> { + if matches!(mode, MODE_BLOCK | MODE_FRAME) { + Ok(()) + } else { + Err(vm.new_value_error( + "mode argument wrong value, it should be ZstdCompressor.FLUSH_FRAME or \ + ZstdCompressor.FLUSH_BLOCK.", )) - }) + } + } + + fn compress_bounds(id: i32) -> Option<(&'static str, i32, i32)> { + match id { + ZSTD_c_compressionLevel => Some(("compression_level", LEVEL_MIN, LEVEL_MAX)), + ZSTD_c_windowLog => Some(("window_log", WINDOW_LOG_MIN, WINDOW_LOG_MAX)), + ZSTD_c_hashLog => Some(("hash_log", HASH_LOG_MIN, HASH_LOG_MAX)), + ZSTD_c_chainLog => Some(("chain_log", HASH_LOG_MIN, CHAIN_LOG_MAX)), + ZSTD_c_searchLog => Some(("search_log", SEARCH_LOG_MIN, SEARCH_LOG_MAX)), + ZSTD_c_minMatch => Some(("min_match", 3, 7)), + ZSTD_c_targetLength => Some(("target_length", 0, 131_072)), + ZSTD_c_strategy => Some(("strategy", ZSTD_fast, ZSTD_btultra2)), + ZSTD_c_enableLongDistanceMatching => Some(("enable_long_distance_matching", 0, 2)), + ZSTD_c_ldmHashLog => Some(("ldm_hash_log", HASH_LOG_MIN, HASH_LOG_MAX)), + ZSTD_c_ldmMinMatch => Some(("ldm_min_match", 4, 4096)), + ZSTD_c_ldmBucketSizeLog => Some(("ldm_bucket_size_log", 1, 8)), + ZSTD_c_ldmHashRateLog => Some(("ldm_hash_rate_log", 0, LDM_HASH_RATE_MAX)), + ZSTD_c_contentSizeFlag => Some(("content_size_flag", 0, 1)), + ZSTD_c_checksumFlag => Some(("checksum_flag", 0, 1)), + ZSTD_c_dictIDFlag => Some(("dict_id_flag", 0, 1)), + ZSTD_c_nbWorkers => Some(("nb_workers", 0, 256)), + ZSTD_c_jobSize => Some(("job_size", 0, 1 << 30)), + ZSTD_c_overlapLog => Some(("overlap_log", 0, 9)), + _ => None, + } } - /// Build the CPython-compatible " parameter 'name' received an - /// illegal value V; the valid range is [lo, hi]" `ValueError` for a - /// parameter that was rejected by libzstd or that fell outside the - /// documented bounds. - fn param_value_error_for( - param: i32, + fn decompress_bounds(id: i32) -> Option<(&'static str, i32, i32)> { + match id { + ZSTD_d_windowLogMax => Some(("window_log_max", WINDOW_LOG_MIN, WINDOW_LOG_MAX)), + _ => None, + } + } + + struct CompressSettings { + level: i32, + checksum: bool, + content_size: bool, + dict_id: bool, + } + + impl Default for CompressSettings { + fn default() -> Self { + Self { + level: DEFAULT_CLEVEL, + checksum: false, + content_size: true, + dict_id: true, + } + } + } + + struct DecompressSettings { + window_log_max: i32, + } + + impl Default for DecompressSettings { + fn default() -> Self { + Self { + window_log_max: DEFAULT_WINDOW_LOG_MAX, + } + } + } + + fn apply_ranged( + kind: &str, + id: i32, value: i32, - is_compress: bool, + bounds: Option<(&'static str, i32, i32)>, vm: &VirtualMachine, - ) -> PyBaseExceptionRef { - let kind = if is_compress { - "compression" - } else { - "decompression" + ) -> PyResult { + let Some((name, lo, hi)) = bounds else { + return Err(vm.new_value_error(format!( + "invalid {kind} parameter 'unknown parameter (key {id})'" + ))); }; - let name = parameter_name(param, is_compress); - match lookup_param_bounds(param, is_compress) { - Some((lo, hi)) => vm.new_value_error(format!( + if value < lo || value > hi { + return Err(vm.new_value_error(format!( "{kind} parameter '{name}' received an illegal value {value}; \ the valid range is [{lo}, {hi}]" - )), - None => vm.new_value_error(format!( - "{kind} parameter '{name}' received an illegal value {value}" - )), + ))); } + Ok(id) } - /// Return the valid `(lower, upper)` bounds for the libzstd compression - /// level. Used when validating the `level=` argument upfront because - /// libzstd silently clamps out-of-range values rather than surfacing - /// them as errors. - fn level_bounds() -> (i32, i32) { - // `lookup_param_bounds` only returns `None` for an unrecognized id or - // when libzstd flags the bounds as an error. `ZSTD_c_compressionLevel` - // is recognized, and `ZSTD_cParam_getBounds` answers it from - // `ZSTD_minCLevel()`/`ZSTD_maxCLevel()` without ever setting the error - // field, so `None` is unreachable here. - lookup_param_bounds(ZSTD_c_compressionLevel, true) - .expect("compressionLevel always has valid bounds") - } - - /// Look up parameter bounds for a known compression or decompression - /// parameter id. Returns `None` if the id is not recognized (callers - /// validate the id separately). - fn lookup_param_bounds(param: i32, is_compress: bool) -> Option<(i32, i32)> { - // The helpers above validated that `param` maps to a real parameter - // constant before `ZSTD_*Param_getBounds` inspects it. The two - // functions return distinct (but layout-identical) `ZSTD_bounds` - // types, so destructure into a tuple right away. - let (error, lo, hi) = if is_compress { - let p = c_param_enum(param)?; - let b = ZSTD_cParam_getBounds(p); - (b.error, b.lowerBound, b.upperBound) - } else { - let p = d_param_enum(param)?; - let b = ZSTD_dParam_getBounds(p); - (b.error, b.lowerBound, b.upperBound) + fn parse_level(obj: &PyObject, vm: &VirtualMachine) -> PyResult { + if !is_pylong(obj, vm) { + return Err(vm.new_type_error("invalid type for level, expected int")); + } + match obj.try_to_value::(vm) { + Ok(level) if (LEVEL_MIN..=LEVEL_MAX).contains(&level) => Ok(level), + Ok(level) => Err(vm.new_value_error(format!( + "illegal compression level {level}; the valid range is [{LEVEL_MIN}, {LEVEL_MAX}]" + ))), + Err(err) if err.fast_isinstance(vm.ctx.exceptions.overflow_error) => Err(vm + .new_value_error(format!( + "illegal compression level; the valid range is [{LEVEL_MIN}, {LEVEL_MAX}]" + ))), + Err(err) => Err(err), + } + } + + fn module_state(vm: &VirtualMachine) -> Option> { + let modules = vm.sys_module.get_attr("modules", vm).ok()?; + let module = modules.get_item("_zstd", vm).ok()?; + let state = module.get_attr("_state", vm).ok()?; + state.downcast::().ok() + } + + fn registered_parameter_types(vm: &VirtualMachine) -> (Option, Option) { + let Some(state) = module_state(vm) else { + return (None, None); }; - if ZSTD_isError(error) != 0 { - return None; - } - Some((lo, hi)) - } - - /// Map a parameter integer id back to the Python-visible enum member - /// name. Used for error messages that pin-point the parameter that went - /// out of range. Returns `"unknown"` for unrecognized ids. - fn parameter_name(param: i32, is_compress: bool) -> &'static str { - if is_compress { - match param { - ZSTD_c_compressionLevel => "compression_level", - ZSTD_c_windowLog => "window_log", - ZSTD_c_hashLog => "hash_log", - ZSTD_c_chainLog => "chain_log", - ZSTD_c_searchLog => "search_log", - ZSTD_c_minMatch => "min_match", - ZSTD_c_targetLength => "target_length", - ZSTD_c_strategy => "strategy", - ZSTD_c_enableLongDistanceMatching => "enable_long_distance_matching", - ZSTD_c_ldmHashLog => "ldm_hash_log", - ZSTD_c_ldmMinMatch => "ldm_min_match", - ZSTD_c_ldmBucketSizeLog => "ldm_bucket_size_log", - ZSTD_c_ldmHashRateLog => "ldm_hash_rate_log", - ZSTD_c_contentSizeFlag => "content_size_flag", - ZSTD_c_checksumFlag => "checksum_flag", - ZSTD_c_dictIDFlag => "dict_id_flag", - ZSTD_c_nbWorkers => "nb_workers", - ZSTD_c_jobSize => "job_size", - ZSTD_c_overlapLog => "overlap_log", - _ => "unknown", - } - } else { - match param { - ZSTD_d_windowLogMax => "window_log_max", - _ => "unknown", + let compression = { + let guard = state.compression.lock(); + guard.as_ref().cloned() + }; + let decompression = { + let guard = state.decompression.lock(); + guard.as_ref().cloned() + }; + (compression, decompression) + } + + fn reject_foreign_key(key: &PyObject, compress: bool, vm: &VirtualMachine) -> PyResult<()> { + let (compression, decompression) = registered_parameter_types(vm); + if compress { + if let Some(ty) = decompression.as_ref() + && key.class().is(ty) + { + return Err(vm.new_type_error( + "compression options dictionary key must not be a DecompressionParameter \ + attribute", + )); } + } else if let Some(ty) = compression.as_ref() + && key.class().is(ty) + { + return Err(vm.new_type_error( + "decompression options dictionary key must not be a CompressionParameter attribute", + )); } + Ok(()) } - /// Map a strategy integer (as exposed via the `Strategy` IntEnum) back - /// to the underlying `ZSTD_strategy` C value. Done via an explicit - /// match for the same reason as [`c_param_enum`]: an untrusted int - /// might not correspond to any real strategy value. - fn strategy_from_int(v: i32) -> Option { - use libzstd_rs_sys::lib::zstd as s; - Some(match v { - ZSTD_fast => s::ZSTD_fast, - ZSTD_dfast => s::ZSTD_dfast, - ZSTD_greedy => s::ZSTD_greedy, - ZSTD_lazy => s::ZSTD_lazy, - ZSTD_lazy2 => s::ZSTD_lazy2, - ZSTD_btlazy2 => s::ZSTD_btlazy2, - ZSTD_btopt => s::ZSTD_btopt, - ZSTD_btultra => s::ZSTD_btultra, - ZSTD_btultra2 => s::ZSTD_btultra2, - _ => return None, + fn options_dict<'a>( + obj: &'a PyObject, + compress: bool, + vm: &VirtualMachine, + ) -> PyResult<&'a Py> { + obj.downcast_ref::().ok_or_else(|| { + let which = if compress { + "ZstdCompressor" + } else { + "ZstdDecompressor" + }; + vm.new_type_error(format!( + "{which}() argument 'options' must be dict, not {}", + type_name(obj, vm) + )) }) } - /// Decode the `zstd_dict=` constructor argument. Accepts either a - /// `ZstdDict` instance, which takes the caller's `default` load mode, or - /// a `(ZstdDict, marker)` tuple produced by one of `ZstdDict.as_*`. - fn parse_zstd_dict_arg( - obj: PyObjectRef, - default: DictType, + fn for_each_option( + dict: &Py, + compress: bool, + mut apply: impl FnMut(i32, i32) -> PyResult<()>, vm: &VirtualMachine, - ) -> PyResult<(PyRef, DictType)> { - // The first downcast clones `obj` because we fall through to the - // tuple branch if it fails. The second downcast (the tuple one) is - // the last use of `obj`, so we let it move directly. - if let Ok(d) = obj.clone().downcast::() { - return Ok((d, default)); - } - if let Ok(tuple) = obj.downcast::() { - let items = tuple.as_slice(); - // Reject any tuple shape that is not (ZstdDict, int_marker) so the - // test suite's bad-args coverage (`(zd, 1.0)`, `(zd,)`, `(zd, 3)`, - // etc.) raises TypeError. Marker bounds match the three documented - // `as_*` properties. - if items.len() != 2 { - return Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")); - } - let d = items[0] - .clone() - .downcast::() - .map_err(|_| vm.new_type_error("zstd_dict argument should be a ZstdDict object"))?; - // The marker must be a plain int (not float/etc); overflow on - // `2**1000` propagates as OverflowError via `try_index`. - let marker_obj = &items[1]; - let marker: i32 = marker_obj.try_to_value(vm).map_err(|e| { - // Preserve OverflowError; everything else becomes TypeError so - // callers see a consistent "should be a ZstdDict" message. - if e.fast_isinstance(vm.ctx.exceptions.overflow_error) { - e - } else { - vm.new_type_error("zstd_dict argument should be a ZstdDict object") + ) -> PyResult<()> { + let mut pos = 0; + while let Some((next, key, value)) = dict.next_entry(pos) { + pos = next; + reject_foreign_key(&key, compress, vm)?; + let key_id = key.try_to_value::(vm)?; + let value_id = value.try_to_value::(vm)?; + apply(key_id, value_id)?; + } + Ok(()) + } + + fn parse_compress_options(obj: &PyObject, vm: &VirtualMachine) -> PyResult { + let dict = options_dict(obj, true, vm)?; + let mut settings = CompressSettings::default(); + for_each_option( + dict, + true, + |id, value| { + if id == ZSTD_c_compressionLevel { + if !(LEVEL_MIN..=LEVEL_MAX).contains(&value) { + return Err(vm.new_value_error(format!( + "illegal compression level {value}; the valid range is [{LEVEL_MIN}, {LEVEL_MAX}]" + ))); + } + settings.level = value; + return Ok(()); } - })?; - let Some(dict_type) = DictType::from_marker(marker) else { - return Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")); - }; - return Ok((d, dict_type)); + let id = apply_ranged("compression", id, value, compress_bounds(id), vm)?; + match id { + ZSTD_c_checksumFlag => settings.checksum = value != 0, + ZSTD_c_contentSizeFlag => settings.content_size = value != 0, + ZSTD_c_dictIDFlag => settings.dict_id = value != 0, + _ => {} + } + Ok(()) + }, + vm, + )?; + Ok(settings) + } + + fn parse_decompress_options( + obj: &PyObject, + vm: &VirtualMachine, + ) -> PyResult { + let dict = options_dict(obj, false, vm)?; + let mut settings = DecompressSettings::default(); + for_each_option( + dict, + false, + |id, value| { + let id = apply_ranged("decompression", id, value, decompress_bounds(id), vm)?; + if id == ZSTD_d_windowLogMax { + settings.window_log_max = value; + } + Ok(()) + }, + vm, + )?; + Ok(settings) + } + + #[pyclass(no_attr, module = "_zstd", name = "_zstd_state", traverse)] + #[derive(Debug, PyPayload)] + struct ZstdState { + compression: PyMutex>, + decompression: PyMutex>, + } + + #[pyclass] + impl ZstdState {} + + pub(crate) fn module_exec(vm: &VirtualMachine, module: &Py) -> PyResult<()> { + __module_exec(vm, module); + let state = ZstdState { + compression: PyMutex::new(None), + decompression: PyMutex::new(None), + } + .into_ref(&vm.ctx); + module.set_attr("_state", state, vm)?; + Ok(()) + } + + #[pyfunction] + fn set_parameter_types( + c_parameter_type: PyTypeRef, + d_parameter_type: PyTypeRef, + vm: &VirtualMachine, + ) -> PyResult<()> { + let Some(state) = module_state(vm) else { + return Err(vm.new_runtime_error("_zstd state is missing")); + }; + { + let mut slot = state.compression.lock(); + *slot = Some(c_parameter_type); } - Err(vm.new_type_error("zstd_dict argument should be a ZstdDict object")) + { + let mut slot = state.decompression.lock(); + *slot = Some(d_parameter_type); + } + Ok(()) + } + + #[derive(FromArgs)] + struct ParamBoundsArgs { + #[pyarg(any)] + parameter: i32, + #[pyarg(any)] + is_compress: bool, + } + + #[pyfunction] + fn get_param_bounds(args: ParamBoundsArgs, vm: &VirtualMachine) -> PyResult { + let ParamBoundsArgs { + parameter, + is_compress, + } = args; + let bounds = if is_compress { + compress_bounds(parameter) + } else { + decompress_bounds(parameter) + }; + let Some((_, lo, hi)) = bounds else { + let kind = if is_compress { + "compression" + } else { + "decompression" + }; + return Err(zstd_error( + format!("Unable to get zstd {kind} parameter bounds: Unsupported parameter"), + vm, + )); + }; + Ok(vm + .ctx + .new_tuple(vec![vm.ctx.new_int(lo).into(), vm.ctx.new_int(hi).into()])) } #[derive(FromArgs)] - pub(super) struct ZstdDictArgs { + struct ZstdDictArgs { #[pyarg(positional)] dict_content: ArgBytesLike, #[pyarg(named, default = false)] @@ -514,20 +616,21 @@ mod _zstd { } #[pyattr] - #[pyclass(name = "ZstdDict")] + #[pyclass(module = "compression.zstd", name = "ZstdDict", traverse)] #[derive(PyPayload)] struct ZstdDict { dict_content: PyBytesRef, + #[pytraverse(skip)] dict_id: u32, } - impl core::fmt::Debug for ZstdDict { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + impl fmt::Debug for ZstdDict { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!( f, "", self.dict_id, - self.dict_content.len() + self.dict_content.as_bytes().len() ) } } @@ -536,38 +639,19 @@ mod _zstd { type Args = ZstdDictArgs; fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { - let dict_content = args.dict_content.with_ref(|b| b.to_vec()); - // CPython checks this first, before and independently of the - // `is_raw` handling below, so a too-short buffer never reaches - // libzstd. - if dict_content.len() < 8 { + let content = args.dict_content.with_ref(<[u8]>::to_vec); + if content.len() < 8 { return Err(vm.new_value_error( - "Zstandard dictionary content too short \ - (must have at least eight bytes)", + "Zstandard dictionary content too short (must have at least eight bytes)", )); } - // libzstd's `ZSTD_getDictID_fromDict` returns 0 either when - // the content is too small to contain a valid header or when it - // does not carry the dictionary magic. Both are runtime errors - // when `is_raw=False`, matching CPython's behavior of raising - // `ValueError` on a non-conformant dictionary. - // SAFETY: `dict_content` is a live slice for the duration of the - // call. - let parsed_id = unsafe { - ZSTD_getDictID_fromDict(dict_content.as_ptr().cast(), dict_content.len()) - }; - if !args.is_raw && parsed_id == 0 { - return Err(vm.new_value_error( - "ZSTD_DICT_MAGIC_NUMBER not found, dict_content cannot be a 'raw content' \ - dictionary. To create a raw content dictionary, pass is_raw=True.", - )); + let dict_id = read_dict_id(&content); + if !args.is_raw && dict_id == 0 { + return Err(vm.new_value_error("invalid Zstandard dictionary")); } - // Raw dictionaries still get a non-zero `dict_id` whenever their - // contents happen to look like a valid dict (this is the - // documented behavior tested in `test_is_raw`). Ok(Self { - dict_content: vm.ctx.new_bytes(dict_content), - dict_id: parsed_id, + dict_content: vm.ctx.new_bytes(content), + dict_id, }) } } @@ -578,21 +662,23 @@ mod _zstd { Ok(format!( "", zelf.dict_id, - zelf.dict_content.len() + zelf.dict_content.as_bytes().len() )) } } - impl rustpython_vm::types::AsMapping for ZstdDict { - fn as_mapping() -> &'static rustpython_vm::protocol::PyMappingMethods { - static AS_MAPPING: rustpython_vm::protocol::PyMappingMethods = - rustpython_vm::protocol::PyMappingMethods { - length: Some(|mapping, _vm| { - Ok(ZstdDict::mapping_downcast(mapping).dict_content.len()) - }), - subscript: None, - ass_subscript: None, - }; + impl AsMapping for ZstdDict { + fn as_mapping() -> &'static PyMappingMethods { + static AS_MAPPING: PyMappingMethods = PyMappingMethods { + length: Some(|mapping, _vm| { + Ok(ZstdDict::mapping_downcast(mapping) + .dict_content + .as_bytes() + .len()) + }), + subscript: None, + ass_subscript: None, + }; &AS_MAPPING } } @@ -600,1468 +686,1005 @@ mod _zstd { #[pyclass(with(Constructor, Representable, AsMapping))] impl ZstdDict { #[pygetset] - fn dict_content(&self) -> PyBytesRef { - self.dict_content.clone() + fn dict_content(zelf: &Py) -> PyBytesRef { + zelf.dict_content.clone() } #[pygetset] - fn dict_id(&self) -> u32 { - self.dict_id + fn dict_id(zelf: &Py) -> u32 { + zelf.dict_id + } + + fn marker_tuple(zelf: PyRef, marker: i32, vm: &VirtualMachine) -> PyTupleRef { + vm.ctx + .new_tuple(vec![zelf.into(), vm.ctx.new_int(marker).into()]) } #[pygetset] fn as_digested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { - vm.ctx.new_tuple(vec![ - zelf.into(), - vm.ctx.new_int(DictType::Digested as i32).into(), - ]) + Self::marker_tuple(zelf, DIGESTED, vm) } #[pygetset] fn as_undigested_dict(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { - vm.ctx.new_tuple(vec![ - zelf.into(), - vm.ctx.new_int(DictType::Undigested as i32).into(), - ]) + Self::marker_tuple(zelf, UNDIGESTED, vm) } #[pygetset] fn as_prefix(zelf: PyRef, vm: &VirtualMachine) -> PyTupleRef { - vm.ctx.new_tuple(vec![ - zelf.into(), - vm.ctx.new_int(DictType::Prefix as i32).into(), - ]) - } - } - - /// The flush modes `ZstdCompressor.compress()` and - /// `ZstdCompressor.flush()` accept, mirrored as the class attributes - /// `CONTINUE`, `FLUSH_BLOCK` and `FLUSH_FRAME` by `extend_class` below. - /// - /// The discriminants are libzstd's own `ZSTD_e_continue`, `ZSTD_e_flush` - /// and `ZSTD_e_end` values, which is what CPython's `_zstd` exposes under - /// these attribute names. They are pinned here because they cross the - /// Python boundary as plain ints, both as those class attributes and as - /// the value the `last_mode` getset hands back. - #[derive(Copy, Clone, PartialEq, Eq)] - enum CompressMode { - Continue = 0, - FlushBlock = 1, - FlushFrame = 2, - } - - impl CompressMode { - /// Decode a user-supplied `mode=` argument. `None` for anything - /// outside the three documented modes. - fn from_int(mode: i32) -> Option { - Some(match mode { - m if m == Self::Continue as i32 => Self::Continue, - m if m == Self::FlushBlock as i32 => Self::FlushBlock, - m if m == Self::FlushFrame as i32 => Self::FlushFrame, - _ => return None, - }) - } - - /// The libzstd end directive this mode drives the streaming API with. - fn end_directive(self) -> ZSTD_EndDirective { - match self { - Self::Continue => ZSTD_e_continue, - Self::FlushBlock => ZSTD_e_flush, - Self::FlushFrame => ZSTD_e_end, - } + Self::marker_tuple(zelf, PREFIX, vm) } } - /// Owning wrapper around a raw libzstd compression context. Frees the - /// context with `ZSTD_freeCCtx` on drop. - struct CCtx(*mut ZSTD_CCtx); + fn zstd_dict_type_error(vm: &VirtualMachine) -> PyBaseExceptionRef { + vm.new_type_error("zstd_dict argument should be a ZstdDict object.") + } - // SAFETY: the wrapped context is only ever accessed while holding the - // owning pyclass's `PyMutex` (`CompressorState`), so no concurrent - // access can occur. - unsafe impl Send for CCtx {} + enum DictArg { + Bare(PyRef), + Marked(PyRef, i32), + } - impl CCtx { - fn create(vm: &VirtualMachine) -> PyResult { - // SAFETY: `ZSTD_createCCtx` has no preconditions; it returns NULL - // only when its allocation of the context struct fails, so that is - // the sole case to handle. - let ptr = unsafe { ZSTD_createCCtx() }; - if ptr.is_null() { - return Err(vm.new_memory_error("failed to allocate a zstd context")); - } - Ok(Self(ptr)) + fn parse_dict_arg(obj: &PyObject, vm: &VirtualMachine) -> PyResult { + if obj.downcast_ref::().is_some() { + let dict = obj + .to_owned() + .downcast::() + .map_err(|_| zstd_dict_type_error(vm))?; + return Ok(DictArg::Bare(dict)); } - - /// Forward a `(parameter, value)` pair to libzstd, passing the value - /// straight through like CPython does. The error code is returned on - /// failure so the caller can pick the right error mapping. - fn set_parameter(&mut self, param: ZSTD_cParameter, value: c_int) -> Result<(), usize> { - // SAFETY: `self.0` is a live context owned by this wrapper. - let code = unsafe { ZSTD_CCtx_setParameter(self.0, param, value) }; - if ZSTD_isError(code) != 0 { - Err(code) - } else { - Ok(()) + if let Some(tuple) = obj.downcast_ref_if_exact::(vm) { + let items = tuple.as_slice(); + if items.len() == 2 && items[0].downcast_ref::().is_some() { + let dict = items[0] + .to_owned() + .downcast::() + .map_err(|_| zstd_dict_type_error(vm))?; + let marker_obj = &items[1]; + if is_pylong(marker_obj, vm) { + let marker = match marker_obj.try_to_value::(vm) { + Ok(marker) => marker, + Err(err) if err.fast_isinstance(vm.ctx.exceptions.overflow_error) => { + return Err(err); + } + Err(_) => return Err(zstd_dict_type_error(vm)), + }; + if matches!(marker, DIGESTED | UNDIGESTED | PREFIX) { + return Ok(DictArg::Marked(dict, marker)); + } + } } } + Err(zstd_dict_type_error(vm)) + } - /// Return the context to a frame boundary, discarding a partially - /// written frame but keeping parameters and dictionary. libzstd - /// documents a session-only reset as infallible. - fn reset_session(&mut self) { - // SAFETY: `self.0` is a live context owned by this wrapper. - unsafe { ZSTD_CCtx_reset(self.0, ZSTD_ResetDirective::ZSTD_reset_session_only) }; - } + enum LoadedDict { + Dict(Box), + Prefix(Vec), + } - fn set_pledged_src_size(&mut self, pledged: Option) -> Result<(), usize> { - // libzstd represents "unknown" as `ZSTD_CONTENTSIZE_UNKNOWN`. - let pledged = pledged.map_or(ZSTD_CONTENTSIZE_UNKNOWN, |v| v); - // SAFETY: `self.0` is a live context owned by this wrapper. - let code = unsafe { ZSTD_CCtx_setPledgedSrcSize(self.0, pledged) }; - if ZSTD_isError(code) != 0 { - Err(code) - } else { - Ok(()) + fn load_dict(spec: DictArg, compress: bool, vm: &VirtualMachine) -> PyResult { + let (dict, marker) = match spec { + DictArg::Bare(dict) => { + let marker = if compress { UNDIGESTED } else { DIGESTED }; + (dict, marker) } - } + DictArg::Marked(dict, marker) => (dict, marker), + }; + let content = dict.dict_content.as_bytes(); + let loaded = match marker { + PREFIX => LoadedDict::Prefix(content.to_vec()), + DIGESTED if dict.dict_id == 0 => { + LoadedDict::Dict(Box::new(Dictionary::raw(content.to_vec()))) + } + DIGESTED => match Dictionary::from_bytes(content) { + Ok(parsed) => LoadedDict::Dict(Box::new(parsed)), + Err(_) => { + let message = if compress { + DIGESTED_COMPRESS_FAILED + } else { + DIGESTED_DECOMPRESS_FAILED + }; + return Err(zstd_error(message.to_owned(), vm)); + } + }, + _ => { + let parsed = match Dictionary::from_bytes(content) { + Ok(parsed) => parsed, + Err(_) => Dictionary::raw(content.to_vec()), + }; + LoadedDict::Dict(Box::new(parsed)) + } + }; + Ok(loaded) } - impl Drop for CCtx { - fn drop(&mut self) { - // SAFETY: `self.0` was created by `ZSTD_createCCtx` and is freed - // exactly once, here. - unsafe { ZSTD_freeCCtx(self.0) }; - } + enum AppliedDict { + None, + Dict { + dict: Box, + write_id: bool, + }, + Prefix(Vec), } - /// Owning wrapper around a raw libzstd decompression context. Frees the - /// context with `ZSTD_freeDCtx` on drop. - struct DCtx(*mut ZSTD_DCtx); + struct Session { + comp: Box, + produced: u64, + } - // SAFETY: same reasoning as `CCtx`: access is serialized by the owning - // pyclass's `PyMutex` (`DecompressorState`). - unsafe impl Send for DCtx {} + struct CompState { + settings: CompressSettings, + applied: AppliedDict, + pledge: Option, + session: Option, + pending: Vec, + block_max: usize, + last_mode_value: i32, + } - impl DCtx { - fn create(vm: &VirtualMachine) -> PyResult { - // SAFETY: `ZSTD_createDCtx` has no preconditions; it returns NULL - // only when its allocation of the context struct fails, so that is - // the sole case to handle. - let ptr = unsafe { ZSTD_createDCtx() }; - if ptr.is_null() { - return Err(vm.new_memory_error("failed to allocate a zstd context")); - } - Ok(Self(ptr)) + impl CompState { + fn fail_session(&mut self) { + self.session = None; + self.pending.clear(); + self.last_mode_value = MODE_FRAME; } - /// See [`CCtx::reset_session`]. - fn reset_session(&mut self) { - // SAFETY: `self.0` is a live context owned by this wrapper. - unsafe { ZSTD_DCtx_reset(self.0, ZSTD_ResetDirective::ZSTD_reset_session_only) }; + fn finish_frame(&mut self) { + self.session = None; + self.pending.clear(); + self.pledge = None; + if matches!(self.applied, AppliedDict::Prefix(_)) { + self.applied = AppliedDict::None; + } } - /// See [`CCtx::set_parameter`]. - fn set_parameter(&mut self, param: ZSTD_dParameter, value: c_int) -> Result<(), usize> { - // SAFETY: `self.0` is a live context owned by this wrapper. - let code = unsafe { ZSTD_DCtx_setParameter(self.0, param, value) }; - if ZSTD_isError(code) != 0 { - Err(code) + fn frame_size_header(&self, len: usize) -> Option { + if self.settings.content_size { + Some(u64::try_from(len).unwrap_or(u64::MAX)) } else { - Ok(()) + None } } - } - - impl Drop for DCtx { - fn drop(&mut self) { - // SAFETY: `self.0` was created by `ZSTD_createDCtx` and is freed - // exactly once, here. - unsafe { ZSTD_freeDCtx(self.0) }; - } - } - - /// Owning wrapper around a digested compression dictionary. Frees the - /// dictionary with `ZSTD_freeCDict` on drop. - struct CDict(*mut ZSTD_CDict); - // SAFETY: the dictionary is only referenced by the `CCtx` stored - // alongside it in `CompressorState`, which the pyclass's `PyMutex` - // serializes access to. - unsafe impl Send for CDict {} - - impl CDict { - /// Build a digested dictionary; libzstd copies the bytes into its own - /// storage. Returns `None` when libzstd rejects the content (NULL), - /// e.g. for a corrupted dictionary. - fn try_create(bytes: &[u8], level: c_int) -> Option { - // SAFETY: `bytes` is a live slice for the duration of the call, - // and `ZSTD_createCDict` copies what it needs. - let ptr = unsafe { ZSTD_createCDict(bytes.as_ptr().cast(), bytes.len(), level) }; - if ptr.is_null() { None } else { Some(Self(ptr)) } + fn stream_header(&self) -> Option { + if self.settings.content_size { + self.pledge + } else { + None + } } - } - impl Drop for CDict { - fn drop(&mut self) { - // SAFETY: `self.0` was created by `ZSTD_createCDict` and is freed - // exactly once, here. - unsafe { ZSTD_freeCDict(self.0) }; + fn apply_dict(&self, comp: &mut Compressor) -> Result<(), Error> { + match &self.applied { + AppliedDict::Dict { dict, write_id } => { + comp.set_dictionary(dict)?; + comp.set_write_dict_id(*write_id && dict.id() != 0); + } + AppliedDict::Prefix(prefix) => { + comp.set_prefix(prefix)?; + } + AppliedDict::None => {} + } + Ok(()) } - } - /// Owning wrapper around a digested decompression dictionary. Frees the - /// dictionary with `ZSTD_freeDDict` on drop. - struct DDict(*mut ZSTD_DDict); - - // SAFETY: same reasoning as `CDict` (guarded by the decompressor's - // `PyMutex` via `DecompressorState`). - unsafe impl Send for DDict {} - - impl DDict { - /// See [`CDict::try_create`]. - fn try_create(bytes: &[u8]) -> Option { - // SAFETY: `bytes` is a live slice for the duration of the call, - // and `ZSTD_createDDict` copies what it needs. - let ptr = unsafe { ZSTD_createDDict(bytes.as_ptr().cast(), bytes.len()) }; - if ptr.is_null() { None } else { Some(Self(ptr)) } + fn open_session(&mut self, header: Option) -> Result<(), Error> { + let mut comp = Compressor::with_options( + CompressOptions { + level: codec_level(self.settings.level), + checksum: self.settings.checksum, + }, + None, + )?; + if let Some(size) = header { + comp.set_pledged_src_size(size); + } + self.apply_dict(&mut comp)?; + self.session = Some(Session { + comp: Box::new(comp), + produced: 0, + }); + Ok(()) } - } - impl Drop for DDict { - fn drop(&mut self) { - // SAFETY: `self.0` was created by `ZSTD_createDDict` and is freed - // exactly once, here. - unsafe { ZSTD_freeDDict(self.0) }; + // Empty input with nothing buffered must not start a frame, except for + // an explicit end, which is a one-shot empty frame. A one-shot end + // writes the input length into the header and ignores a stored pledge. + // Short CONTINUE input stays buffered: the first stream call emits a + // header, and a short write has to leave the underlying file untouched. + fn compress(&mut self, input: &[u8], mode: i32) -> Result, Error> { + if input.is_empty() + && self.session.is_none() + && self.pending.is_empty() + && mode != MODE_FRAME + { + self.last_mode_value = mode; + return Ok(Vec::new()); + } + let oneshot = self.session.is_none() && self.pending.is_empty() && mode == MODE_FRAME; + if !input.is_empty() { + self.pending.extend_from_slice(input); + } + if self.session.is_none() + && mode == MODE_CONTINUE + && self.pending.len() < self.block_max + { + self.last_mode_value = mode; + return Ok(Vec::new()); + } + if self.session.is_none() { + let header_len = self.pending.len(); + let header = if oneshot { + self.frame_size_header(header_len) + } else { + self.stream_header() + }; + if let Err(err) = self.open_session(header) { + self.fail_session(); + return Err(err); + } + } + let flush = flush_kind(mode); + let pending = core::mem::take(&mut self.pending); + let pumped = match self.session.as_mut() { + Some(session) => pump(session, &pending, flush), + None => Err(Error::Corruption), + }; + let bytes = match pumped { + Ok(bytes) => bytes, + Err(err) => { + self.fail_session(); + return Err(err); + } + }; + if mode == MODE_FRAME { + let Some(produced) = self.session.as_ref().map(|session| session.produced) else { + self.fail_session(); + return Err(Error::Corruption); + }; + if !oneshot && self.pledge.is_some_and(|pledged| pledged != produced) { + self.fail_session(); + return Err(Error::ContentSizeMismatch); + } + self.finish_frame(); + } + self.last_mode_value = mode; + Ok(bytes) + } + } + + fn pump(session: &mut Session, input: &[u8], flush: Flush) -> Result, Error> { + let added = u64::try_from(input.len()).unwrap_or(u64::MAX); + session.produced = session.produced.saturating_add(added); + let mut out = Vec::new(); + let mut buf = vec![0u8; compress_stream_out_size()]; + let mut pending = input; + loop { + let status = session.comp.stream(pending, &mut buf, flush)?; + pending = &[]; + out.extend_from_slice(&buf[..status.output_produced]); + if status.done || status.output_produced < buf.len() { + break; + } } + Ok(out) } - /// Internal state of a `ZstdCompressor`. Holds the libzstd context, the - /// last mode used (for `last_mode` and `set_pledged_input_size` - /// validation), and the dictionary handles that the context may reference - /// internally. Field order matters here: Rust drops in declaration order, - /// so `cctx` is freed first; the held `CDict` (if any) and the source - /// `PyRef` go away afterwards, which is the safe order for - /// teardown. - struct CompressorState { - cctx: CCtx, - /// Cached digested dictionary. The CCtx references this via - /// `ZSTD_CCtx_refCDict`, so it must outlive the CCtx (handled by - /// Rust's field drop order: `cctx` drops first). - _cdict: Option, - /// Keeps the ZstdDict's bytes alive for `ZSTD_CCtx_refPrefix` mode. - _dict: Option>, - last_mode: CompressMode, - } - - impl CompressorState { - /// Recover from a failed `compress()`/`flush()`, as CPython does: - /// drop the half-written frame and report `last_mode == FLUSH_FRAME`. - /// Without this the context stays mid-frame and every later call - /// fails with "Operation not authorized at current processing stage". - fn reset_session(&mut self) { - self.last_mode = CompressMode::FlushFrame; - self.cctx.reset_session(); - } + #[derive(FromArgs)] + struct ZstdCompressorArgs { + #[pyarg(any, optional, py_default = "None")] + level: OptionalOption, + #[pyarg(any, optional, py_default = "None")] + options: OptionalOption, + #[pyarg(any, optional, py_default = "None")] + zstd_dict: OptionalOption, } #[pyattr] - #[pyclass(name = "ZstdCompressor")] + #[pyclass(module = "compression.zstd", name = "ZstdCompressor", traverse)] #[derive(PyPayload)] struct ZstdCompressor { - state: PyMutex, + dict_ref: Option>, + #[pytraverse(skip)] + last_mode: AtomicI32, + #[pytraverse(skip)] + state: PyMutex, } - impl core::fmt::Debug for ZstdCompressor { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - write!(f, "_zstd.ZstdCompressor") + impl fmt::Debug for ZstdCompressor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("compression.zstd.ZstdCompressor") } } - #[derive(FromArgs)] - pub(super) struct ZstdCompressorArgs { - #[pyarg(any, optional)] - level: OptionalOption, - #[pyarg(any, optional)] - options: OptionalOption, - #[pyarg(any, optional)] - zstd_dict: OptionalOption, - } - - /// Decode the public `mode` int accepted by `ZstdCompressor.compress()`, - /// which takes all three modes. - fn compress_mode_from_int(mode: i32, vm: &VirtualMachine) -> PyResult { - CompressMode::from_int(mode).ok_or_else(|| { - let (cont, block, frame) = ( - CompressMode::Continue as i32, - CompressMode::FlushBlock as i32, - CompressMode::FlushFrame as i32, - ); - vm.new_value_error(format!( - "mode argument wrong value, it should be one of \ - ZstdCompressor.CONTINUE ({cont}), \ - ZstdCompressor.FLUSH_BLOCK ({block}), or \ - ZstdCompressor.FLUSH_FRAME ({frame})" - )) - }) - } - impl Constructor for ZstdCompressor { type Args = ZstdCompressorArgs; fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { - let level_opt = args.level.flatten(); - let options_opt = args.options.flatten(); - let dict_opt = args.zstd_dict.flatten(); - - if level_opt.is_some() && options_opt.is_some() { + let level = args.level.flatten(); + let options = args.options.flatten(); + let zstd_dict = args.zstd_dict.flatten(); + if level.is_some() && options.is_some() { return Err(vm.new_type_error("Only one of level or options should be used.")); } - - let mut cctx = CCtx::create(vm)?; - - // Mirror CPython's `self->compression_level`: the level libzstd - // will actually compress at, which a digested dictionary has to - // be built for. - let mut level = ZSTD_CLEVEL_DEFAULT; - - if let Some(level_obj) = level_opt { - level = parse_compression_level(&level_obj, vm)?; - cctx.set_parameter(ZSTD_cParameter::ZSTD_c_compressionLevel, level) - .map_err(|_| param_value_error_for(ZSTD_c_compressionLevel, level, true, vm))?; + let mut settings = CompressSettings::default(); + if let Some(level) = level.as_ref() { + settings.level = parse_level(level, vm)?; } - - if let Some(options_obj) = options_opt - && let Some(from_options) = apply_options(&mut cctx, options_obj, true, vm)? - { - level = from_options; + if let Some(options) = options.as_ref() { + settings = parse_compress_options(options, vm)?; } - - let state = - build_compressor_state(cctx, dict_opt, level, CompressMode::FlushFrame, vm)?; + let (dict_ref, applied) = match zstd_dict.as_ref() { + Some(obj) => { + let spec = parse_dict_arg(obj, vm)?; + let dict_ref = match &spec { + DictArg::Bare(dict) | DictArg::Marked(dict, _) => dict.clone(), + }; + let loaded = load_dict(spec, true, vm)?; + let applied = match loaded { + LoadedDict::Dict(dict) => AppliedDict::Dict { + dict, + write_id: settings.dict_id, + }, + LoadedDict::Prefix(bytes) => AppliedDict::Prefix(bytes), + }; + (Some(dict_ref), applied) + } + None => (None, AppliedDict::None), + }; Ok(Self { - state: PyMutex::new(state), + dict_ref, + last_mode: AtomicI32::new(MODE_FRAME), + state: PyMutex::new(CompState { + block_max: block_capacity(settings.level), + settings, + applied, + pledge: None, + session: None, + pending: Vec::new(), + last_mode_value: MODE_FRAME, + }), }) } } - /// Parse and validate a compression `level` argument. libzstd silently - /// clamps out-of-range levels rather than erroring, but CPython surfaces - /// them as `ValueError`, and bigints become `ValueError` (not - /// `OverflowError`) for the same reason. This helper centralizes both - /// conversions so the constructor stays linear. - fn parse_compression_level(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { - let (lo, hi) = level_bounds(); - let level: i32 = obj.try_to_value(vm).map_err(|e| { - if e.fast_isinstance(vm.ctx.exceptions.overflow_error) { - vm.new_value_error(format!( - "illegal compression level; the valid range is [{lo}, {hi}]" - )) - } else { - e - } - })?; - if level < lo || level > hi { - return Err(vm.new_value_error(format!( - "illegal compression level {level}; the valid range is [{lo}, {hi}]" - ))); - } - Ok(level) + #[derive(FromArgs)] + struct CompressCallArgs { + #[pyarg(any)] + data: ArgBytesLike, + #[pyarg(any, default = 0)] + mode: i32, } - /// Drain an `options=` dict onto either a `CCtx` or a `DCtx`. Validates - /// each key/value pair (rejects wrong enum kind, rejects floats, rejects - /// out-of-range values) so the constructor's flow stays a single line. - fn apply_options( - ctx: &mut dyn ParamSetter, - options_obj: PyObjectRef, - is_compress: bool, - vm: &VirtualMachine, - ) -> PyResult> { - let mut level = None; - let dict = options_obj - .downcast::() - .map_err(|_| vm.new_type_error("options must be a dict"))?; - // Resolve the parameter enum class that is invalid for this direction - // (registered by the pure-Python wrapper via set_parameter_types) once, - // so a key from the wrong family yields a clear TypeError naming the - // type. `None` when the wrapper never ran (e.g. `_zstd` used directly). - let (wrong_kind_attr, kind) = if is_compress { - ("_decompression_parameter_type", "compression") - } else { - ("_compression_parameter_type", "decompression") - }; - let wrong_kind = vm.get_attribute_opt(vm.import("_zstd", 0)?, wrong_kind_attr)?; - for (k, v) in dict { - // Reject a key from the wrong parameter family before any numeric - // coercion, so the error names the type rather than giving a - // generic out-of-range message. - check_wrong_param_kind(&k, wrong_kind.as_ref(), kind, vm)?; - let key_int: i32 = k.try_to_value(vm)?; - let val_int: i32 = v.try_to_value(vm)?; - // Only the compression level is range-checked here. libzstd - // silently clamps it (`ZSTD_cParam_clampBounds`) instead of - // reporting an error, so without this check the `ValueError` - // that `test_compress_parameters` requires would never be - // raised. - // - // Every other parameter is left to libzstd, which does its own - // bounds checking and whose error code becomes the same - // `ValueError` via `ParamSetter::apply`. A blanket range test - // here would reject values CPython accepts: 0 means "use the - // default" for most parameters, the boolean flags take any - // non-zero value, and `nb_workers`/`job_size`/`overlap_log` are - // clamped rather than rejected. - if is_compress - && key_int == ZSTD_c_compressionLevel - && let Some((lo, hi)) = lookup_param_bounds(key_int, is_compress) - && (val_int < lo || val_int > hi) - { - return Err(param_value_error_for(key_int, val_int, is_compress, vm)); - } - ctx.apply(key_int, val_int, vm)?; - // Track the level the compressor ends up with: a digested - // dictionary has to be built at that level (see `load_dict`). - if is_compress && key_int == ZSTD_c_compressionLevel { - level = Some(val_int); - } - } - Ok(level) + #[derive(FromArgs)] + struct FlushCallArgs { + #[pyarg(any, default = 2)] + mode: i32, } - /// Trait wrapper over `CCtx::set_parameter` and `DCtx::set_parameter` so - /// `apply_options` can drive either context without duplicated code. - /// `apply` validates the (id, value) pair via `cparameter_from_int` / - /// `dparameter_from_int`, then forwards to libzstd. - trait ParamSetter { - fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()>; - } + #[pyclass(with(Constructor))] + impl ZstdCompressor { + fn locked_compress(zelf: &Py, input: &[u8], mode: i32) -> Result, Error> { + let mut state = zelf.state.lock(); + let result = state.compress(input, mode); + zelf.last_mode + .store(state.last_mode_value, Ordering::Relaxed); + result + } - impl ParamSetter for CCtx { - fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()> { - let p = cparameter_from_int(param, value, vm)?; - self.set_parameter(p, value) - .map_err(|_| param_value_error_for(param, value, true, vm))?; - Ok(()) + #[pymethod] + fn compress( + zelf: &Py, + args: CompressCallArgs, + vm: &VirtualMachine, + ) -> PyResult> { + let CompressCallArgs { data, mode } = args; + check_compress_mode(mode, vm)?; + let input = data.with_ref(<[u8]>::to_vec); + vm.allow_threads(|| Self::locked_compress(zelf, &input, mode)) + .map_err(|err| compress_error(err, vm)) } - } - impl ParamSetter for DCtx { - fn apply(&mut self, param: i32, value: i32, vm: &VirtualMachine) -> PyResult<()> { - let p = dparameter_from_int(param, vm)?; - self.set_parameter(p, value) - .map_err(|_| param_value_error_for(param, value, false, vm))?; - Ok(()) + #[pymethod] + fn flush(zelf: &Py, args: FlushCallArgs, vm: &VirtualMachine) -> PyResult> { + let FlushCallArgs { mode } = args; + check_flush_mode(mode, vm)?; + vm.allow_threads(|| Self::locked_compress(zelf, &[], mode)) + .map_err(|err| compress_error(err, vm)) } - } - /// Trait that captures the only differences between how the compressor - /// and decompressor consume a dictionary: the name of the type that - /// appears in error messages, the eager-validation constructor for the - /// digested variant, and the three ways of attaching it to the context. - /// The `Err` payload of each method is a libzstd error code. - trait DictLoader { - type Digested; - const KIND_NAME: &'static str; - /// Load mode applied when the caller passes a bare `ZstdDict` instead - /// of one of the `as_*` tuples. CPython compresses with an undigested - /// dictionary by default and decompresses with a digested one. - const DEFAULT_DICT_TYPE: DictType; - /// `level` is the compression level the digested dictionary must be - /// built for; the decompressor has no level and ignores it. - fn try_create_digested(bytes: &[u8], level: c_int) -> Option; - fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize>; - fn load_undigested(&mut self, bytes: &[u8]) -> Result<(), usize>; - fn ref_prefix(&mut self, bytes: &[u8]) -> Result<(), usize>; - } - - impl DictLoader for CCtx { - type Digested = CDict; - const KIND_NAME: &'static str = "ZSTD_CDict"; - const DEFAULT_DICT_TYPE: DictType = DictType::Undigested; - fn try_create_digested(bytes: &[u8], level: c_int) -> Option { - // The level matters: `ZSTD_CCtx_refCDict` takes its compression - // parameters from the CDict, so a CDict built at the default - // level would override whatever the caller asked for. - CDict::try_create(bytes, level) - } - fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize> { - // SAFETY: both handles are live; `load_dict`'s safety contract - // keeps `dict` alive at least as long as `self`. - let code = unsafe { ZSTD_CCtx_refCDict(self.0, dict.0.cast_const()) }; - if ZSTD_isError(code) != 0 { - Err(code) + #[pymethod] + fn set_pledged_input_size( + zelf: &Py, + size: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult<()> { + let pledge = parse_pledge(&size, vm)?; + let bad_mode = vm.allow_threads(|| { + let mut state = zelf.state.lock(); + if zelf.last_mode.load(Ordering::Relaxed) != MODE_FRAME { + true + } else { + state.pledge = pledge; + false + } + }); + if bad_mode { + Err(vm.new_value_error( + "set_pledged_input_size() method must be called when last_mode == FLUSH_FRAME", + )) } else { Ok(()) } } - fn load_undigested(&mut self, bytes: &[u8]) -> Result<(), usize> { - // SAFETY: `bytes` is a live slice for the duration of the call; - // its contents are copied into the context. - let code = - unsafe { ZSTD_CCtx_loadDictionary(self.0, bytes.as_ptr().cast(), bytes.len()) }; - if ZSTD_isError(code) != 0 { - Err(code) - } else { - Ok(()) - } + + #[pygetset] + fn last_mode(zelf: &Py) -> i32 { + zelf.last_mode.load(Ordering::Relaxed) } - fn ref_prefix(&mut self, bytes: &[u8]) -> Result<(), usize> { - // SAFETY: `bytes` is a live slice for the duration of the call. - // libzstd keeps a raw pointer to the bytes afterwards; - // `load_dict`'s safety contract keeps them alive at least as - // long as `self`. - let code = unsafe { ZSTD_CCtx_refPrefix(self.0, bytes.as_ptr().cast(), bytes.len()) }; - if ZSTD_isError(code) != 0 { - Err(code) - } else { - Ok(()) - } + + #[extend_class] + fn extend_class(ctx: &Context, class: &'static Py) { + class.set_attr( + ctx.intern_str("CONTINUE"), + ctx.new_int(MODE_CONTINUE).into(), + ); + class.set_attr( + ctx.intern_str("FLUSH_BLOCK"), + ctx.new_int(MODE_BLOCK).into(), + ); + class.set_attr( + ctx.intern_str("FLUSH_FRAME"), + ctx.new_int(MODE_FRAME).into(), + ); } } - impl DictLoader for DCtx { - type Digested = DDict; - const KIND_NAME: &'static str = "ZSTD_DDict"; - const DEFAULT_DICT_TYPE: DictType = DictType::Digested; - fn try_create_digested(bytes: &[u8], _level: c_int) -> Option { - DDict::try_create(bytes) + fn parse_pledge(size: &PyObject, vm: &VirtualMachine) -> PyResult> { + if vm.is_none(size) { + return Ok(None); } - fn ref_digested(&mut self, dict: &Self::Digested) -> Result<(), usize> { - // SAFETY: both handles are live; `load_dict`'s safety contract - // keeps `dict` alive at least as long as `self`. - let code = unsafe { ZSTD_DCtx_refDDict(self.0, dict.0.cast_const()) }; - if ZSTD_isError(code) != 0 { - Err(code) - } else { - Ok(()) - } + if !is_pylong(size, vm) { + return Err(pledge_error(vm)); } - fn load_undigested(&mut self, bytes: &[u8]) -> Result<(), usize> { - // SAFETY: `bytes` is a live slice for the duration of the call; - // its contents are copied into the context. - let code = - unsafe { ZSTD_DCtx_loadDictionary(self.0, bytes.as_ptr().cast(), bytes.len()) }; - if ZSTD_isError(code) != 0 { - Err(code) - } else { - Ok(()) - } - } - fn ref_prefix(&mut self, bytes: &[u8]) -> Result<(), usize> { - // SAFETY: `bytes` is a live slice for the duration of the call. - // libzstd keeps a raw pointer to the bytes afterwards; - // `load_dict`'s safety contract keeps them alive at least as - // long as `self`. - let code = unsafe { ZSTD_DCtx_refPrefix(self.0, bytes.as_ptr().cast(), bytes.len()) }; - if ZSTD_isError(code) != 0 { - Err(code) - } else { - Ok(()) - } + let int_obj = size.try_index(vm)?; + match int_obj.as_bigint().to_u64() { + Some(value) if value < u64::MAX - 1 => Ok(Some(value)), + _ => Err(pledge_error(vm)), } } - /// Return value of `load_dict`: the digested `CDict`/`DDict` (if any) - /// and the `PyRef` we hold to keep the dictionary bytes alive - /// while `ref_prefix` may point into them. - type DictLoadResult = PyResult<(Option, Option>)>; - - /// Common path for attaching a dictionary to either context type. Returns - /// the digested `CDict`/`DDict` (if the caller used digested mode) plus - /// the `PyRef` whose bytes libzstd's `ref_prefix` may point into. - /// - /// # Safety - /// - /// libzstd stores the dictionary as a raw pointer that bypasses Rust's - /// lifetime tracking. The caller must keep both returned values alive at - /// least as long as `ctx`: - /// - /// - In `digested` mode, `ctx` holds a raw pointer to the returned - /// `L::Digested`; dropping it before `ctx` is use-after-free. - /// - In `prefix` mode, `ctx` holds a raw pointer into the bytes owned by - /// the returned `PyRef`; dropping the `PyRef` before `ctx` - /// is use-after-free. - /// - /// In `undigested` mode the bytes are copied into `ctx`, so neither - /// return value carries a safety obligation — but the caller cannot tell - /// the modes apart, so it must keep both alive regardless. - unsafe fn load_dict( - ctx: &mut L, - dict_obj: Option, - level: c_int, - vm: &VirtualMachine, - ) -> DictLoadResult { - let Some(dict_obj) = dict_obj else { - return Ok((None, None)); - }; - let (zdict, dict_type) = parse_zstd_dict_arg(dict_obj, L::DEFAULT_DICT_TYPE, vm)?; - let bad_dict_err = || -> PyBaseExceptionRef { - new_zstd_error( - format!( - "Failed to load the {} instance from corrupted Zstandard dictionary content.", - L::KIND_NAME - ), - vm, - ) - }; - let dict_bytes = zdict.dict_content.as_bytes(); - let mut digested = None; - match dict_type { - DictType::Prefix => { - ctx.ref_prefix(dict_bytes).map_err(|_| bad_dict_err())?; - } - DictType::Digested => { - // Build the digested dict eagerly so a corrupted dictionary - // surfaces as a `ZstdError` at construction time, not when - // the first compress/decompress call runs. - let d = L::try_create_digested(dict_bytes, level).ok_or_else(bad_dict_err)?; - ctx.ref_digested(&d).map_err(|_| bad_dict_err())?; - digested = Some(d); - } - DictType::Undigested => { - // Copy the bytes into the context. Whether a corrupted - // dictionary is caught here is up to libzstd and differs by - // direction: `ZSTD_DCtx_loadDictionary` digests the content - // eagerly and fails now, while `ZSTD_CCtx_loadDictionary` - // defers that work and accepts it. - ctx.load_undigested(dict_bytes) - .map_err(|_| bad_dict_err())?; - } - } - Ok((digested, Some(zdict))) + fn pledge_error(vm: &VirtualMachine) -> PyBaseExceptionRef { + let limit = u64::MAX - 1; + vm.new_value_error(format!( + "size argument should be a positive int less than {limit}" + )) } - /// Build a fully-initialized `CompressorState` from a freshly-created - /// `CCtx` and an optional dictionary argument. This is the safe interface - /// that `unsafe fn load_dict` was waiting for: by assembling the struct - /// here, both invariants `load_dict` documents become structural and a - /// safe-Rust caller cannot split the pieces apart. - fn build_compressor_state( - mut cctx: CCtx, - dict_obj: Option, - level: c_int, - last_mode: CompressMode, - vm: &VirtualMachine, - ) -> PyResult { - // SAFETY: `load_dict` requires its two return values to outlive `ctx`. - // We satisfy that by moving `cctx` and both return values into - // `CompressorState` in one expression — Rust drops the struct's - // fields in declaration order, so on teardown `cctx` is dropped - // first, releasing its raw pointers before `_cdict` (digested mode) - // and `_dict` (prefix mode) are freed. `CompressorState` is private - // to this module and is never destructured, so no safe caller can - // reorder the drops. - let (cdict, dict) = unsafe { load_dict::(&mut cctx, dict_obj, level, vm) }?; - Ok(CompressorState { - cctx, - _cdict: cdict, - _dict: dict, - last_mode, - }) + enum DecApplied { + None, + Dict(Box), + Prefix(Vec), } - /// Build a fully-initialized `DecompressorState`. See - /// [`build_compressor_state`] for the safety reasoning; - /// `DecompressorState`'s field order plays the same role here. - fn build_decompressor_state( - mut dctx: DCtx, - dict_obj: Option, - vm: &VirtualMachine, - ) -> PyResult { - // SAFETY: see [`build_compressor_state`]. - // The decompressor has no compression level; `DCtx` ignores it. - let (ddict, dict) = - unsafe { load_dict::(&mut dctx, dict_obj, ZSTD_CLEVEL_DEFAULT, vm) }?; - Ok(DecompressorState { - dctx, - _ddict: ddict, - _dict: dict, - eof: false, - needs_input: true, - unused_data: vm.ctx.empty_bytes.clone(), - input_buffer: Vec::new(), - }) + struct DecState { + window_max: u64, + applied: DecApplied, + src: Vec, + fed: usize, + dec: Option>, + frame_len: Option, + out: Vec, + out_pos: usize, + finished: bool, + probed: usize, + needs_input: bool, } - /// Drive `ZSTD_compressStream2` until the input is fully consumed and, - /// for flush/end directives, the internal buffers report zero remaining - /// bytes. Grows the output `Vec` by `ZSTD_CStreamOutSize` chunks as - /// needed. - fn do_compress( - state: &mut CompressorState, - data: &[u8], - end_op: ZSTD_EndDirective, - vm: &VirtualMachine, - ) -> PyResult> { - // Release the GIL for the duration of the compression loop. Safety: - // `data` is an immutable borrow of a local `Vec` in the caller, - // `state` is held under the compressor's `PyMutex` (no other Python - // thread can touch it), the `_dict` bytes referenced by libzstd are - // an immutable `Vec` inside a `PyRef` (other readers fine), - // and the output `Vec` is local to this function. No Python object - // access happens inside the closure — error codes are surfaced as - // `usize` and converted into exceptions after re-attaching. - let is_end = end_op != ZSTD_e_continue; - let chunk_size = ZSTD_CStreamOutSize().max(1); - let result: Result, usize> = vm.allow_threads(|| { - let mut output: Vec = Vec::new(); - let mut input = ZSTD_inBuffer { - src: data.as_ptr().cast(), - size: data.len(), - pos: 0, + impl DecState { + fn reset_session(&mut self) { + self.src.clear(); + self.fed = 0; + self.dec = None; + self.frame_len = None; + self.out.clear(); + self.out_pos = 0; + self.finished = false; + self.probed = 0; + self.needs_input = true; + } + + fn is_eof(&self) -> bool { + self.finished && self.out.is_empty() + } + + fn unused(&self) -> &[u8] { + // Bytes after this frame stay hidden until its output is drained. + if !self.is_eof() { + return &[]; + } + let start = self.frame_len.unwrap_or(self.src.len()).min(self.src.len()); + &self.src[start..] + } + + fn clear_prefix(&mut self) { + if matches!(self.applied, DecApplied::Prefix(_)) { + self.applied = DecApplied::None; + } + } + + fn make_decoder(&self) -> Box { + let mut dec = Decompressor::with_options(DecompressOptions { + window_max: self.window_max, + force_ignore_checksum: false, + }); + match &self.applied { + DecApplied::Dict(dict) => dec.set_dictionary(dict.as_ref().clone()), + DecApplied::Prefix(prefix) => dec.set_prefix(prefix), + DecApplied::None => {} + } + Box::new(dec) + } + + fn classify(&mut self) -> Result<(), Error> { + if self.dec.is_some() || self.finished || self.src.len() < 4 { + return Ok(()); + } + let Some(magic) = read_u32_at(&self.src, 0) else { + return Err(Error::Corruption); }; - loop { - let prev_len = output.len(); - output.reserve(chunk_size); - let remaining = { - let spare = output.spare_capacity_mut(); - let mut out_buf = ZSTD_outBuffer { - dst: spare.as_mut_ptr().cast(), - size: spare.len(), - pos: 0, - }; - // SAFETY: `state.cctx` is a live context; `out_buf` - // covers the Vec's spare capacity, which is valid for - // `spare.len()` bytes; `input` points to `data`, which - // outlives the call. - let code = unsafe { - ZSTD_compressStream2(state.cctx.0, &mut out_buf, &mut input, end_op) - }; - if ZSTD_isError(code) != 0 { - return Err(code); - } - // SAFETY: libzstd wrote `out_buf.pos` bytes - // (`out_buf.pos <= spare.len()`) into the spare capacity. - unsafe { output.set_len(prev_len + out_buf.pos) }; - code + if is_skippable(magic) { + if self.src.len() < 8 { + return Ok(()); + } + let Some(body_len) = read_u32_at(&self.src, 4) else { + return Err(Error::Corruption); }; - let consumed_all = input.pos == input.size; - // Stop when input is fully consumed and, for flush/end - // directives, libzstd reports that all internal buffers have - // been drained (remaining == 0). Otherwise loop; the next - // `reserve` will grow the output if we hit the previous cap. - if consumed_all && (!is_end || remaining == 0) { - break Ok(output); + let Ok(body) = usize::try_from(body_len) else { + return Err(Error::Corruption); + }; + let Some(total) = 8usize.checked_add(body) else { + return Err(Error::Corruption); + }; + if self.src.len() < total { + return Ok(()); } + self.frame_len = Some(total); + self.finished = true; + self.clear_prefix(); + return Ok(()); } - }); - result.map_err(|c| catch_zstd_error(c, vm)) - } + if magic != MAGIC { + return Err(Error::BadMagic); + } + let dec = self.make_decoder(); + self.dec = Some(dec); + Ok(()) + } - #[pyclass(with(Constructor))] - impl ZstdCompressor { - #[pymethod] - fn compress(&self, args: CompressMethodArgs, vm: &VirtualMachine) -> PyResult> { - let mode = match args.mode { - Some(mode) => compress_mode_from_int(mode, vm)?, - None => CompressMode::Continue, - }; - let end_op = mode.end_directive(); - let data = args.data.with_ref(|b| b.to_vec()); - let mut state = self.state.lock(); - match do_compress(&mut state, &data, end_op, vm) { - Ok(out) => { - state.last_mode = mode; - Ok(out) + fn probe(&mut self) -> Result<(), Error> { + if self.frame_len.is_some() || self.probed == self.src.len() { + return Ok(()); + } + match find_frame_compressed_size(&self.src) { + Ok(0) => Err(Error::Corruption), + Ok(len) => { + self.frame_len = Some(len); + self.probed = self.src.len(); + Ok(()) } - Err(e) => { - state.reset_session(); - Err(e) + Err(Error::UnexpectedEof) => { + self.probed = self.src.len(); + Ok(()) } + Err(err) => Err(err), } } - #[pymethod] - fn flush(&self, args: FlushMethodArgs, vm: &VirtualMachine) -> PyResult> { - // `flush()` takes only the two flushing modes, and rejects - // everything else — `CONTINUE` included — with its own message - // rather than the one `compress()` uses. - let mode = match args.mode.map(CompressMode::from_int) { - None => CompressMode::FlushFrame, - Some(Some(mode @ (CompressMode::FlushBlock | CompressMode::FlushFrame))) => mode, - Some(_) => { - let (block, frame) = ( - CompressMode::FlushBlock as i32, - CompressMode::FlushFrame as i32, - ); - return Err(vm.new_value_error(format!( - "mode argument wrong value, it should be \ - ZstdCompressor.FLUSH_FRAME ({frame}) or \ - ZstdCompressor.FLUSH_BLOCK ({block})" - ))); - } + fn drain(&mut self, chunk: Vec, end: bool) -> Result<(), Error> { + let mut dec = match self.dec.take() { + Some(dec) => dec, + None => return Err(Error::Corruption), }; - let end_op = mode.end_directive(); - let mut state = self.state.lock(); - match do_compress(&mut state, &[], end_op, vm) { - Ok(out) => { - state.last_mode = mode; - Ok(out) - } - Err(e) => { - state.reset_session(); - Err(e) + let mut buf = vec![0u8; decompress_stream_out_size()]; + let mut input: &[u8] = &chunk; + let result = loop { + match dec.stream(input, &mut buf, end) { + Ok(status) => { + input = &[]; + self.out.extend_from_slice(&buf[..status.output_produced]); + if status.done { + self.finished = true; + break Ok(()); + } + if status.output_produced < buf.len() { + break Ok(()); + } + } + Err(err) => break Err(err), } + }; + self.dec = Some(dec); + if self.finished { + self.clear_prefix(); } + result } - #[pymethod] - fn set_pledged_input_size(&self, size: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - // Parse the argument *before* taking the lock: `try_index` can run a - // Python `__index__`, and doing that while holding `self.state` would - // let a re-entrant call into this compressor deadlock. CPython - // likewise converts the argument before touching the compressor. - // - // Python passes `None` to mean "unknown"; libzstd represents that - // internally as `ZSTD_CONTENTSIZE_UNKNOWN` (`u64::MAX`), which - // `set_pledged_src_size` maps `None` to. libzstd also reserves - // `ZSTD_CONTENTSIZE_ERROR` (`u64::MAX - 1`), so a concrete size must - // be strictly less than that; reject anything else up front so - // callers see the documented `ValueError`, not a libzstd-level error. - let pledged: Option = if vm.is_none(&size) { - None - } else { - const LIMIT: u64 = u64::MAX - 1; - let err = || { - vm.new_value_error(format!( - "size argument should be a positive int less than {LIMIT}" - )) - }; - // `try_to_primitive` fails (OverflowError) for negatives and for - // values above `u64::MAX`; the explicit check covers the rest of - // the reserved range. - let v: u64 = size - .try_index(vm)? - .try_to_primitive(vm) - .map_err(|_| err())?; - if v >= LIMIT { - return Err(err()); - } - Some(v) + fn feed(&mut self) -> Result<(), Error> { + if self.dec.is_none() || self.finished { + return Ok(()); + } + let (input_end, end_flag) = match self.frame_len { + Some(len) => (len.min(self.src.len()), self.src.len() >= len), + None => (self.src.len(), false), }; - let mut state = self.state.lock(); - if state.last_mode != CompressMode::FlushFrame { - return Err(vm.new_value_error( - "set_pledged_input_size() method must be called when last_mode == FLUSH_FRAME", - )); + if self.fed < input_end { + let chunk = self.src[self.fed..input_end].to_vec(); + self.fed = input_end; + self.drain(chunk, end_flag) + } else if end_flag { + self.drain(Vec::new(), true) + } else { + Ok(()) } - state - .cctx - .set_pledged_src_size(pledged) - .map_err(|c| catch_zstd_error(c, vm))?; - Ok(()) } - #[pygetset] - fn last_mode(&self) -> i32 { - self.state.lock().last_mode as i32 + fn drive(&mut self) -> Result<(), Error> { + if self.finished { + return Ok(()); + } + self.classify()?; + if self.finished || self.dec.is_none() { + return Ok(()); + } + self.probe()?; + self.feed() + } + + // needs_input stays true unless this call filled max_length, or the + // frame is already finished. Returning b'' with needs_input false and + // eof false makes the reader spin. + fn take(&mut self, max_length: Option) -> Vec { + let available = self.out.len().saturating_sub(self.out_pos); + let take = match max_length { + Some(max) => available.min(max), + None => available, + }; + let start = self.out_pos; + let end = start + take; + let bytes = self.out[start..end].to_vec(); + self.out_pos = end; + if self.out_pos == self.out.len() { + self.out.clear(); + self.out_pos = 0; + } + self.needs_input = if self.finished { + false + } else { + match max_length { + Some(max) => take != max, + None => true, + } + }; + bytes } - // Install class-level constants `CONTINUE`, `FLUSH_BLOCK`, and - // `FLUSH_FRAME` so callers can reference them as - // `ZstdCompressor.FLUSH_FRAME` (as the Python `ZstdFile` wrapper - // does). - #[extend_class] - fn extend_class(ctx: &Context, class: &'static Py) { - class.set_attr( - ctx.intern_str("CONTINUE"), - ctx.new_int(CompressMode::Continue as i32).into(), - ); - class.set_attr( - ctx.intern_str("FLUSH_BLOCK"), - ctx.new_int(CompressMode::FlushBlock as i32).into(), - ); - class.set_attr( - ctx.intern_str("FLUSH_FRAME"), - ctx.new_int(CompressMode::FlushFrame as i32).into(), - ); + fn pull(&mut self, input: &[u8], max_length: Option) -> Result, Error> { + self.src.extend_from_slice(input); + if !self.finished { + self.drive()?; + } + Ok(self.take(max_length)) } } - #[derive(FromArgs)] - pub(super) struct CompressMethodArgs { - /// Positional-or-keyword, as in CPython's clinic signature for - /// `_zstd.ZstdCompressor.compress` (no `/` marker). - #[pyarg(any)] - data: ArgBytesLike, - #[pyarg(any, optional)] - mode: Option, + enum DecFail { + Eof, + Codec(Error), } #[derive(FromArgs)] - pub(super) struct FlushMethodArgs { - #[pyarg(any, optional)] - mode: Option, - } - - /// Internal state of a `ZstdDecompressor`. The CPython decompressor is - /// single-frame: once we hit end-of-frame, additional bytes go into - /// `unused_data` and further `decompress` calls raise `EOFError`. Field - /// drop order matters here for the same reason as in `CompressorState`: - /// the `dctx` is freed first and must give up its internal pointers - /// before any referenced `DDict`/`PyRef` is dropped. - struct DecompressorState { - dctx: DCtx, - /// Cached decompression dictionary referenced by the DCtx via - /// `ZSTD_DCtx_refDDict`. - _ddict: Option, - _dict: Option>, - eof: bool, - needs_input: bool, - /// Bytes that arrived after the end of the first frame. - unused_data: PyBytesRef, - /// Input bytes buffered because the previous `decompress` call ran - /// into its `max_length` cap before consuming them all. - input_buffer: Vec, - } - - impl DecompressorState { - /// Recover from a failed `decompress()`, mirroring CPython's - /// `decompressor_reset_session_lock_held`. Without it the DCtx stays - /// wedged on the corrupt frame and every later call re-reports the - /// same error, even for valid input. - fn reset_session(&mut self, vm: &VirtualMachine) { - self.input_buffer.clear(); - self.unused_data = vm.ctx.empty_bytes.clone(); - self.needs_input = true; - self.eof = false; - self.dctx.reset_session(); - } + struct ZstdDecompressorArgs { + #[pyarg(any, optional, py_default = "None")] + zstd_dict: OptionalOption, + #[pyarg(any, optional, py_default = "None")] + options: OptionalOption, } #[pyattr] - #[pyclass(name = "ZstdDecompressor")] + #[pyclass(module = "compression.zstd", name = "ZstdDecompressor", traverse)] #[derive(PyPayload)] struct ZstdDecompressor { - state: PyMutex, + dict_ref: Option>, + #[pymember] + #[pytraverse(skip)] + eof: AtomicBool, + #[pymember] + #[pytraverse(skip)] + needs_input: AtomicBool, + #[pytraverse(skip)] + state: PyMutex, } - impl core::fmt::Debug for ZstdDecompressor { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - write!(f, "_zstd.ZstdDecompressor") + impl fmt::Debug for ZstdDecompressor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("compression.zstd.ZstdDecompressor") } } - #[derive(FromArgs)] - pub(super) struct ZstdDecompressorArgs { - #[pyarg(any, optional)] - zstd_dict: OptionalOption, - #[pyarg(any, optional)] - options: OptionalOption, - } - impl Constructor for ZstdDecompressor { type Args = ZstdDecompressorArgs; fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { - let dict_opt = args.zstd_dict.flatten(); - let options_opt = args.options.flatten(); - - let mut dctx = DCtx::create(vm)?; - - if let Some(options_obj) = options_opt { - apply_options(&mut dctx, options_obj, false, vm)?; - } - - let state = build_decompressor_state(dctx, dict_opt, vm)?; + let zstd_dict = args.zstd_dict.flatten(); + let options = args.options.flatten(); + let (dict_ref, applied) = match zstd_dict.as_ref() { + Some(obj) => { + let spec = parse_dict_arg(obj, vm)?; + let dict_ref = match &spec { + DictArg::Bare(dict) | DictArg::Marked(dict, _) => dict.clone(), + }; + let loaded = load_dict(spec, false, vm)?; + let applied = match loaded { + LoadedDict::Dict(dict) => DecApplied::Dict(dict), + LoadedDict::Prefix(bytes) => DecApplied::Prefix(bytes), + }; + (Some(dict_ref), applied) + } + None => (None, DecApplied::None), + }; + let window_log_max = match options.as_ref() { + Some(options) => parse_decompress_options(options, vm)?.window_log_max, + None => DEFAULT_WINDOW_LOG_MAX, + }; Ok(Self { - state: PyMutex::new(state), + dict_ref, + eof: AtomicBool::new(false), + needs_input: AtomicBool::new(true), + state: PyMutex::new(DecState { + window_max: window_max_bytes(window_log_max), + applied, + src: Vec::new(), + fed: 0, + dec: None, + frame_len: None, + out: Vec::new(), + out_pos: 0, + finished: false, + probed: 0, + needs_input: true, + }), }) } } - #[derive(FromArgs)] - pub(super) struct DecompressMethodArgs { - /// Positional-or-keyword, as in CPython's clinic signature for - /// `_zstd.ZstdDecompressor.decompress`. - #[pyarg(any)] - data: ArgBytesLike, - #[pyarg(any, default = -1)] - max_length: isize, - } - - /// Drive `decompress_stream` until either the frame ends, the input is - /// exhausted with no more output coming, or `max_length` bytes have been - /// produced. Sets the various `state` flags to reflect the new situation. - /// - /// Loop control: we always keep going while either the input still has - /// bytes to feed OR the previous call filled the output buffer (which - /// indicates libzstd had more to emit but ran out of room). We only stop - /// short of a frame boundary when the input is exhausted AND libzstd had - /// room left in the output buffer, which means it is genuinely waiting - /// for more compressed bytes. - fn do_decompress( - state: &mut DecompressorState, - new_data: &[u8], - max_length: Option, - vm: &VirtualMachine, - ) -> PyResult> { - // Combine any buffered leftover input with the new data so the - // decompressor sees one contiguous stream. `Cow` avoids the - // allocation when there is no leftover. - let work_data: alloc::borrow::Cow<'_, [u8]> = if state.input_buffer.is_empty() { - alloc::borrow::Cow::Borrowed(new_data) - } else { - let mut combined = Vec::with_capacity(state.input_buffer.len() + new_data.len()); - combined.extend_from_slice(&state.input_buffer); - combined.extend_from_slice(new_data); - alloc::borrow::Cow::Owned(combined) - }; - - let chunk_size = ZSTD_DStreamOutSize().max(1); - // Release the GIL for the streaming loop. Safety: see `do_compress`; - // the closure captures only Rust-owned buffers and `&mut state` (held - // under the decompressor's `PyMutex`), and surfaces error codes as - // `usize` so we can build the exception after re-attaching. - let loop_result: Result<(Vec, bool, usize), usize> = vm.allow_threads(|| { - let mut input = ZSTD_inBuffer { - src: work_data.as_ptr().cast(), - size: work_data.len(), - pos: 0, - }; - let mut output: Vec = Vec::new(); - // Reusable scratch buffer for each decompress_stream call. We need - // an exact-size output buffer because `Vec::reserve` may - // over-allocate; reporting the full Vec capacity to libzstd would - // let it write past `max_length`. A cap below one chunk bounds - // how much of it can ever be used, so do not allocate more. - let scratch_size = match max_length { - Some(maxl) => maxl.min(chunk_size), - None => chunk_size, - }; - let mut scratch: Vec = vec![0u8; scratch_size]; - let mut hit_max = false; - let mut iteration = 0usize; - - let outcome = loop { - iteration += 1; - // Honor `max_length`: stop growing the output buffer once - // we have produced enough. When the cap is zero, hand - // libzstd a zero-size buffer instead: it then consumes - // input without emitting, so zero-output frames (skippable - // frame, empty content frame) still complete while real - // output stays inside libzstd until a later call has room. - // CPython's decompressor uses the same zero-size mechanism. - let grow = match max_length { - Some(maxl) if output.len() >= maxl && iteration > 1 => { - hit_max = true; - break Ok(()); - } - Some(maxl) if output.len() >= maxl => 0, - Some(maxl) => (maxl - output.len()).min(chunk_size), - None => chunk_size, - }; - let code; - let written; - { - let slot = &mut scratch[..grow]; - let mut out_buf = ZSTD_outBuffer { - dst: slot.as_mut_ptr().cast(), - size: slot.len(), - pos: 0, - }; - // SAFETY: `state.dctx` is a live context; `out_buf` - // covers `slot`; `input` points into `work_data`, which - // outlives the call. - code = unsafe { ZSTD_decompressStream(state.dctx.0, &mut out_buf, &mut input) }; - written = out_buf.pos; - } - output.extend_from_slice(&scratch[..written]); - if ZSTD_isError(code) != 0 { - break Err(code); + #[pyclass(with(Constructor))] + impl ZstdDecompressor { + #[pymethod] + fn decompress( + zelf: &Py, + args: DecompressorArgs, + vm: &VirtualMachine, + ) -> PyResult> { + if zelf.eof.load(Ordering::Relaxed) { + return Err(vm.new_eof_error("Already at the end of a Zstandard frame.")); + } + let max_length = args.max_length(); + let data = args.data().to_vec(); + let outcome = vm.allow_threads(|| { + let mut state = zelf.state.lock(); + if state.is_eof() { + return Err(DecFail::Eof); } - if code == 0 { - // Frame fully decompressed; the decompressor is at EOF. - state.eof = true; - break Ok(()); + match state.pull(&data, max_length) { + Ok(bytes) => { + zelf.eof.store(state.is_eof(), Ordering::Relaxed); + zelf.needs_input.store(state.needs_input, Ordering::Relaxed); + Ok(bytes) + } + Err(err) => { + state.reset_session(); + zelf.eof.store(false, Ordering::Relaxed); + zelf.needs_input.store(true, Ordering::Relaxed); + Err(DecFail::Codec(err)) + } } - // Only meaningful with a non-zero buffer: libzstd had more - // to emit but ran out of room. A zero-size call (the - // max_length == 0 probe) is not "full". - let output_was_full = grow > 0 && written == grow; - let input_consumed = input.pos == input.size; - - if let Some(maxl) = max_length - && output.len() >= maxl - && iteration > 1 - { - hit_max = true; - break Ok(()); + }); + match outcome { + Ok(bytes) => Ok(bytes), + Err(DecFail::Eof) => { + Err(vm.new_eof_error("Already at the end of a Zstandard frame.")) } + Err(DecFail::Codec(err)) => Err(decompress_error(err, vm)), + } + } - // Input is gone and libzstd had room to write but did - // not, which means the frame is incomplete and the - // caller has to supply more input. - if input_consumed && !output_was_full { - break Ok(()); - } + #[pygetset] + fn unused_data(zelf: &Py, vm: &VirtualMachine) -> PyBytesRef { + let unused = { + let state = zelf.state.lock(); + state.unused().to_vec() }; - outcome.map(|()| (output, hit_max, input.pos)) - }); - - let (output, hit_max, consumed) = match loop_result { - Ok(v) => v, - Err(code) => { - state.reset_session(vm); - return Err(catch_zstd_error(code, vm)); + if unused.is_empty() { + vm.ctx.empty_bytes.clone() + } else { + vm.ctx.new_bytes(unused) } - }; + } + } - let remaining = &work_data[consumed..]; + enum SampleSize { + Value(usize), + Mismatch, + } - if state.eof { - if !remaining.is_empty() { - state.unused_data = vm.ctx.new_bytes(remaining.to_vec()); + fn sample_size(obj: &PyObject, vm: &VirtualMachine) -> PyResult { + let int_obj = match obj.try_index(vm) { + Ok(int_obj) => int_obj, + Err(err) if err.fast_isinstance(vm.ctx.exceptions.overflow_error) => { + return Ok(SampleSize::Mismatch); } - state.input_buffer.clear(); - state.needs_input = false; - } else if hit_max { - // Output cap reached with input still pending. Buffer the rest - // and report `needs_input == false` so the caller knows to call - // `decompress(b'', max_length=...)` to drain it. - state.input_buffer = remaining.to_vec(); - state.needs_input = false; - } else if max_length == Some(0) { - // Caller explicitly asked for zero output bytes. Keep whatever - // input is left around for the next call and signal that they - // do not need to feed more right now. CPython's decompressor - // treats `max_length=0` as "stop here without losing state". - state.input_buffer = remaining.to_vec(); - state.needs_input = false; - } else { - // All input consumed but the frame is not complete; the caller - // should provide more data on the next call. - state.input_buffer.clear(); - state.needs_input = true; + Err(err) => return Err(err), + }; + match int_obj.try_to_primitive_in_range::(vm) { + Ok(value) => Ok(SampleSize::Value(value)), + Err(err) if err.fast_isinstance(vm.ctx.exceptions.overflow_error) => { + Ok(SampleSize::Mismatch) + } + Err(err) => Err(err), } - - Ok(output) } - #[pyclass(with(Constructor))] - impl ZstdDecompressor { - #[pymethod] - fn decompress(&self, args: DecompressMethodArgs, vm: &VirtualMachine) -> PyResult> { - let data_vec = args.data.with_ref(|b| b.to_vec()); - let max_length = if args.max_length < 0 { - None - } else { - Some(args.max_length as usize) + fn split_samples<'a>( + bytes: &'a [u8], + sizes: &PyTupleRef, + vm: &VirtualMachine, + ) -> PyResult> { + if u32::try_from(sizes.as_slice().len()).is_err() { + return Err( + vm.new_value_error(format!("The number of samples should be <= {}.", u32::MAX)) + ); + } + let mut samples = Vec::with_capacity(sizes.as_slice().len()); + let mut offset = 0usize; + for item in sizes.as_slice() { + let size = match sample_size(item, vm)? { + SampleSize::Mismatch => return Err(vm.new_value_error(SAMPLE_MISMATCH)), + SampleSize::Value(size) => size, }; - let mut state = self.state.lock(); - if state.eof { - return Err(vm.new_exception_msg( - vm.ctx.exceptions.eof_error.to_owned(), - "Already at the end of a Zstandard frame.".to_owned().into(), - )); + let Some(next) = offset.checked_add(size) else { + return Err(vm.new_value_error(SAMPLE_MISMATCH)); + }; + if next > bytes.len() { + return Err(vm.new_value_error(SAMPLE_MISMATCH)); } - do_decompress(&mut state, &data_vec, max_length, vm) - } - - #[pygetset] - fn eof(&self) -> bool { - self.state.lock().eof + samples.push(&bytes[offset..next]); + offset = next; } - - #[pygetset] - fn needs_input(&self) -> bool { - self.state.lock().needs_input + if offset != bytes.len() { + return Err(vm.new_value_error(SAMPLE_MISMATCH)); } + Ok(samples) + } - #[pygetset] - fn unused_data(&self) -> PyBytesRef { - self.state.lock().unused_data.clone() - } + fn dictionary_failure(finalize: bool, vm: &VirtualMachine) -> PyBaseExceptionRef { + let verb = if finalize { "finalize" } else { "train" }; + zstd_error( + format!("Unable to {verb} the Zstandard dictionary: Src size is incorrect"), + vm, + ) } - #[pyfunction] - fn get_frame_size(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { - let buf = frame_buffer.with_ref(|b| b.to_vec()); - // SAFETY: `buf` is a live slice for the duration of the call. - let size = unsafe { ZSTD_findFrameCompressedSize(buf.as_ptr().cast(), buf.len()) }; - if ZSTD_isError(size) != 0 { - return Err(new_zstd_error( - "Error when finding the compressed size of a Zstandard frame. \ - Ensure the frame_buffer argument starts from the beginning of a frame, \ - and its length not less than this complete frame.", - vm, - )); - } - Ok(size) + fn forced_dict_id(custom: &[u8]) -> u32 { + let id = read_dict_id(custom).wrapping_add(1); + if id == 0 { 1 } else { id } } - #[pyfunction] - fn get_frame_info( - frame_buffer: ArgBytesLike, + fn train_or_finalize( + custom: Option<&[u8]>, + samples_bytes: &[u8], + samples_sizes: &PyTupleRef, + dict_size: isize, vm: &VirtualMachine, - ) -> PyResult<(PyObjectRef, u32)> { - let buf = frame_buffer.with_ref(|b| b.to_vec()); - // SAFETY: `buf` is a live slice for the duration of the call. - let content_size = unsafe { ZSTD_getFrameContentSize(buf.as_ptr().cast(), buf.len()) }; - if content_size == ZSTD_CONTENTSIZE_ERROR { - return Err(new_zstd_error( - "Error when getting information from the header of a Zstandard frame. \ - Ensure the frame_buffer argument starts from the beginning of a frame, \ - and its length not less than the frame header (6~18 bytes).", - vm, - )); + ) -> PyResult> { + if dict_size <= 0 { + return Err(vm.new_value_error(DICT_SIZE_NONPOSITIVE)); } - let content_size_obj: PyObjectRef = if content_size == ZSTD_CONTENTSIZE_UNKNOWN { - vm.ctx.none() - } else { - vm.ctx.new_int(content_size).into() + let samples = split_samples(samples_bytes, samples_sizes, vm)?; + let Ok(max_dict) = usize::try_from(dict_size) else { + return Err(vm.new_value_error(DICT_SIZE_NONPOSITIVE)); + }; + let dict_id = custom.map(forced_dict_id); + let opts = TrainOptions { + max_dict, + dict_id, + ..TrainOptions::fastcover() }; - // SAFETY: `buf` is a live slice for the duration of the call. - let dict_id = unsafe { ZSTD_getDictID_fromFrame(buf.as_ptr().cast(), buf.len()) }; - Ok((content_size_obj, dict_id)) + match vm.allow_threads(|| train(&samples, opts)) { + Ok(bytes) => Ok(bytes), + Err(_) => Err(dictionary_failure(custom.is_some(), vm)), + } } - #[derive(FromArgs)] - pub(super) struct TrainDictArgs { - /// Concatenated sample bytes. Must be a `bytes` object, not - /// `bytearray` or another buffer type, to match CPython's strict - /// type-checking on this argument. - #[pyarg(positional)] + #[pyfunction] + fn train_dict( samples_bytes: PyBytesRef, - /// A tuple of integer sample sizes that partition `samples_bytes`. - /// Lists and other iterables are not accepted. - #[pyarg(positional)] samples_sizes: PyTupleRef, - /// Maximum size of the returned dictionary, in bytes. Must be a - /// positive `int`. - #[pyarg(positional)] - dict_size: PyObjectRef, - } - - /// Collect the elements of `tuple` into a `Vec`, validating that - /// each element is a non-negative int that fits in `usize`. Used by both - /// `train_dict` and `finalize_dict` for the `samples_sizes` argument. - /// - /// Floats (and any object whose `__index__` slot is missing) raise - /// `TypeError`; values that do not fit `usize` raise `ValueError` so - /// the test suite's `(2**1000,)` / `(-1,)` coverage holds. - fn parse_sample_sizes(tuple: PyTupleRef, vm: &VirtualMachine) -> PyResult> { - let items = tuple.as_slice(); - let mut out = Vec::with_capacity(items.len()); - for item in items { - let idx = item.try_index(vm)?; - let v: usize = idx - .try_to_primitive(vm) - .map_err(|_| vm.new_value_error("sample size out of range for size_t"))?; - out.push(v); - } - Ok(out) - } - - /// Convert a Python `int` to a positive `isize`. Rejects floats (via - /// `try_index`) and non-positive values; bigints that don't fit `isize` - /// propagate as `OverflowError`. Used for the `dict_size` argument of - /// `train_dict` and `finalize_dict`, which must always be a strictly - /// positive int. - fn parse_positive_dict_size(obj: &PyObjectRef, vm: &VirtualMachine) -> PyResult { - let idx = obj.try_index(vm)?; - // `try_to_primitive::` raises `OverflowError` on bigints out of - // range; pass that through verbatim so the test suite's - // `assertRaises(OverflowError)` coverage matches. - let v: isize = idx.try_to_primitive(vm)?; - if v <= 0 { - return Err(vm.new_value_error("dict_size must be positive")); - } - Ok(v as usize) - } - - /// Sum the per-sample sizes and check they exactly cover `expected_total`, - /// rejecting overflow. Any call into libzstd with these sizes must not - /// let a wrapping sum sneak past the equality check, since libzstd would - /// then read past the samples buffer. - fn check_sample_sizes_match( - sizes: &[usize], - expected_total: usize, + dict_size: isize, vm: &VirtualMachine, - ) -> PyResult<()> { - let mismatch = || -> PyBaseExceptionRef { - vm.new_value_error("The samples size tuple doesn't match the concatenation's size") - }; - let total = sizes - .iter() - .try_fold(0usize, |a, &b| a.checked_add(b)) - .ok_or_else(mismatch)?; - if total != expected_total { - return Err(mismatch()); - } - Ok(()) + ) -> PyResult> { + train_or_finalize( + None, + samples_bytes.as_bytes(), + &samples_sizes, + dict_size, + vm, + ) } #[pyfunction] - fn train_dict(args: TrainDictArgs, vm: &VirtualMachine) -> PyResult> { - let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; - let samples_buffer = args.samples_bytes.as_bytes().to_vec(); - let sizes = parse_sample_sizes(args.samples_sizes, vm)?; - check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; - // `dict_size` is caller-controlled, so reserve fallibly: an - // infallible `Vec::with_capacity` aborts the whole interpreter on a - // failed allocation, where CPython raises `MemoryError`. - let mut dict_buffer: Vec = Vec::new(); - dict_buffer - .try_reserve_exact(dict_size) - .map_err(|_| vm.new_memory_error("failed to allocate the dictionary buffer"))?; - // SAFETY: `dict_buffer`'s spare capacity is valid for writes of up to - // `dict_size` bytes; the samples and sizes buffers are live slices - // for the duration of the call. - let written = unsafe { - ZDICT_trainFromBuffer( - dict_buffer.as_mut_ptr().cast(), - dict_size, - samples_buffer.as_ptr().cast(), - sizes.as_ptr(), - sizes.len() as u32, - ) - }; - if ZDICT_isError(written) != 0 { - // SAFETY: `ZDICT_getErrorName` returns a pointer to a static, - // NUL-terminated error string from libzstd's error table. - let name = unsafe { CStr::from_ptr(ZDICT_getErrorName(written)) }; - return Err(new_zstd_error(name.to_string_lossy(), vm)); - } - // SAFETY: `ZDICT_trainFromBuffer` wrote `written` (<= dict_size) - // bytes into the buffer's spare capacity. - unsafe { dict_buffer.set_len(written) }; - Ok(dict_buffer) - } - - #[derive(FromArgs)] - pub(super) struct FinalizeDictArgs { - /// Raw "starting" dictionary content to finalize. Must be `bytes` - /// (not `bytearray`) to match CPython. - #[pyarg(positional)] + fn finalize_dict( custom_dict_bytes: PyBytesRef, - /// Concatenated sample bytes used to derive the dictionary's - /// statistics tables. Must be `bytes`. - #[pyarg(positional)] samples_bytes: PyBytesRef, - /// Tuple of integer sample sizes partitioning `samples_bytes`. - #[pyarg(positional)] samples_sizes: PyTupleRef, - /// Maximum size of the finalized dictionary, in bytes. Positive int. - #[pyarg(positional)] - dict_size: PyObjectRef, - /// Compression level the dictionary will be tuned for. Must be int. - #[pyarg(positional)] - compression_level: PyObjectRef, - } - - #[pyfunction] - fn finalize_dict(args: FinalizeDictArgs, vm: &VirtualMachine) -> PyResult> { - let dict_size = parse_positive_dict_size(&args.dict_size, vm)?; - let compression_level: i32 = args.compression_level.try_to_value(vm)?; - let custom_dict = args.custom_dict_bytes.as_bytes().to_vec(); - let samples_buffer = args.samples_bytes.as_bytes().to_vec(); - let sizes = parse_sample_sizes(args.samples_sizes, vm)?; - check_sample_sizes_match(&sizes, samples_buffer.len(), vm)?; - - // Reserve fallibly before filling: see `train_dict`. - let mut dict_buffer: Vec = Vec::new(); - dict_buffer - .try_reserve_exact(dict_size) - .map_err(|_| vm.new_memory_error("failed to allocate the dictionary buffer"))?; - dict_buffer.resize(dict_size, 0); - let params = ZDICT_params_t { - compressionLevel: compression_level, - notificationLevel: 0, - dictID: 0, - }; - - // SAFETY: All pointers point into Rust-owned, properly sized buffers - // that outlive the FFI call. ZDICT_finalizeDictionary just reads from - // the sample/dict buffers and writes into `dict_buffer`. - let written = unsafe { - ZDICT_finalizeDictionary( - dict_buffer.as_mut_ptr() as *mut _, - dict_buffer.len(), - custom_dict.as_ptr() as *const _, - custom_dict.len(), - samples_buffer.as_ptr() as *const _, - sizes.as_ptr(), - sizes.len() as u32, - params, - ) - }; - if ZDICT_isError(written) != 0 { - let err_ptr = ZDICT_getErrorName(written); - let msg = if err_ptr.is_null() { - "zstd dictionary finalization failed".to_string() - } else { - // SAFETY: `ZDICT_getErrorName` returns a pointer to a static, - // NUL-terminated error string from libzstd's error table. - unsafe { CStr::from_ptr(err_ptr) } - .to_string_lossy() - .into_owned() - }; - return Err(new_zstd_error(msg, vm)); - } - dict_buffer.truncate(written); - Ok(dict_buffer) + dict_size: isize, + compression_level: i32, + vm: &VirtualMachine, + ) -> PyResult> { + let _ = compression_level; + train_or_finalize( + Some(custom_dict_bytes.as_bytes()), + samples_bytes.as_bytes(), + &samples_sizes, + dict_size, + vm, + ) } - #[derive(FromArgs)] - pub(super) struct ParamBoundsArgs { - #[pyarg(any)] - parameter: i32, - /// Positional-or-keyword, matching CPython's clinic signature - /// `get_param_bounds(parameter, is_compress)`. - #[pyarg(any)] - is_compress: bool, + fn frame_info_error(vm: &VirtualMachine) -> PyBaseExceptionRef { + zstd_error( + "Error when getting information from the header of a Zstandard frame. \ + Ensure the frame_buffer argument starts from the beginning of a frame, \ + and its length is not less than the frame header (6~18 bytes)." + .to_owned(), + vm, + ) } #[pyfunction] - fn get_param_bounds(args: ParamBoundsArgs, vm: &VirtualMachine) -> PyResult<(c_int, c_int)> { - let unknown = || -> PyBaseExceptionRef { - let kind = if args.is_compress { - "compression" - } else { - "decompression" - }; - vm.new_value_error(format!( - "invalid {kind} parameter 'unknown parameter (key {})'", - args.parameter - )) - }; - // Validate the id via the same safe enum-lookup helpers used in - // `lookup_param_bounds`, then call libzstd directly so we can - // distinguish a libzstd-reported error from our own "unknown". The - // two `getBounds` functions return distinct (but layout-identical) - // `ZSTD_bounds` types, so destructure into a tuple right away. - let (error, lo, hi) = if args.is_compress { - let p = c_param_enum(args.parameter).ok_or_else(unknown)?; - let b = ZSTD_cParam_getBounds(p); - (b.error, b.lowerBound, b.upperBound) - } else { - let p = d_param_enum(args.parameter).ok_or_else(unknown)?; - let b = ZSTD_dParam_getBounds(p); - (b.error, b.lowerBound, b.upperBound) + fn get_frame_info(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { + let kind = frame_buffer.with_ref(get_frame_header); + let kind = kind.map_err(|_| frame_info_error(vm))?; + let (size, dict_id) = match kind { + FrameKind::Zstd(header) => { + let size = match header.content_size { + Some(value) => vm.ctx.new_int(value).into(), + None => vm.ctx.none(), + }; + (size, header.dict_id.unwrap_or(0)) + } + FrameKind::Skippable { .. } => (vm.ctx.none(), 0), }; - if ZSTD_isError(error) != 0 { - return Err(catch_zstd_error(error, vm)); - } - Ok((lo, hi)) + Ok(vm.ctx.new_tuple(vec![size, vm.ctx.new_int(dict_id).into()])) } - // Register the `CompressionParameter` / `DecompressionParameter` enum - // classes defined by the pure-Python wrapper so [`check_wrong_param_kind`] - // can reject a key from the wrong parameter family by identity. The types - // are stashed as private `_zstd` module attributes — the RustPython - // equivalent of the module state CPython keeps these in. The wrapper in - // `Lib/compression/zstd/__init__.py` calls this exactly once at import. #[pyfunction] - fn set_parameter_types( - c_parameter_type: PyTypeRef, - d_parameter_type: PyTypeRef, - vm: &VirtualMachine, - ) -> PyResult<()> { - let module = vm.import("_zstd", 0)?; - module.set_attr("_compression_parameter_type", c_parameter_type, vm)?; - module.set_attr("_decompression_parameter_type", d_parameter_type, vm)?; - Ok(()) + fn get_frame_size(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { + let result = frame_buffer.with_ref(find_frame_compressed_size); + match result { + Ok(size) => Ok(size), + Err(err) => Err(zstd_error( + format!( + "Error when finding the compressed size of a Zstandard frame. \ + Ensure the frame_buffer argument starts from the beginning of a frame, \ + and its length is not less than this complete frame. \ + Zstd error message: {}.", + frame_size_reason(err) + ), + vm, + )), + } } } From d9130bd6b7ce71eb04ec1ed1c63dfbe716e1cc6d Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Fri, 2 Oct 2026 00:47:30 +0900 Subject: [PATCH 22/25] Use libzstd-rs-sys for _zstd Replace rusty_zstd 0.2.5 with libzstd-rs-sys rev 7afaf79f14f4. Streaming compression, decompression, frame queries, and dictionary training go through that crate. zstd_version is 1.5.8. Assisted-by: Grok:grok-4.7 --- Cargo.lock | 17 +- Cargo.toml | 2 +- crates/stdlib/Cargo.toml | 2 +- crates/stdlib/src/zstd.rs | 1703 +++++++++++++++++++------------------ 4 files changed, 882 insertions(+), 842 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bf9705c85ad..13760b77871 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2321,6 +2321,15 @@ dependencies = [ "zlib-rs", ] +[[package]] +name = "libzstd-rs-sys" +version = "0.0.1-prerelease.2" +source = "git+https://github.com/trifectatechfoundation/libzstd-rs-sys?rev=7afaf79f14f4a5ba4722f66d83791d30ebc47e65#7afaf79f14f4a5ba4722f66d83791d30ebc47e65" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -4024,6 +4033,7 @@ dependencies = [ "itertools 0.15.0", "libc", "libsqlite3-sys", + "libzstd-rs-sys", "malachite-bigint", "memchr", "mt19937", @@ -4051,7 +4061,6 @@ dependencies = [ "rustpython-ruff_text_size", "rustpython-unicode", "rustpython-vm", - "rusty_zstd", "scopeguard", "tcl-sys", "tk-sys", @@ -4180,12 +4189,6 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" -[[package]] -name = "rusty_zstd" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc2318a1441e10cf436f0bf98615216b0a5410b6f546548b50faa170c7b4a1a6" - [[package]] name = "rustyline" version = "18.0.1" diff --git a/Cargo.toml b/Cargo.toml index 449dbda25dd..60d1b1682c8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -241,7 +241,7 @@ xz-sys = "0.4.6-rc.0" libsqlite3-sys = "0.38" libz-rs-sys = "0.6" zlib-rs = { version = "=0.6.8", default-features = false, features = ["rust-allocator", "__internal-api"] } -rusty_zstd = "0.2.5" +libzstd-rs-sys = { git = "https://github.com/trifectatechfoundation/libzstd-rs-sys", rev = "7afaf79f14f4a5ba4722f66d83791d30ebc47e65" } lock_api = "0.4" log = "0.4.30" lz4_flex = "0.14" diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index 1e095fb9663..b6cd0ba4841 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -71,7 +71,7 @@ constant_time_eq = { workspace = true } # compression adler32 = { workspace = true } -rusty_zstd = { workspace = true } +libzstd-rs-sys = { workspace = true } # tkinter tk-sys = { workspace = true, optional = true } diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index e2d240dda27..69fd96ae0a4 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -9,7 +9,7 @@ pub(crate) use _zstd::module_def; #[pymodule] mod _zstd { - #![allow(non_upper_case_globals)] + #![allow(non_upper_case_globals, clippy::upper_case_acronyms)] use crate::compression::DecompressorArgs; use crate::vm::{ @@ -23,45 +23,60 @@ mod _zstd { protocol::PyMappingMethods, types::{AsMapping, Constructor, Representable}, }; - use alloc::boxed::Box; + use core::ffi::{CStr, c_void}; use core::fmt; use core::sync::atomic::{AtomicBool, AtomicI32, Ordering}; - use num_traits::ToPrimitive; - use rusty_zstd::{ - CompressOptions, Compressor, DEFAULT_CLEVEL, DEFAULT_WINDOW_MAX, DecompressOptions, - Decompressor, Dictionary, Error, Flush, FrameKind, MAGIC, MAGIC_DICTIONARY, - MAGIC_SKIPPABLE_MAX, MAGIC_SKIPPABLE_MIN, MAX_CLEVEL, MIN_CLEVEL, TrainOptions, - compress_stream_out_size, compression_params, decompress_stream_out_size, - find_frame_compressed_size, get_frame_header, train, + use libzstd_rs_sys::lib::zdict::{ZDICT_finalizeDictionary, ZDICT_params_t}; + use libzstd_rs_sys::lib::zstd::{ZSTD_e_continue, ZSTD_e_end, ZSTD_e_flush}; + use libzstd_rs_sys::{ + ZDICT_isError, ZDICT_trainFromBuffer, ZSTD_CCtx, ZSTD_CCtx_loadDictionary, + ZSTD_CCtx_refCDict, ZSTD_CCtx_refPrefix, ZSTD_CCtx_reset, ZSTD_CCtx_setParameter, + ZSTD_CCtx_setPledgedSrcSize, ZSTD_CONTENTSIZE_ERROR, ZSTD_CONTENTSIZE_UNKNOWN, + ZSTD_DCtx_loadDictionary, ZSTD_DCtx_refDDict, ZSTD_DCtx_refPrefix, ZSTD_DCtx_reset, + ZSTD_DCtx_setParameter, ZSTD_DStreamOutSize as dstream_out_size, ZSTD_ResetDirective, + ZSTD_VERSION_NUMBER, ZSTD_cParam_getBounds, ZSTD_cParameter, ZSTD_compressBound, + ZSTD_compressStream2, ZSTD_createCCtx, ZSTD_createCDict, ZSTD_createDCtx, ZSTD_createDDict, + ZSTD_dParam_getBounds, ZSTD_dParameter, ZSTD_decompressStream, + ZSTD_findFrameCompressedSize, ZSTD_freeCCtx, ZSTD_freeCDict, ZSTD_freeDCtx, ZSTD_freeDDict, + ZSTD_getDictID_fromDict, ZSTD_getDictID_fromFrame, ZSTD_getErrorName, + ZSTD_getFrameContentSize, ZSTD_inBuffer, ZSTD_isError, ZSTD_maxCLevel, ZSTD_minCLevel, + ZSTD_outBuffer, }; + use num_traits::ToPrimitive; + use std::collections::HashMap; + + const _: () = assert!(ZSTD_VERSION_NUMBER == 10_508); + const _: () = assert!(dstream_out_size() == 131_072); const MODE_CONTINUE: i32 = 0; const MODE_BLOCK: i32 = 1; const MODE_FRAME: i32 = 2; - const LEVEL_MIN: i32 = -131_072; - const LEVEL_MAX: i32 = 22; + const LEVEL_MIN: i32 = ZSTD_minCLevel(); + const LEVEL_MAX: i32 = ZSTD_maxCLevel(); const DIGESTED: i32 = 0; const UNDIGESTED: i32 = 1; const PREFIX: i32 = 2; - const WINDOW_LOG_MIN: i32 = 10; - const WINDOW_LOG_MAX: i32 = if cfg!(target_pointer_width = "32") { - 30 - } else { - 31 - }; - const HASH_LOG_MIN: i32 = 6; - const HASH_LOG_MAX: i32 = 30; - const CHAIN_LOG_MAX: i32 = if cfg!(target_pointer_width = "32") { - 29 - } else { - 30 - }; - const SEARCH_LOG_MIN: i32 = 1; - const SEARCH_LOG_MAX: i32 = WINDOW_LOG_MAX - 1; - const LDM_HASH_RATE_MAX: i32 = WINDOW_LOG_MAX - HASH_LOG_MIN; - const DEFAULT_WINDOW_LOG_MAX: i32 = 27; + const OUTPUT_BLOCKS: [usize; 17] = [ + 32 * 1024, + 64 * 1024, + 256 * 1024, + 1024 * 1024, + 4 * 1024 * 1024, + 8 * 1024 * 1024, + 16 * 1024 * 1024, + 16 * 1024 * 1024, + 32 * 1024 * 1024, + 32 * 1024 * 1024, + 32 * 1024 * 1024, + 32 * 1024 * 1024, + 64 * 1024 * 1024, + 64 * 1024 * 1024, + 128 * 1024 * 1024, + 128 * 1024 * 1024, + 256 * 1024 * 1024, + ]; const SAMPLE_MISMATCH: &str = "The samples size tuple doesn't match the concatenation's size."; const DICT_SIZE_NONPOSITIVE: &str = "dict_size argument should be positive number."; @@ -69,15 +84,16 @@ mod _zstd { "Failed to create a ZSTD_CDict instance from Zstandard dictionary content."; const DIGESTED_DECOMPRESS_FAILED: &str = "Failed to create a ZSTD_DDict instance from Zstandard dictionary content."; + const BUFFER_FAILED: &str = "Unable to allocate output buffer."; #[pyattr] - const zstd_version: &str = "1.5.7"; + const zstd_version: &str = "1.5.8"; #[pyattr] - const zstd_version_number: i32 = 10_507; + const zstd_version_number: i32 = ZSTD_VERSION_NUMBER as i32; #[pyattr] - const ZSTD_CLEVEL_DEFAULT: i32 = DEFAULT_CLEVEL; + const ZSTD_CLEVEL_DEFAULT: i32 = libzstd_rs_sys::ZSTD_CLEVEL_DEFAULT; #[pyattr] - const ZSTD_DStreamOutSize: i32 = 131_072; + const ZSTD_DStreamOutSize: i32 = dstream_out_size() as i32; #[pyattr] const ZSTD_c_compressionLevel: i32 = 100; @@ -139,6 +155,19 @@ mod _zstd { #[pyattr] const ZSTD_btultra2: i32 = 9; + const _: () = assert!( + c_parameter_id(libzstd_rs_sys::ZSTD_cParameter::ZSTD_c_compressionLevel) + == ZSTD_c_compressionLevel as u32 + ); + const _: () = assert!( + c_parameter_id(libzstd_rs_sys::ZSTD_cParameter::ZSTD_c_nbWorkers) + == ZSTD_c_nbWorkers as u32 + ); + const _: () = assert!( + d_parameter_id(libzstd_rs_sys::ZSTD_dParameter::ZSTD_d_windowLogMax) + == ZSTD_d_windowLogMax as u32 + ); + #[pyattr(once, name = "ZstdError")] fn zstd_error_type(vm: &VirtualMachine) -> PyTypeRef { vm.ctx.new_exception_type_with_doc( @@ -153,236 +182,262 @@ mod _zstd { vm.new_exception_msg(vm.class("_zstd", "ZstdError"), message.into()) } - fn compress_reason(err: Error) -> &'static str { - match err { - Error::ContentSizeMismatch | Error::UnexpectedEof => "Src size is incorrect", - Error::BadMagic => "Unknown frame descriptor", - Error::ChecksumMismatch => "Restored data doesn't match checksum", - Error::WindowTooLarge | Error::ContentSizeTooLarge => { - "Frame requires too much memory for decoding" - } - Error::ReservedBitSet | Error::UnusedBitSet => "Unsupported frame parameter", - Error::DictionaryNeeded { .. } | Error::DictionaryMismatch { .. } => { - "Dictionary mismatch" - } - Error::ReservedBlockType - | Error::BlockTooLarge - | Error::TrailingBytes - | Error::Unimplemented - | Error::Corruption - | Error::InvalidLevel => "Data corruption detected", - } + enum CallKind { + Compress, + Decompress, + Pledge, + LoadCompress, + LoadDecompress, + Level, + Train, + Finalize, + BoundsCompress, + BoundsDecompress, + FrameSize, + } + + enum CodecError { + NoMemory, + Buffer, + Static(&'static str), + Named(CallKind, usize), + Eof, } - fn decompress_reason(err: Error) -> &'static str { - match err { - Error::BadMagic => "Unknown frame descriptor", - Error::ChecksumMismatch => "Restored data doesn't match checksum", - Error::WindowTooLarge | Error::ContentSizeTooLarge => { - "Frame requires too much memory for decoding" - } - Error::ReservedBitSet | Error::UnusedBitSet => "Unsupported frame parameter", - Error::DictionaryNeeded { .. } | Error::DictionaryMismatch { .. } => { - "Dictionary mismatch" - } - Error::ContentSizeMismatch - | Error::UnexpectedEof - | Error::ReservedBlockType - | Error::BlockTooLarge - | Error::TrailingBytes - | Error::Unimplemented - | Error::Corruption - | Error::InvalidLevel => "Data corruption detected", + fn error_name(code: usize) -> String { + let ptr = ZSTD_getErrorName(code); + if ptr.is_null() { + return String::new(); } + unsafe { CStr::from_ptr(ptr) } + .to_string_lossy() + .into_owned() } - fn compress_error(err: Error, vm: &VirtualMachine) -> PyBaseExceptionRef { - zstd_error( - format!( - "Unable to compress Zstandard data: {}", - compress_reason(err) - ), - vm, - ) - } - - fn decompress_error(err: Error, vm: &VirtualMachine) -> PyBaseExceptionRef { - zstd_error( - format!( - "Unable to decompress Zstandard data: {}", - decompress_reason(err) - ), - vm, - ) - } - - fn frame_size_reason(err: Error) -> &'static str { + fn raise_codec(err: CodecError, vm: &VirtualMachine) -> PyBaseExceptionRef { match err { - Error::BadMagic => "Unknown frame descriptor", - Error::UnexpectedEof | Error::ContentSizeMismatch => "Src size is incorrect", - other => decompress_reason(other), + CodecError::NoMemory => vm.no_memory_error(), + CodecError::Buffer => vm.new_memory_error(BUFFER_FAILED), + CodecError::Static(message) => zstd_error(message.to_owned(), vm), + CodecError::Eof => vm.new_eof_error("Already at the end of a Zstandard frame."), + CodecError::Named(kind, code) => { + let name = error_name(code); + let message = match kind { + CallKind::Compress => format!("Unable to compress Zstandard data: {name}"), + CallKind::Decompress => { + format!("Unable to decompress Zstandard data: {name}") + } + CallKind::Pledge => { + format!("Unable to set pledged uncompressed content size: {name}") + } + CallKind::LoadCompress => format!( + "Unable to load Zstandard dictionary or prefix for compression: {name}" + ), + CallKind::LoadDecompress => format!( + "Unable to load Zstandard dictionary or prefix for decompression: {name}" + ), + CallKind::Level => format!("Unable to set zstd compression level: {name}"), + CallKind::Train => { + format!("Unable to train the Zstandard dictionary: {name}") + } + CallKind::Finalize => { + format!("Unable to finalize the Zstandard dictionary: {name}") + } + CallKind::BoundsCompress => { + format!("Unable to get zstd compression parameter bounds: {name}") + } + CallKind::BoundsDecompress => { + format!("Unable to get zstd decompression parameter bounds: {name}") + } + CallKind::FrameSize => format!( + "Error when finding the compressed size of a Zstandard frame. \ + Ensure the frame_buffer argument starts from the beginning of a frame, \ + and its length is not less than this complete frame. \ + Zstd error message: {name}." + ), + }; + zstd_error(message, vm) + } } } - fn is_pylong(obj: &PyObject, vm: &VirtualMachine) -> bool { - obj.downcast_ref::().is_some() || obj.fast_isinstance(vm.ctx.types.int_type) + fn is_error(code: usize) -> bool { + ZSTD_isError(code) != 0 } - fn type_name(obj: &PyObject, vm: &VirtualMachine) -> String { - obj.class() - .fully_qualified_name(vm) - .unwrap_or_else(|_| "object".to_owned()) - } - - fn read_u32_at(buf: &[u8], offset: usize) -> Option { - let end = offset.checked_add(4)?; - let bytes = buf.get(offset..end)?; - let chunk = <[u8; 4]>::try_from(bytes).ok()?; - Some(u32::from_le_bytes(chunk)) + const fn c_parameter_id(param: ZSTD_cParameter) -> u32 { + unsafe { core::mem::transmute(param) } } - fn read_dict_id(content: &[u8]) -> u32 { - match read_u32_at(content, 0) { - Some(magic) if magic == MAGIC_DICTIONARY => read_u32_at(content, 4).unwrap_or(0), - _ => 0, - } + const fn d_parameter_id(param: ZSTD_dParameter) -> u32 { + unsafe { core::mem::transmute(param) } } - fn is_skippable(magic: u32) -> bool { - (MAGIC_SKIPPABLE_MIN..=MAGIC_SKIPPABLE_MAX).contains(&magic) + fn c_parameter(id: i32) -> ZSTD_cParameter { + unsafe { core::mem::transmute(id as u32) } } - const BLOCKSIZE_MAX_BYTES: usize = 128 * 1024; - - fn codec_level(level: i32) -> i32 { - level.clamp(MIN_CLEVEL, MAX_CLEVEL) - } - - fn block_capacity(level: i32) -> usize { - let log = - compression_params(codec_level(level), None).map_or(17, |params| params.window_log); - let window = 1usize.checked_shl(log.min(31)).unwrap_or(usize::MAX); - window.min(BLOCKSIZE_MAX_BYTES) - } - - fn window_max_bytes(log: i32) -> u64 { - let Ok(log) = u32::try_from(log) else { - return DEFAULT_WINDOW_MAX; - }; - 1u64.checked_shl(log).unwrap_or(u64::MAX) + fn d_parameter(id: i32) -> ZSTD_dParameter { + unsafe { core::mem::transmute(id as u32) } } - fn flush_kind(mode: i32) -> Flush { + fn end_op(mode: i32) -> libzstd_rs_sys::ZSTD_EndDirective { match mode { - MODE_BLOCK => Flush::Flush, - MODE_FRAME => Flush::End, - _ => Flush::Continue, + MODE_BLOCK => ZSTD_e_flush, + MODE_FRAME => ZSTD_e_end, + _ => ZSTD_e_continue, } } - fn check_compress_mode(mode: i32, vm: &VirtualMachine) -> PyResult<()> { - if matches!(mode, MODE_CONTINUE | MODE_BLOCK | MODE_FRAME) { - Ok(()) + fn ffi_bytes(bytes: &[u8], scratch: &u8) -> (*const c_void, usize) { + if bytes.is_empty() { + (core::ptr::from_ref(scratch).cast(), 0) } else { - Err(vm.new_value_error( - "mode argument wrong value, it should be one of ZstdCompressor.CONTINUE, \ - ZstdCompressor.FLUSH_BLOCK, ZstdCompressor.FLUSH_FRAME.", - )) + (bytes.as_ptr().cast(), bytes.len()) } } - fn check_flush_mode(mode: i32, vm: &VirtualMachine) -> PyResult<()> { - if matches!(mode, MODE_BLOCK | MODE_FRAME) { - Ok(()) - } else { - Err(vm.new_value_error( - "mode argument wrong value, it should be ZstdCompressor.FLUSH_FRAME or \ - ZstdCompressor.FLUSH_BLOCK.", - )) + fn in_buffer(src: &[u8], scratch: &u8) -> ZSTD_inBuffer { + let (ptr, size) = ffi_bytes(src, scratch); + ZSTD_inBuffer { + src: ptr, + size, + pos: 0, } } - fn compress_bounds(id: i32) -> Option<(&'static str, i32, i32)> { - match id { - ZSTD_c_compressionLevel => Some(("compression_level", LEVEL_MIN, LEVEL_MAX)), - ZSTD_c_windowLog => Some(("window_log", WINDOW_LOG_MIN, WINDOW_LOG_MAX)), - ZSTD_c_hashLog => Some(("hash_log", HASH_LOG_MIN, HASH_LOG_MAX)), - ZSTD_c_chainLog => Some(("chain_log", HASH_LOG_MIN, CHAIN_LOG_MAX)), - ZSTD_c_searchLog => Some(("search_log", SEARCH_LOG_MIN, SEARCH_LOG_MAX)), - ZSTD_c_minMatch => Some(("min_match", 3, 7)), - ZSTD_c_targetLength => Some(("target_length", 0, 131_072)), - ZSTD_c_strategy => Some(("strategy", ZSTD_fast, ZSTD_btultra2)), - ZSTD_c_enableLongDistanceMatching => Some(("enable_long_distance_matching", 0, 2)), - ZSTD_c_ldmHashLog => Some(("ldm_hash_log", HASH_LOG_MIN, HASH_LOG_MAX)), - ZSTD_c_ldmMinMatch => Some(("ldm_min_match", 4, 4096)), - ZSTD_c_ldmBucketSizeLog => Some(("ldm_bucket_size_log", 1, 8)), - ZSTD_c_ldmHashRateLog => Some(("ldm_hash_rate_log", 0, LDM_HASH_RATE_MAX)), - ZSTD_c_contentSizeFlag => Some(("content_size_flag", 0, 1)), - ZSTD_c_checksumFlag => Some(("checksum_flag", 0, 1)), - ZSTD_c_dictIDFlag => Some(("dict_id_flag", 0, 1)), - ZSTD_c_nbWorkers => Some(("nb_workers", 0, 256)), - ZSTD_c_jobSize => Some(("job_size", 0, 1 << 30)), - ZSTD_c_overlapLog => Some(("overlap_log", 0, 9)), - _ => None, + struct CCtx(*mut ZSTD_CCtx); + // The context moves between threads; the owning mutex is what serializes use. + unsafe impl Send for CCtx {} + impl Drop for CCtx { + fn drop(&mut self) { + if !self.0.is_null() { + unsafe { + ZSTD_freeCCtx(self.0); + } + self.0 = core::ptr::null_mut(); + } } } - fn decompress_bounds(id: i32) -> Option<(&'static str, i32, i32)> { - match id { - ZSTD_d_windowLogMax => Some(("window_log_max", WINDOW_LOG_MIN, WINDOW_LOG_MAX)), - _ => None, + struct DCtx(*mut libzstd_rs_sys::ZSTD_DCtx); + unsafe impl Send for DCtx {} + impl Drop for DCtx { + fn drop(&mut self) { + if !self.0.is_null() { + unsafe { + ZSTD_freeDCtx(self.0); + } + self.0 = core::ptr::null_mut(); + } } } - struct CompressSettings { - level: i32, - checksum: bool, - content_size: bool, - dict_id: bool, - } - - impl Default for CompressSettings { - fn default() -> Self { - Self { - level: DEFAULT_CLEVEL, - checksum: false, - content_size: true, - dict_id: true, + struct CDict(*mut libzstd_rs_sys::ZSTD_CDict); + unsafe impl Send for CDict {} + impl Drop for CDict { + fn drop(&mut self) { + if !self.0.is_null() { + unsafe { + ZSTD_freeCDict(self.0); + } + self.0 = core::ptr::null_mut(); } } } - struct DecompressSettings { - window_log_max: i32, + struct DDict(*mut libzstd_rs_sys::ZSTD_DDict); + unsafe impl Send for DDict {} + impl Drop for DDict { + fn drop(&mut self) { + if !self.0.is_null() { + unsafe { + ZSTD_freeDDict(self.0); + } + self.0 = core::ptr::null_mut(); + } + } } - impl Default for DecompressSettings { - fn default() -> Self { - Self { - window_log_max: DEFAULT_WINDOW_LOG_MAX, - } + fn known_parameter_name(compress: bool, key: i32) -> Option<&'static str> { + if compress { + let name = match key { + ZSTD_c_compressionLevel => "compression_level", + ZSTD_c_windowLog => "window_log", + ZSTD_c_hashLog => "hash_log", + ZSTD_c_chainLog => "chain_log", + ZSTD_c_searchLog => "search_log", + ZSTD_c_minMatch => "min_match", + ZSTD_c_targetLength => "target_length", + ZSTD_c_strategy => "strategy", + ZSTD_c_enableLongDistanceMatching => "enable_long_distance_matching", + ZSTD_c_ldmHashLog => "ldm_hash_log", + ZSTD_c_ldmMinMatch => "ldm_min_match", + ZSTD_c_ldmBucketSizeLog => "ldm_bucket_size_log", + ZSTD_c_ldmHashRateLog => "ldm_hash_rate_log", + ZSTD_c_contentSizeFlag => "content_size_flag", + ZSTD_c_checksumFlag => "checksum_flag", + ZSTD_c_dictIDFlag => "dict_id_flag", + ZSTD_c_nbWorkers => "nb_workers", + ZSTD_c_jobSize => "job_size", + ZSTD_c_overlapLog => "overlap_log", + _ => return None, + }; + Some(name) + } else if key == ZSTD_d_windowLogMax { + Some("window_log_max") + } else { + None + } + } + + fn query_bounds(compress: bool, key: i32) -> Result<(i32, i32), usize> { + let (error, lower, upper) = if compress { + let bounds = ZSTD_cParam_getBounds(c_parameter(key)); + (bounds.error, bounds.lowerBound, bounds.upperBound) + } else { + let bounds = ZSTD_dParam_getBounds(d_parameter(key)); + (bounds.error, bounds.lowerBound, bounds.upperBound) + }; + if is_error(error) { + Err(error) + } else { + Ok((lower, upper)) } } - fn apply_ranged( - kind: &str, - id: i32, + fn parameter_value_error( + compress: bool, + key: i32, value: i32, - bounds: Option<(&'static str, i32, i32)>, vm: &VirtualMachine, - ) -> PyResult { - let Some((name, lo, hi)) = bounds else { - return Err(vm.new_value_error(format!( - "invalid {kind} parameter 'unknown parameter (key {id})'" - ))); + ) -> PyBaseExceptionRef { + let kind = if compress { + "compression" + } else { + "decompression" }; - if value < lo || value > hi { - return Err(vm.new_value_error(format!( + let name = match known_parameter_name(compress, key) { + Some(name) => name.to_owned(), + None => format!("unknown parameter (key {key})"), + }; + match query_bounds(compress, key) { + Err(_) => vm.new_value_error(format!("invalid {kind} parameter '{name}'")), + Ok((lower, upper)) => vm.new_value_error(format!( "{kind} parameter '{name}' received an illegal value {value}; \ - the valid range is [{lo}, {hi}]" - ))); + the valid range is [{lower}, {upper}]" + )), } - Ok(id) + } + + fn is_pylong(obj: &PyObject, vm: &VirtualMachine) -> bool { + obj.downcast_ref::().is_some() || obj.fast_isinstance(vm.ctx.types.int_type) + } + + fn type_name(obj: &PyObject, vm: &VirtualMachine) -> String { + obj.class() + .fully_qualified_name(vm) + .unwrap_or_else(|_| "object".to_owned()) } fn parse_level(obj: &PyObject, vm: &VirtualMachine) -> PyResult { @@ -390,10 +445,7 @@ mod _zstd { return Err(vm.new_type_error("invalid type for level, expected int")); } match obj.try_to_value::(vm) { - Ok(level) if (LEVEL_MIN..=LEVEL_MAX).contains(&level) => Ok(level), - Ok(level) => Err(vm.new_value_error(format!( - "illegal compression level {level}; the valid range is [{LEVEL_MIN}, {LEVEL_MAX}]" - ))), + Ok(level) => Ok(level), Err(err) if err.fast_isinstance(vm.ctx.exceptions.overflow_error) => Err(vm .new_value_error(format!( "illegal compression level; the valid range is [{LEVEL_MIN}, {LEVEL_MAX}]" @@ -480,57 +532,6 @@ mod _zstd { Ok(()) } - fn parse_compress_options(obj: &PyObject, vm: &VirtualMachine) -> PyResult { - let dict = options_dict(obj, true, vm)?; - let mut settings = CompressSettings::default(); - for_each_option( - dict, - true, - |id, value| { - if id == ZSTD_c_compressionLevel { - if !(LEVEL_MIN..=LEVEL_MAX).contains(&value) { - return Err(vm.new_value_error(format!( - "illegal compression level {value}; the valid range is [{LEVEL_MIN}, {LEVEL_MAX}]" - ))); - } - settings.level = value; - return Ok(()); - } - let id = apply_ranged("compression", id, value, compress_bounds(id), vm)?; - match id { - ZSTD_c_checksumFlag => settings.checksum = value != 0, - ZSTD_c_contentSizeFlag => settings.content_size = value != 0, - ZSTD_c_dictIDFlag => settings.dict_id = value != 0, - _ => {} - } - Ok(()) - }, - vm, - )?; - Ok(settings) - } - - fn parse_decompress_options( - obj: &PyObject, - vm: &VirtualMachine, - ) -> PyResult { - let dict = options_dict(obj, false, vm)?; - let mut settings = DecompressSettings::default(); - for_each_option( - dict, - false, - |id, value| { - let id = apply_ranged("decompression", id, value, decompress_bounds(id), vm)?; - if id == ZSTD_d_windowLogMax { - settings.window_log_max = value; - } - Ok(()) - }, - vm, - )?; - Ok(settings) - } - #[pyclass(no_attr, module = "_zstd", name = "_zstd_state", traverse)] #[derive(Debug, PyPayload)] struct ZstdState { @@ -586,25 +587,23 @@ mod _zstd { parameter, is_compress, } = args; - let bounds = if is_compress { - compress_bounds(parameter) - } else { - decompress_bounds(parameter) - }; - let Some((_, lo, hi)) = bounds else { + let (lower, upper) = query_bounds(is_compress, parameter).map_err(|code| { let kind = if is_compress { - "compression" + CallKind::BoundsCompress } else { - "decompression" + CallKind::BoundsDecompress }; - return Err(zstd_error( - format!("Unable to get zstd {kind} parameter bounds: Unsupported parameter"), - vm, - )); - }; - Ok(vm - .ctx - .new_tuple(vec![vm.ctx.new_int(lo).into(), vm.ctx.new_int(hi).into()])) + raise_codec(CodecError::Named(kind, code), vm) + })?; + Ok(vm.ctx.new_tuple(vec![ + vm.ctx.new_int(lower).into(), + vm.ctx.new_int(upper).into(), + ])) + } + + struct DictCache { + cdicts: HashMap, + ddict: Option, } #[derive(FromArgs)] @@ -619,6 +618,9 @@ mod _zstd { #[pyclass(module = "compression.zstd", name = "ZstdDict", traverse)] #[derive(PyPayload)] struct ZstdDict { + // Freed before `dict_content` so digested dictionaries do not outlive the bytes. + #[pytraverse(skip)] + cache: PyMutex, dict_content: PyBytesRef, #[pytraverse(skip)] dict_id: u32, @@ -645,11 +647,16 @@ mod _zstd { "Zstandard dictionary content too short (must have at least eight bytes)", )); } - let dict_id = read_dict_id(&content); + let dict_id = + unsafe { ZSTD_getDictID_fromDict(content.as_ptr().cast(), content.len()) }; if !args.is_raw && dict_id == 0 { return Err(vm.new_value_error("invalid Zstandard dictionary")); } Ok(Self { + cache: PyMutex::new(DictCache { + cdicts: HashMap::new(), + ddict: None, + }), dict_content: vm.ctx.new_bytes(content), dict_id, }) @@ -725,6 +732,18 @@ mod _zstd { Marked(PyRef, i32), } + impl DictArg { + fn parts(self, compress: bool) -> (PyRef, i32) { + match self { + Self::Bare(dict) => { + let marker = if compress { UNDIGESTED } else { DIGESTED }; + (dict, marker) + } + Self::Marked(dict, marker) => (dict, marker), + } + } + } + fn parse_dict_arg(obj: &PyObject, vm: &VirtualMachine) -> PyResult { if obj.downcast_ref::().is_some() { let dict = obj @@ -758,217 +777,304 @@ mod _zstd { Err(zstd_dict_type_error(vm)) } - enum LoadedDict { - Dict(Box), - Prefix(Vec), + fn cdict_at_level( + dict: &ZstdDict, + level: i32, + ) -> Result<*const libzstd_rs_sys::ZSTD_CDict, CodecError> { + { + let cache = dict.cache.lock(); + if let Some(existing) = cache.cdicts.get(&level) { + return Ok(existing.0); + } + } + let bytes = dict.dict_content.as_bytes(); + let created = unsafe { ZSTD_createCDict(bytes.as_ptr().cast(), bytes.len(), level) }; + if created.is_null() { + return Err(CodecError::Static(DIGESTED_COMPRESS_FAILED)); + } + let mut cache = dict.cache.lock(); + if let Some(existing) = cache.cdicts.get(&level) { + unsafe { + ZSTD_freeCDict(created); + } + return Ok(existing.0); + } + let ptr = created as *const libzstd_rs_sys::ZSTD_CDict; + cache.cdicts.insert(level, CDict(created)); + Ok(ptr) } - fn load_dict(spec: DictArg, compress: bool, vm: &VirtualMachine) -> PyResult { - let (dict, marker) = match spec { - DictArg::Bare(dict) => { - let marker = if compress { UNDIGESTED } else { DIGESTED }; - (dict, marker) + fn ddict_of(dict: &ZstdDict) -> Result<*const libzstd_rs_sys::ZSTD_DDict, CodecError> { + { + let cache = dict.cache.lock(); + if let Some(existing) = cache.ddict.as_ref() { + return Ok(existing.0); } - DictArg::Marked(dict, marker) => (dict, marker), - }; - let content = dict.dict_content.as_bytes(); - let loaded = match marker { - PREFIX => LoadedDict::Prefix(content.to_vec()), - DIGESTED if dict.dict_id == 0 => { - LoadedDict::Dict(Box::new(Dictionary::raw(content.to_vec()))) + } + let bytes = dict.dict_content.as_bytes(); + let created = unsafe { ZSTD_createDDict(bytes.as_ptr().cast(), bytes.len()) }; + if created.is_null() { + return Err(CodecError::Static(DIGESTED_DECOMPRESS_FAILED)); + } + let mut cache = dict.cache.lock(); + if let Some(existing) = cache.ddict.as_ref() { + unsafe { + ZSTD_freeDDict(created); } - DIGESTED => match Dictionary::from_bytes(content) { - Ok(parsed) => LoadedDict::Dict(Box::new(parsed)), - Err(_) => { - let message = if compress { - DIGESTED_COMPRESS_FAILED - } else { - DIGESTED_DECOMPRESS_FAILED - }; - return Err(zstd_error(message.to_owned(), vm)); - } - }, - _ => { - let parsed = match Dictionary::from_bytes(content) { - Ok(parsed) => parsed, - Err(_) => Dictionary::raw(content.to_vec()), - }; - LoadedDict::Dict(Box::new(parsed)) + return Ok(existing.0); + } + let ptr = created as *const libzstd_rs_sys::ZSTD_DDict; + cache.ddict = Some(DDict(created)); + Ok(ptr) + } + + struct OutBuf { + data: Vec, + filled: usize, + block_size: usize, + written_in_block: usize, + blocks: usize, + max_length: Option, + } + + impl OutBuf { + fn new(block: usize, max_length: Option) -> Result { + let mut buf = Self { + data: Vec::new(), + filled: 0, + block_size: block, + written_in_block: 0, + blocks: 1, + max_length, + }; + buf.reserve()?; + Ok(buf) + } + + fn reserve(&mut self) -> Result<(), CodecError> { + let need = self.filled.saturating_add(self.block_size); + if self.data.capacity() < need { + let extra = need - self.data.capacity(); + self.data + .try_reserve(extra) + .map_err(|_| CodecError::Buffer)?; } - }; - Ok(loaded) + Ok(()) + } + + fn as_zstd(&mut self, scratch: &mut u8) -> ZSTD_outBuffer { + let dst = if self.block_size == 0 { + core::ptr::from_mut(scratch).cast() + } else { + unsafe { self.data.as_mut_ptr().add(self.filled).cast() } + }; + ZSTD_outBuffer { + dst, + size: self.block_size, + pos: self.written_in_block, + } + } + + fn sync(&mut self, pos: usize) { + let pos = pos.min(self.block_size); + self.written_in_block = pos; + let new_len = self.filled + pos; + if new_len > self.data.len() { + // `pos` bytes at the end of the current block were written by zstd. + unsafe { self.data.set_len(new_len) }; + } else if new_len < self.data.len() { + self.data.truncate(new_len); + } + } + + fn grow(&mut self) -> Result<(), CodecError> { + self.filled += self.written_in_block; + self.written_in_block = 0; + self.block_size = next_block(self.blocks, self.max_length, self.filled)?; + self.blocks += 1; + self.reserve() + } + + fn into_vec(self) -> Vec { + self.data + } } - enum AppliedDict { - None, - Dict { - dict: Box, - write_id: bool, - }, - Prefix(Vec), + fn next_block( + index: usize, + max_length: Option, + allocated: usize, + ) -> Result { + let mut block = OUTPUT_BLOCKS + .get(index) + .copied() + .unwrap_or(OUTPUT_BLOCKS[OUTPUT_BLOCKS.len() - 1]); + if let Some(max) = max_length { + let rest = max.saturating_sub(allocated); + if rest == 0 { + return Err(CodecError::Buffer); + } + block = block.min(rest); + } + let limit = isize::MAX as usize; + if allocated > limit || block > limit - allocated { + return Err(CodecError::Buffer); + } + Ok(block) } - struct Session { - comp: Box, - produced: u64, + fn decompressor_out(max_length: Option) -> Result { + let block = match max_length { + Some(max) if max < OUTPUT_BLOCKS[0] => max, + _ => OUTPUT_BLOCKS[0], + }; + OutBuf::new(block, max_length) } struct CompState { - settings: CompressSettings, - applied: AppliedDict, - pledge: Option, - session: Option, - pending: Vec, - block_max: usize, - last_mode_value: i32, + cctx: CCtx, + use_multithread: bool, + compression_level: i32, } impl CompState { - fn fail_session(&mut self) { - self.session = None; - self.pending.clear(); - self.last_mode_value = MODE_FRAME; - } - - fn finish_frame(&mut self) { - self.session = None; - self.pending.clear(); - self.pledge = None; - if matches!(self.applied, AppliedDict::Prefix(_)) { - self.applied = AppliedDict::None; + fn new() -> Result { + let cctx = unsafe { ZSTD_createCCtx() }; + if cctx.is_null() { + return Err(CodecError::Static("Unable to create ZSTD_CCtx instance.")); } + Ok(Self { + cctx: CCtx(cctx), + use_multithread: false, + compression_level: ZSTD_CLEVEL_DEFAULT, + }) } - fn frame_size_header(&self, len: usize) -> Option { - if self.settings.content_size { - Some(u64::try_from(len).unwrap_or(u64::MAX)) - } else { - None + fn reset_session(&mut self) { + unsafe { + ZSTD_CCtx_reset(self.cctx.0, ZSTD_ResetDirective::ZSTD_reset_session_only); } } - fn stream_header(&self) -> Option { - if self.settings.content_size { - self.pledge - } else { - None + fn set_level(&mut self, level: i32, vm: &VirtualMachine) -> PyResult<()> { + if !(LEVEL_MIN..=LEVEL_MAX).contains(&level) { + return Err(vm.new_value_error(format!( + "illegal compression level {level}; the valid range is [{LEVEL_MIN}, {LEVEL_MAX}]" + ))); } - } - - fn apply_dict(&self, comp: &mut Compressor) -> Result<(), Error> { - match &self.applied { - AppliedDict::Dict { dict, write_id } => { - comp.set_dictionary(dict)?; - comp.set_write_dict_id(*write_id && dict.id() != 0); - } - AppliedDict::Prefix(prefix) => { - comp.set_prefix(prefix)?; - } - AppliedDict::None => {} + self.compression_level = level; + let code = unsafe { + ZSTD_CCtx_setParameter(self.cctx.0, c_parameter(ZSTD_c_compressionLevel), level) + }; + if is_error(code) { + return Err(raise_codec(CodecError::Named(CallKind::Level, code), vm)); } Ok(()) } - fn open_session(&mut self, header: Option) -> Result<(), Error> { - let mut comp = Compressor::with_options( - CompressOptions { - level: codec_level(self.settings.level), - checksum: self.settings.checksum, + fn apply_options(&mut self, obj: &PyObject, vm: &VirtualMachine) -> PyResult<()> { + let dict = options_dict(obj, true, vm)?; + for_each_option( + dict, + true, + |key, value| { + if key == ZSTD_c_compressionLevel { + return self.set_level(value, vm); + } + if key == ZSTD_c_nbWorkers && value != 0 { + self.use_multithread = true; + } + let code = + unsafe { ZSTD_CCtx_setParameter(self.cctx.0, c_parameter(key), value) }; + if is_error(code) { + return Err(parameter_value_error(true, key, value, vm)); + } + Ok(()) }, - None, - )?; - if let Some(size) = header { - comp.set_pledged_src_size(size); - } - self.apply_dict(&mut comp)?; - self.session = Some(Session { - comp: Box::new(comp), - produced: 0, - }); - Ok(()) + vm, + ) } - // Empty input with nothing buffered must not start a frame, except for - // an explicit end, which is a one-shot empty frame. A one-shot end - // writes the input length into the header and ignores a stored pledge. - // Short CONTINUE input stays buffered: the first stream call emits a - // header, and a short write has to leave the underlying file untouched. - fn compress(&mut self, input: &[u8], mode: i32) -> Result, Error> { - if input.is_empty() - && self.session.is_none() - && self.pending.is_empty() - && mode != MODE_FRAME - { - self.last_mode_value = mode; - return Ok(Vec::new()); - } - let oneshot = self.session.is_none() && self.pending.is_empty() && mode == MODE_FRAME; - if !input.is_empty() { - self.pending.extend_from_slice(input); - } - if self.session.is_none() - && mode == MODE_CONTINUE - && self.pending.len() < self.block_max - { - self.last_mode_value = mode; - return Ok(Vec::new()); - } - if self.session.is_none() { - let header_len = self.pending.len(); - let header = if oneshot { - self.frame_size_header(header_len) - } else { - self.stream_header() - }; - if let Err(err) = self.open_session(header) { - self.fail_session(); - return Err(err); + fn load_dict(&mut self, dict: &ZstdDict, marker: i32) -> Result<(), CodecError> { + let code = match marker { + DIGESTED => { + let cdict = cdict_at_level(dict, self.compression_level)?; + unsafe { ZSTD_CCtx_refCDict(self.cctx.0, cdict) } } - } - let flush = flush_kind(mode); - let pending = core::mem::take(&mut self.pending); - let pumped = match self.session.as_mut() { - Some(session) => pump(session, &pending, flush), - None => Err(Error::Corruption), - }; - let bytes = match pumped { - Ok(bytes) => bytes, - Err(err) => { - self.fail_session(); - return Err(err); + PREFIX => { + let bytes = dict.dict_content.as_bytes(); + unsafe { ZSTD_CCtx_refPrefix(self.cctx.0, bytes.as_ptr().cast(), bytes.len()) } } - }; - if mode == MODE_FRAME { - let Some(produced) = self.session.as_ref().map(|session| session.produced) else { - self.fail_session(); - return Err(Error::Corruption); - }; - if !oneshot && self.pledge.is_some_and(|pledged| pledged != produced) { - self.fail_session(); - return Err(Error::ContentSizeMismatch); + _ => { + let bytes = dict.dict_content.as_bytes(); + unsafe { + ZSTD_CCtx_loadDictionary(self.cctx.0, bytes.as_ptr().cast(), bytes.len()) + } } - self.finish_frame(); + }; + if is_error(code) { + Err(CodecError::Named(CallKind::LoadCompress, code)) + } else { + Ok(()) } - self.last_mode_value = mode; - Ok(bytes) } } - fn pump(session: &mut Session, input: &[u8], flush: Flush) -> Result, Error> { - let added = u64::try_from(input.len()).unwrap_or(u64::MAX); - session.produced = session.produced.saturating_add(added); - let mut out = Vec::new(); - let mut buf = vec![0u8; compress_stream_out_size()]; - let mut pending = input; + fn compress_locked( + state: &mut CompState, + src: &[u8], + mode: i32, + mt_continue: bool, + ) -> Result, CodecError> { + let in_scratch = 0u8; + let mut input = in_buffer(src, &in_scratch); + let mut output = if mt_continue { + OutBuf::new(OUTPUT_BLOCKS[0], None)? + } else { + let bound = ZSTD_compressBound(src.len()); + if is_error(bound) || bound > isize::MAX as usize { + return Err(CodecError::NoMemory); + } + OutBuf::new(bound, None)? + }; + let directive = end_op(mode); loop { - let status = session.comp.stream(pending, &mut buf, flush)?; - pending = &[]; - out.extend_from_slice(&buf[..status.output_produced]); - if status.done || status.output_produced < buf.len() { + let mut out_scratch = 0u8; + let mut out = output.as_zstd(&mut out_scratch); + let before = (input.pos, out.pos); + let ret = if mt_continue { + loop { + let ret = unsafe { + ZSTD_compressStream2(state.cctx.0, &mut out, &mut input, directive) + }; + if is_error(ret) || out.pos == out.size || input.pos == input.size { + break ret; + } + } + } else { + unsafe { ZSTD_compressStream2(state.cctx.0, &mut out, &mut input, directive) } + }; + let pos = out.pos; + let filled = pos == out.size; + output.sync(pos); + if is_error(ret) { + return Err(CodecError::Named(CallKind::Compress, ret)); + } + if !mt_continue && ret == 0 { + break; + } + if filled { + output.grow()?; + continue; + } + if mt_continue { break; } + if (input.pos, pos) == before { + return Err(CodecError::Named(CallKind::Compress, ret)); + } } - Ok(out) + Ok(output.into_vec()) } #[derive(FromArgs)] @@ -985,11 +1091,12 @@ mod _zstd { #[pyclass(module = "compression.zstd", name = "ZstdCompressor", traverse)] #[derive(PyPayload)] struct ZstdCompressor { + // `state` is first so the context is freed before the dictionary it references. + #[pytraverse(skip)] + state: PyMutex, dict_ref: Option>, #[pytraverse(skip)] last_mode: AtomicI32, - #[pytraverse(skip)] - state: PyMutex, } impl fmt::Debug for ZstdCompressor { @@ -1008,43 +1115,25 @@ mod _zstd { if level.is_some() && options.is_some() { return Err(vm.new_type_error("Only one of level or options should be used.")); } - let mut settings = CompressSettings::default(); + let mut dict_ref: Option> = None; + let mut comp = CompState::new().map_err(|err| raise_codec(err, vm))?; if let Some(level) = level.as_ref() { - settings.level = parse_level(level, vm)?; + let level = parse_level(level, vm)?; + comp.set_level(level, vm)?; } if let Some(options) = options.as_ref() { - settings = parse_compress_options(options, vm)?; + comp.apply_options(options, vm)?; + } + if let Some(obj) = zstd_dict.as_ref() { + let (dict, marker) = parse_dict_arg(obj, vm)?.parts(true); + let loaded = vm.allow_threads(|| comp.load_dict(&dict, marker)); + loaded.map_err(|err| raise_codec(err, vm))?; + dict_ref = Some(dict); } - let (dict_ref, applied) = match zstd_dict.as_ref() { - Some(obj) => { - let spec = parse_dict_arg(obj, vm)?; - let dict_ref = match &spec { - DictArg::Bare(dict) | DictArg::Marked(dict, _) => dict.clone(), - }; - let loaded = load_dict(spec, true, vm)?; - let applied = match loaded { - LoadedDict::Dict(dict) => AppliedDict::Dict { - dict, - write_id: settings.dict_id, - }, - LoadedDict::Prefix(bytes) => AppliedDict::Prefix(bytes), - }; - (Some(dict_ref), applied) - } - None => (None, AppliedDict::None), - }; Ok(Self { + state: PyMutex::new(comp), dict_ref, last_mode: AtomicI32::new(MODE_FRAME), - state: PyMutex::new(CompState { - block_max: block_capacity(settings.level), - settings, - applied, - pledge: None, - session: None, - pending: Vec::new(), - last_mode_value: MODE_FRAME, - }), }) } } @@ -1063,16 +1152,51 @@ mod _zstd { mode: i32, } - #[pyclass(with(Constructor))] - impl ZstdCompressor { - fn locked_compress(zelf: &Py, input: &[u8], mode: i32) -> Result, Error> { - let mut state = zelf.state.lock(); - let result = state.compress(input, mode); - zelf.last_mode - .store(state.last_mode_value, Ordering::Relaxed); - result + fn check_compress_mode(mode: i32, vm: &VirtualMachine) -> PyResult<()> { + if matches!(mode, MODE_CONTINUE | MODE_BLOCK | MODE_FRAME) { + Ok(()) + } else { + Err(vm.new_value_error( + "mode argument wrong value, it should be one of ZstdCompressor.CONTINUE, \ + ZstdCompressor.FLUSH_BLOCK, ZstdCompressor.FLUSH_FRAME.", + )) + } + } + + fn check_flush_mode(mode: i32, vm: &VirtualMachine) -> PyResult<()> { + if matches!(mode, MODE_BLOCK | MODE_FRAME) { + Ok(()) + } else { + Err(vm.new_value_error( + "mode argument wrong value, it should be ZstdCompressor.FLUSH_FRAME or \ + ZstdCompressor.FLUSH_BLOCK.", + )) + } + } + + fn drive_compress( + zelf: &Py, + src: &[u8], + mode: i32, + allow_mt: bool, + ) -> Result, CodecError> { + let mut state = zelf.state.lock(); + let mt_continue = allow_mt && state.use_multithread && mode == MODE_CONTINUE; + match compress_locked(&mut state, src, mode, mt_continue) { + Ok(bytes) => { + zelf.last_mode.store(mode, Ordering::Release); + Ok(bytes) + } + Err(err) => { + state.reset_session(); + zelf.last_mode.store(MODE_FRAME, Ordering::Release); + Err(err) + } } + } + #[pyclass(with(Constructor))] + impl ZstdCompressor { #[pymethod] fn compress( zelf: &Py, @@ -1082,16 +1206,16 @@ mod _zstd { let CompressCallArgs { data, mode } = args; check_compress_mode(mode, vm)?; let input = data.with_ref(<[u8]>::to_vec); - vm.allow_threads(|| Self::locked_compress(zelf, &input, mode)) - .map_err(|err| compress_error(err, vm)) + vm.allow_threads(|| drive_compress(zelf, &input, mode, true)) + .map_err(|err| raise_codec(err, vm)) } #[pymethod] fn flush(zelf: &Py, args: FlushCallArgs, vm: &VirtualMachine) -> PyResult> { let FlushCallArgs { mode } = args; check_flush_mode(mode, vm)?; - vm.allow_threads(|| Self::locked_compress(zelf, &[], mode)) - .map_err(|err| compress_error(err, vm)) + vm.allow_threads(|| drive_compress(zelf, &[], mode, false)) + .map_err(|err| raise_codec(err, vm)) } #[pymethod] @@ -1100,28 +1224,31 @@ mod _zstd { size: PyObjectRef, vm: &VirtualMachine, ) -> PyResult<()> { - let pledge = parse_pledge(&size, vm)?; - let bad_mode = vm.allow_threads(|| { - let mut state = zelf.state.lock(); - if zelf.last_mode.load(Ordering::Relaxed) != MODE_FRAME { - true + let pledged = parse_pledge(&size, vm)?; + let result = vm.allow_threads(|| { + let state = zelf.state.lock(); + if zelf.last_mode.load(Ordering::Acquire) != MODE_FRAME { + return Err(None); + } + let code = unsafe { ZSTD_CCtx_setPledgedSrcSize(state.cctx.0, pledged) }; + if is_error(code) { + Err(Some(code)) } else { - state.pledge = pledge; - false + Ok(()) } }); - if bad_mode { - Err(vm.new_value_error( + match result { + Ok(()) => Ok(()), + Err(None) => Err(vm.new_value_error( "set_pledged_input_size() method must be called when last_mode == FLUSH_FRAME", - )) - } else { - Ok(()) + )), + Err(Some(code)) => Err(raise_codec(CodecError::Named(CallKind::Pledge, code), vm)), } } #[pygetset] fn last_mode(zelf: &Py) -> i32 { - zelf.last_mode.load(Ordering::Relaxed) + zelf.last_mode.load(Ordering::Acquire) } #[extend_class] @@ -1141,248 +1268,159 @@ mod _zstd { } } - fn parse_pledge(size: &PyObject, vm: &VirtualMachine) -> PyResult> { + fn parse_pledge(size: &PyObject, vm: &VirtualMachine) -> PyResult { if vm.is_none(size) { - return Ok(None); + return Ok(ZSTD_CONTENTSIZE_UNKNOWN); } if !is_pylong(size, vm) { - return Err(pledge_error(vm)); + return Err(vm.new_type_error("an integer is required")); } let int_obj = size.try_index(vm)?; match int_obj.as_bigint().to_u64() { - Some(value) if value < u64::MAX - 1 => Ok(Some(value)), - _ => Err(pledge_error(vm)), + Some(value) if value < ZSTD_CONTENTSIZE_ERROR => Ok(value), + _ => Err(vm.new_value_error(format!( + "size argument should be a positive int less than {ZSTD_CONTENTSIZE_ERROR}" + ))), } } - fn pledge_error(vm: &VirtualMachine) -> PyBaseExceptionRef { - let limit = u64::MAX - 1; - vm.new_value_error(format!( - "size argument should be a positive int less than {limit}" - )) - } - - enum DecApplied { - None, - Dict(Box), - Prefix(Vec), - } - struct DecState { - window_max: u64, - applied: DecApplied, - src: Vec, - fed: usize, - dec: Option>, - frame_len: Option, - out: Vec, - out_pos: usize, - finished: bool, - probed: usize, + dctx: DCtx, + pending: Vec, + eof: bool, needs_input: bool, } impl DecState { - fn reset_session(&mut self) { - self.src.clear(); - self.fed = 0; - self.dec = None; - self.frame_len = None; - self.out.clear(); - self.out_pos = 0; - self.finished = false; - self.probed = 0; - self.needs_input = true; - } - - fn is_eof(&self) -> bool { - self.finished && self.out.is_empty() - } - - fn unused(&self) -> &[u8] { - // Bytes after this frame stay hidden until its output is drained. - if !self.is_eof() { - return &[]; - } - let start = self.frame_len.unwrap_or(self.src.len()).min(self.src.len()); - &self.src[start..] - } - - fn clear_prefix(&mut self) { - if matches!(self.applied, DecApplied::Prefix(_)) { - self.applied = DecApplied::None; + fn new() -> Result { + let dctx = unsafe { ZSTD_createDCtx() }; + if dctx.is_null() { + return Err(CodecError::Static("Unable to create ZSTD_DCtx instance.")); } + Ok(Self { + dctx: DCtx(dctx), + pending: Vec::new(), + eof: false, + needs_input: true, + }) } - fn make_decoder(&self) -> Box { - let mut dec = Decompressor::with_options(DecompressOptions { - window_max: self.window_max, - force_ignore_checksum: false, - }); - match &self.applied { - DecApplied::Dict(dict) => dec.set_dictionary(dict.as_ref().clone()), - DecApplied::Prefix(prefix) => dec.set_prefix(prefix), - DecApplied::None => {} + fn reset_session(&mut self) { + self.pending.clear(); + self.eof = false; + self.needs_input = true; + unsafe { + ZSTD_DCtx_reset(self.dctx.0, ZSTD_ResetDirective::ZSTD_reset_session_only); } - Box::new(dec) } - fn classify(&mut self) -> Result<(), Error> { - if self.dec.is_some() || self.finished || self.src.len() < 4 { - return Ok(()); - } - let Some(magic) = read_u32_at(&self.src, 0) else { - return Err(Error::Corruption); - }; - if is_skippable(magic) { - if self.src.len() < 8 { - return Ok(()); - } - let Some(body_len) = read_u32_at(&self.src, 4) else { - return Err(Error::Corruption); - }; - let Ok(body) = usize::try_from(body_len) else { - return Err(Error::Corruption); - }; - let Some(total) = 8usize.checked_add(body) else { - return Err(Error::Corruption); - }; - if self.src.len() < total { - return Ok(()); - } - self.frame_len = Some(total); - self.finished = true; - self.clear_prefix(); - return Ok(()); - } - if magic != MAGIC { - return Err(Error::BadMagic); - } - let dec = self.make_decoder(); - self.dec = Some(dec); - Ok(()) + fn publish(&self, eof: &AtomicBool, needs_input: &AtomicBool) { + eof.store(self.eof, Ordering::Release); + needs_input.store(self.needs_input, Ordering::Release); } - fn probe(&mut self) -> Result<(), Error> { - if self.frame_len.is_some() || self.probed == self.src.len() { - return Ok(()); - } - match find_frame_compressed_size(&self.src) { - Ok(0) => Err(Error::Corruption), - Ok(len) => { - self.frame_len = Some(len); - self.probed = self.src.len(); - Ok(()) - } - Err(Error::UnexpectedEof) => { - self.probed = self.src.len(); + fn apply_options(&mut self, obj: &PyObject, vm: &VirtualMachine) -> PyResult<()> { + let dict = options_dict(obj, false, vm)?; + for_each_option( + dict, + false, + |key, value| { + let code = + unsafe { ZSTD_DCtx_setParameter(self.dctx.0, d_parameter(key), value) }; + if is_error(code) { + return Err(parameter_value_error(false, key, value, vm)); + } Ok(()) - } - Err(err) => Err(err), - } + }, + vm, + ) } - fn drain(&mut self, chunk: Vec, end: bool) -> Result<(), Error> { - let mut dec = match self.dec.take() { - Some(dec) => dec, - None => return Err(Error::Corruption), - }; - let mut buf = vec![0u8; decompress_stream_out_size()]; - let mut input: &[u8] = &chunk; - let result = loop { - match dec.stream(input, &mut buf, end) { - Ok(status) => { - input = &[]; - self.out.extend_from_slice(&buf[..status.output_produced]); - if status.done { - self.finished = true; - break Ok(()); - } - if status.output_produced < buf.len() { - break Ok(()); - } + fn load_dict(&mut self, dict: &ZstdDict, marker: i32) -> Result<(), CodecError> { + let code = match marker { + UNDIGESTED => { + let bytes = dict.dict_content.as_bytes(); + unsafe { + ZSTD_DCtx_loadDictionary(self.dctx.0, bytes.as_ptr().cast(), bytes.len()) } - Err(err) => break Err(err), + } + PREFIX => { + let bytes = dict.dict_content.as_bytes(); + unsafe { ZSTD_DCtx_refPrefix(self.dctx.0, bytes.as_ptr().cast(), bytes.len()) } + } + _ => { + let ddict = ddict_of(dict)?; + unsafe { ZSTD_DCtx_refDDict(self.dctx.0, ddict) } } }; - self.dec = Some(dec); - if self.finished { - self.clear_prefix(); - } - result - } - - fn feed(&mut self) -> Result<(), Error> { - if self.dec.is_none() || self.finished { - return Ok(()); - } - let (input_end, end_flag) = match self.frame_len { - Some(len) => (len.min(self.src.len()), self.src.len() >= len), - None => (self.src.len(), false), - }; - if self.fed < input_end { - let chunk = self.src[self.fed..input_end].to_vec(); - self.fed = input_end; - self.drain(chunk, end_flag) - } else if end_flag { - self.drain(Vec::new(), true) + if is_error(code) { + Err(CodecError::Named(CallKind::LoadDecompress, code)) } else { Ok(()) } } - fn drive(&mut self) -> Result<(), Error> { - if self.finished { - return Ok(()); - } - self.classify()?; - if self.finished || self.dec.is_none() { - return Ok(()); + fn decompress( + &mut self, + data: &[u8], + max_length: Option, + ) -> Result, CodecError> { + if self.eof { + return Err(CodecError::Eof); } - self.probe()?; - self.feed() - } - - // needs_input stays true unless this call filled max_length, or the - // frame is already finished. Returning b'' with needs_input false and - // eof false makes the reader spin. - fn take(&mut self, max_length: Option) -> Vec { - let available = self.out.len().saturating_sub(self.out_pos); - let take = match max_length { - Some(max) => available.min(max), - None => available, - }; - let start = self.out_pos; - let end = start + take; - let bytes = self.out[start..end].to_vec(); - self.out_pos = end; - if self.out_pos == self.out.len() { - self.out.clear(); - self.out_pos = 0; + self.pending + .try_reserve(data.len()) + .map_err(|_| CodecError::Buffer)?; + self.pending.extend_from_slice(data); + + let in_scratch = 0u8; + let mut input = in_buffer(&self.pending, &in_scratch); + let mut output = decompressor_out(max_length)?; + loop { + let mut out_scratch = 0u8; + let mut out = output.as_zstd(&mut out_scratch); + let ret = unsafe { ZSTD_decompressStream(self.dctx.0, &mut out, &mut input) }; + let pos = out.pos; + let filled = pos == out.size; + output.sync(pos); + if is_error(ret) { + return Err(CodecError::Named(CallKind::Decompress, ret)); + } + if ret == 0 { + self.eof = true; + break; + } + if filled { + let produced = self_produced(&output); + if max_length.is_some_and(|max| produced == max) { + break; + } + output.grow()?; + continue; + } + if input.pos == input.size { + break; + } } - self.needs_input = if self.finished { - false + + let produced = output.into_vec(); + if input.pos == input.size { + self.pending.clear(); + let hit = max_length.is_some_and(|max| produced.len() == max); + self.needs_input = !(hit || self.eof); } else { - match max_length { - Some(max) => take != max, - None => true, + self.needs_input = false; + let pos = input.pos; + if pos > 0 { + self.pending.drain(..pos); } - }; - bytes - } - - fn pull(&mut self, input: &[u8], max_length: Option) -> Result, Error> { - self.src.extend_from_slice(input); - if !self.finished { - self.drive()?; } - Ok(self.take(max_length)) + Ok(produced) } } - enum DecFail { - Eof, - Codec(Error), + fn self_produced(output: &OutBuf) -> usize { + output.filled + output.written_in_block } #[derive(FromArgs)] @@ -1397,6 +1435,9 @@ mod _zstd { #[pyclass(module = "compression.zstd", name = "ZstdDecompressor", traverse)] #[derive(PyPayload)] struct ZstdDecompressor { + // `state` is first so the context is freed before the dictionary it references. + #[pytraverse(skip)] + state: PyMutex, dict_ref: Option>, #[pymember] #[pytraverse(skip)] @@ -1404,8 +1445,6 @@ mod _zstd { #[pymember] #[pytraverse(skip)] needs_input: AtomicBool, - #[pytraverse(skip)] - state: PyMutex, } impl fmt::Debug for ZstdDecompressor { @@ -1420,42 +1459,24 @@ mod _zstd { fn py_new(_cls: &Py, args: Self::Args, vm: &VirtualMachine) -> PyResult { let zstd_dict = args.zstd_dict.flatten(); let options = args.options.flatten(); - let (dict_ref, applied) = match zstd_dict.as_ref() { - Some(obj) => { - let spec = parse_dict_arg(obj, vm)?; - let dict_ref = match &spec { - DictArg::Bare(dict) | DictArg::Marked(dict, _) => dict.clone(), - }; - let loaded = load_dict(spec, false, vm)?; - let applied = match loaded { - LoadedDict::Dict(dict) => DecApplied::Dict(dict), - LoadedDict::Prefix(bytes) => DecApplied::Prefix(bytes), - }; - (Some(dict_ref), applied) - } - None => (None, DecApplied::None), - }; - let window_log_max = match options.as_ref() { - Some(options) => parse_decompress_options(options, vm)?.window_log_max, - None => DEFAULT_WINDOW_LOG_MAX, - }; + let mut dict_ref: Option> = None; + let mut dec = DecState::new().map_err(|err| raise_codec(err, vm))?; + if let Some(obj) = zstd_dict.as_ref() { + let (dict, marker) = parse_dict_arg(obj, vm)?.parts(false); + let loaded = vm.allow_threads(|| dec.load_dict(&dict, marker)); + loaded.map_err(|err| raise_codec(err, vm))?; + dict_ref = Some(dict); + } + if let Some(options) = options.as_ref() { + dec.apply_options(options, vm)?; + } + let eof = dec.eof; + let needs_input = dec.needs_input; Ok(Self { + state: PyMutex::new(dec), dict_ref, - eof: AtomicBool::new(false), - needs_input: AtomicBool::new(true), - state: PyMutex::new(DecState { - window_max: window_max_bytes(window_log_max), - applied, - src: Vec::new(), - fed: 0, - dec: None, - frame_len: None, - out: Vec::new(), - out_pos: 0, - finished: false, - probed: 0, - needs_input: true, - }), + eof: AtomicBool::new(eof), + needs_input: AtomicBool::new(needs_input), }) } } @@ -1468,44 +1489,35 @@ mod _zstd { args: DecompressorArgs, vm: &VirtualMachine, ) -> PyResult> { - if zelf.eof.load(Ordering::Relaxed) { - return Err(vm.new_eof_error("Already at the end of a Zstandard frame.")); - } let max_length = args.max_length(); let data = args.data().to_vec(); let outcome = vm.allow_threads(|| { let mut state = zelf.state.lock(); - if state.is_eof() { - return Err(DecFail::Eof); - } - match state.pull(&data, max_length) { + match state.decompress(&data, max_length) { Ok(bytes) => { - zelf.eof.store(state.is_eof(), Ordering::Relaxed); - zelf.needs_input.store(state.needs_input, Ordering::Relaxed); + state.publish(&zelf.eof, &zelf.needs_input); Ok(bytes) } + Err(CodecError::Eof) => Err(CodecError::Eof), Err(err) => { state.reset_session(); - zelf.eof.store(false, Ordering::Relaxed); - zelf.needs_input.store(true, Ordering::Relaxed); - Err(DecFail::Codec(err)) + state.publish(&zelf.eof, &zelf.needs_input); + Err(err) } } }); - match outcome { - Ok(bytes) => Ok(bytes), - Err(DecFail::Eof) => { - Err(vm.new_eof_error("Already at the end of a Zstandard frame.")) - } - Err(DecFail::Codec(err)) => Err(decompress_error(err, vm)), - } + outcome.map_err(|err| raise_codec(err, vm)) } #[pygetset] fn unused_data(zelf: &Py, vm: &VirtualMachine) -> PyBytesRef { let unused = { let state = zelf.state.lock(); - state.unused().to_vec() + if state.eof { + state.pending.clone() + } else { + Vec::new() + } }; if unused.is_empty() { vm.ctx.empty_bytes.clone() @@ -1537,19 +1549,16 @@ mod _zstd { } } - fn split_samples<'a>( - bytes: &'a [u8], - sizes: &PyTupleRef, - vm: &VirtualMachine, - ) -> PyResult> { - if u32::try_from(sizes.as_slice().len()).is_err() { + fn sample_sizes(bytes: &[u8], sizes: &PyTupleRef, vm: &VirtualMachine) -> PyResult> { + let items = sizes.as_slice(); + if u32::try_from(items.len()).is_err() { return Err( vm.new_value_error(format!("The number of samples should be <= {}.", u32::MAX)) ); } - let mut samples = Vec::with_capacity(sizes.as_slice().len()); + let mut chunks = Vec::with_capacity(items.len()); let mut offset = 0usize; - for item in sizes.as_slice() { + for item in items { let size = match sample_size(item, vm)? { SampleSize::Mismatch => return Err(vm.new_value_error(SAMPLE_MISMATCH)), SampleSize::Value(size) => size, @@ -1560,52 +1569,78 @@ mod _zstd { if next > bytes.len() { return Err(vm.new_value_error(SAMPLE_MISMATCH)); } - samples.push(&bytes[offset..next]); + chunks.push(size); offset = next; } if offset != bytes.len() { return Err(vm.new_value_error(SAMPLE_MISMATCH)); } - Ok(samples) - } - - fn dictionary_failure(finalize: bool, vm: &VirtualMachine) -> PyBaseExceptionRef { - let verb = if finalize { "finalize" } else { "train" }; - zstd_error( - format!("Unable to {verb} the Zstandard dictionary: Src size is incorrect"), - vm, - ) - } - - fn forced_dict_id(custom: &[u8]) -> u32 { - let id = read_dict_id(custom).wrapping_add(1); - if id == 0 { 1 } else { id } + Ok(chunks) } - fn train_or_finalize( + fn build_dict( custom: Option<&[u8]>, - samples_bytes: &[u8], - samples_sizes: &PyTupleRef, + samples: &[u8], + sizes: &PyTupleRef, dict_size: isize, + compression_level: i32, vm: &VirtualMachine, ) -> PyResult> { if dict_size <= 0 { return Err(vm.new_value_error(DICT_SIZE_NONPOSITIVE)); } - let samples = split_samples(samples_bytes, samples_sizes, vm)?; - let Ok(max_dict) = usize::try_from(dict_size) else { + let chunks = sample_sizes(samples, sizes, vm)?; + let Ok(cap) = usize::try_from(dict_size) else { return Err(vm.new_value_error(DICT_SIZE_NONPOSITIVE)); }; - let dict_id = custom.map(forced_dict_id); - let opts = TrainOptions { - max_dict, - dict_id, - ..TrainOptions::fastcover() - }; - match vm.allow_threads(|| train(&samples, opts)) { - Ok(bytes) => Ok(bytes), - Err(_) => Err(dictionary_failure(custom.is_some(), vm)), + let mut dst = Vec::new(); + dst.try_reserve_exact(cap) + .map_err(|_| vm.no_memory_error())?; + dst.resize(cap, 0); + let nb = u32::try_from(chunks.len()).unwrap_or(0); + let code = vm.allow_threads(|| { + let scratch = 0u8; + let (samples_ptr, _) = ffi_bytes(samples, &scratch); + let sizes_ptr = if chunks.is_empty() { + core::ptr::null() + } else { + chunks.as_ptr() + }; + unsafe { + if let Some(custom) = custom { + let (custom_ptr, custom_len) = ffi_bytes(custom, &scratch); + ZDICT_finalizeDictionary( + dst.as_mut_ptr().cast(), + cap, + custom_ptr, + custom_len, + samples_ptr, + sizes_ptr, + nb, + ZDICT_params_t { + compressionLevel: compression_level, + notificationLevel: 0, + dictID: 0, + }, + ) + } else { + ZDICT_trainFromBuffer(dst.as_mut_ptr().cast(), cap, samples_ptr, sizes_ptr, nb) + } + } + }); + if ZDICT_isError(code) != 0 { + let kind = if custom.is_some() { + CallKind::Finalize + } else { + CallKind::Train + }; + return Err(raise_codec(CodecError::Named(kind, code), vm)); + } + if code > cap { + return Err(vm.no_memory_error()); } + dst.truncate(code); + Ok(dst) } #[pyfunction] @@ -1615,11 +1650,12 @@ mod _zstd { dict_size: isize, vm: &VirtualMachine, ) -> PyResult> { - train_or_finalize( + build_dict( None, samples_bytes.as_bytes(), &samples_sizes, dict_size, + 0, vm, ) } @@ -1633,12 +1669,12 @@ mod _zstd { compression_level: i32, vm: &VirtualMachine, ) -> PyResult> { - let _ = compression_level; - train_or_finalize( + build_dict( Some(custom_dict_bytes.as_bytes()), samples_bytes.as_bytes(), &samples_sizes, dict_size, + compression_level, vm, ) } @@ -1655,36 +1691,37 @@ mod _zstd { #[pyfunction] fn get_frame_info(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { - let kind = frame_buffer.with_ref(get_frame_header); - let kind = kind.map_err(|_| frame_info_error(vm))?; - let (size, dict_id) = match kind { - FrameKind::Zstd(header) => { - let size = match header.content_size { - Some(value) => vm.ctx.new_int(value).into(), - None => vm.ctx.none(), - }; - (size, header.dict_id.unwrap_or(0)) - } - FrameKind::Skippable { .. } => (vm.ctx.none(), 0), + let bytes = frame_buffer.with_ref(<[u8]>::to_vec); + let scratch = 0u8; + let (ptr, len) = ffi_bytes(&bytes, &scratch); + let size = unsafe { ZSTD_getFrameContentSize(ptr, len) }; + if size == ZSTD_CONTENTSIZE_ERROR { + return Err(frame_info_error(vm)); + } + let dict_id = unsafe { ZSTD_getDictID_fromFrame(ptr, len) }; + let size_obj = if size == ZSTD_CONTENTSIZE_UNKNOWN { + vm.ctx.none() + } else { + vm.ctx.new_int(size).into() }; - Ok(vm.ctx.new_tuple(vec![size, vm.ctx.new_int(dict_id).into()])) + Ok(vm + .ctx + .new_tuple(vec![size_obj, vm.ctx.new_int(dict_id).into()])) } #[pyfunction] fn get_frame_size(frame_buffer: ArgBytesLike, vm: &VirtualMachine) -> PyResult { - let result = frame_buffer.with_ref(find_frame_compressed_size); - match result { - Ok(size) => Ok(size), - Err(err) => Err(zstd_error( - format!( - "Error when finding the compressed size of a Zstandard frame. \ - Ensure the frame_buffer argument starts from the beginning of a frame, \ - and its length is not less than this complete frame. \ - Zstd error message: {}.", - frame_size_reason(err) - ), + let bytes = frame_buffer.with_ref(<[u8]>::to_vec); + let scratch = 0u8; + let (ptr, len) = ffi_bytes(&bytes, &scratch); + let size = unsafe { ZSTD_findFrameCompressedSize(ptr, len) }; + if is_error(size) { + Err(raise_codec( + CodecError::Named(CallKind::FrameSize, size), vm, - )), + )) + } else { + Ok(size) } } } From 4707f1b7f67b8f1b4f3341f17df36ae0f3051431 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 16 Oct 2026 00:12:50 +0900 Subject: [PATCH 23/25] Add zstd context and dictionary names to spell-checker ignore Assisted-by: Claude Code:claude-opus-5-5 --- crates/stdlib/src/zstd.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/stdlib/src/zstd.rs b/crates/stdlib/src/zstd.rs index 69fd96ae0a4..d57b387af6d 100644 --- a/crates/stdlib/src/zstd.rs +++ b/crates/stdlib/src/zstd.rs @@ -1,7 +1,7 @@ // spell-checker:ignore zstd Zstd ZSTD Zstandard dfast btlazy btultra btopt btultra2 btlazy2 // spell-checker:ignore ldm nbworkers windowlog hashlog chainlog searchlog minmatch dictid // spell-checker:ignore checksumflag contentsizeflag dictidflag overlaplog jobsize dstream clevel -// spell-checker:ignore skippable pledged fastcover windowlogmax +// spell-checker:ignore skippable pledged fastcover windowlogmax cctx dctx cdict cdicts ddict zdict ZDICT // cspell:ignore zstd Zstd ZSTD Zstandard dfast btlazy btultra btopt btultra2 btlazy2 ldm nbworkers // cspell:ignore windowlog hashlog chainlog searchlog minmatch dictid checksumflag skippable pledged From bd61e5fd1a2ea1ee3657908787b7feffa4a82cac Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 16 Oct 2026 00:12:54 +0900 Subject: [PATCH 24/25] Exclude _zstd from wasm32-unknown-unknown libzstd-rs-sys imports libc::size_t and libc::ptrdiff_t, which libc does not provide on wasm32-unknown-unknown. Move the dependency to the target table shared with rustls and gate the module with the same cfg. WASI targets keep _zstd. Assisted-by: Claude Code:claude-opus-5-5 --- crates/stdlib/Cargo.toml | 3 ++- crates/stdlib/src/lib.rs | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index b6cd0ba4841..680c2903f12 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -71,7 +71,6 @@ constant_time_eq = { workspace = true } # compression adler32 = { workspace = true } -libzstd-rs-sys = { workspace = true } # tkinter tk-sys = { workspace = true, optional = true } @@ -95,6 +94,8 @@ rustls-pemfile = { workspace = true, optional = true } x509-parser = { workspace = true, optional = true } pem-rfc7468 = { workspace = true, features = ["alloc"], optional = true } webpki-roots = { workspace = true, optional = true } +# libzstd-rs-sys needs libc types that wasm32-unknown-unknown lacks. +libzstd-rs-sys = { workspace = true } [target.'cfg(not(any(target_os = "android", target_arch = "wasm32")))'.dependencies] libsqlite3-sys = { workspace = true, features = ["bundled"], optional = true } diff --git a/crates/stdlib/src/lib.rs b/crates/stdlib/src/lib.rs index b3ddf6eb6b0..7ccb972d03c 100644 --- a/crates/stdlib/src/lib.rs +++ b/crates/stdlib/src/lib.rs @@ -28,6 +28,7 @@ mod lzma; mod zlib; +#[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))] mod zstd; mod blake2; @@ -327,6 +328,7 @@ pub fn stdlib_module_defs(ctx: &Context) -> Vec<&'static builtins::PyModuleDef> ))] uuid::module_def(ctx), zlib::module_def(ctx), + #[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))] zstd::module_def(ctx), ] } From cb083afc6a30f8269677c8f43dfd262a829d4533 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Sat, 17 Oct 2026 01:25:15 +0900 Subject: [PATCH 25/25] Build _zstd on wasm32-unknown-unknown with a libzstd-rs-sys fork Point libzstd-rs-sys at youknowone/libzstd-rs-sys a3e65793, which builds on wasm32-unknown-unknown and disables multithreading where std::thread::spawn is unavailable. Drop the wasm32-unknown-unknown gating of the dependency and the _zstd module. Assisted-by: Claude Code:claude-opus-5-5 --- Cargo.lock | 4 ++-- Cargo.toml | 2 +- crates/stdlib/Cargo.toml | 3 +-- crates/stdlib/src/lib.rs | 2 -- 4 files changed, 4 insertions(+), 7 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 13760b77871..01137547cb4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2324,7 +2324,7 @@ dependencies = [ [[package]] name = "libzstd-rs-sys" version = "0.0.1-prerelease.2" -source = "git+https://github.com/trifectatechfoundation/libzstd-rs-sys?rev=7afaf79f14f4a5ba4722f66d83791d30ebc47e65#7afaf79f14f4a5ba4722f66d83791d30ebc47e65" +source = "git+https://github.com/youknowone/libzstd-rs-sys?rev=a3e65793c11012ad6ae09c5d9b5dd892d2952c01#a3e65793c11012ad6ae09c5d9b5dd892d2952c01" dependencies = [ "cc", "libc", @@ -4674,7 +4674,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.3", + "getrandom 0.3.4", "once_cell", "rustix", "windows-sys 0.52.0", diff --git a/Cargo.toml b/Cargo.toml index 60d1b1682c8..c0f061f3e39 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -241,7 +241,7 @@ xz-sys = "0.4.6-rc.0" libsqlite3-sys = "0.38" libz-rs-sys = "0.6" zlib-rs = { version = "=0.6.8", default-features = false, features = ["rust-allocator", "__internal-api"] } -libzstd-rs-sys = { git = "https://github.com/trifectatechfoundation/libzstd-rs-sys", rev = "7afaf79f14f4a5ba4722f66d83791d30ebc47e65" } +libzstd-rs-sys = { git = "https://github.com/youknowone/libzstd-rs-sys", rev = "a3e65793c11012ad6ae09c5d9b5dd892d2952c01" } lock_api = "0.4" log = "0.4.30" lz4_flex = "0.14" diff --git a/crates/stdlib/Cargo.toml b/crates/stdlib/Cargo.toml index 680c2903f12..b6cd0ba4841 100644 --- a/crates/stdlib/Cargo.toml +++ b/crates/stdlib/Cargo.toml @@ -71,6 +71,7 @@ constant_time_eq = { workspace = true } # compression adler32 = { workspace = true } +libzstd-rs-sys = { workspace = true } # tkinter tk-sys = { workspace = true, optional = true } @@ -94,8 +95,6 @@ rustls-pemfile = { workspace = true, optional = true } x509-parser = { workspace = true, optional = true } pem-rfc7468 = { workspace = true, features = ["alloc"], optional = true } webpki-roots = { workspace = true, optional = true } -# libzstd-rs-sys needs libc types that wasm32-unknown-unknown lacks. -libzstd-rs-sys = { workspace = true } [target.'cfg(not(any(target_os = "android", target_arch = "wasm32")))'.dependencies] libsqlite3-sys = { workspace = true, features = ["bundled"], optional = true } diff --git a/crates/stdlib/src/lib.rs b/crates/stdlib/src/lib.rs index 7ccb972d03c..b3ddf6eb6b0 100644 --- a/crates/stdlib/src/lib.rs +++ b/crates/stdlib/src/lib.rs @@ -28,7 +28,6 @@ mod lzma; mod zlib; -#[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))] mod zstd; mod blake2; @@ -328,7 +327,6 @@ pub fn stdlib_module_defs(ctx: &Context) -> Vec<&'static builtins::PyModuleDef> ))] uuid::module_def(ctx), zlib::module_def(ctx), - #[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))] zstd::module_def(ctx), ] }