ddidderr/zstd-rs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor(compress): move estimator policy to Rust

Browse Source 
The public CCtx and CStream size-estimation APIs still selected row-matchfinder modes and walked compression levels in C, leaving a policy-level implementation behind the Rust workspace-sizing bridge. That also made the C layer responsible for the subtle raw-size_t MAX behavior used when an estimator returns an encoded error.

Move the row-mode selection and shared monotonic memory-budget policy into Rust. Keep the C callbacks responsible for constructing private ZSTD_CCtx_params values and for the distinct CCtx/CStream sizing formulas, so no private parameter layout crosses the ABI. Preserve the signed MIN(compressionLevel, 1) start and raw maxima exactly.

Test Plan:

- cargo +nightly fmt --manifest-path rust/Cargo.toml --all -- --check

- git diff --check

- ulimit -v 41943040; CARGO_BUILD_JOBS=1 make -j1
This commit is contained in:
ddidderr
2026-07-21 11:11:46 +02:00
parent ee1429cf79
commit 1a3971ff10
2 changed files with 301 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lib/compress/zstd_compress.c
+92 -41
View File
@@ -2179,6 +2179,40 @@ size_t ZSTD_rust_estimateCCtxWorkspaceSize(
size_t ZSTD_rust_planCCtxReset(const ZSTD_rustCCtxResetState* state);
size_t ZSTD_rust_maxEstimateCCtxSize(size_t estimate0, size_t estimate1,
size_t estimate2, size_t estimate3);
typedef size_t (*ZSTD_rust_estimateCSizeUsingCParams_f)(
void* context, int useRowMatchFinder);
typedef struct {
void* callbackContext;
int strategy;
ZSTD_rust_estimateCSizeUsingCParams_f estimate;
} ZSTD_rust_estimateCSizeUsingCParamsState;
size_t ZSTD_rust_estimateCSizeUsingCParams(
const ZSTD_rust_estimateCSizeUsingCParamsState* state);
typedef char ZSTD_rust_estimate_csize_using_cparams_state_layout[
(offsetof(ZSTD_rust_estimateCSizeUsingCParamsState, callbackContext) == 0
&& offsetof(ZSTD_rust_estimateCSizeUsingCParamsState, strategy)
== sizeof(void*)
&& offsetof(ZSTD_rust_estimateCSizeUsingCParamsState, estimate)
== 2 * sizeof(void*)
&& sizeof(ZSTD_rust_estimateCSizeUsingCParamsState)
== 3 * sizeof(void*))
? 1 : -1];
typedef size_t (*ZSTD_rust_estimateCSize_f)(void* context, int compressionLevel);
typedef struct {
void* callbackContext;
int compressionLevel;
ZSTD_rust_estimateCSize_f estimate;
} ZSTD_rust_estimateCSizeState;
size_t ZSTD_rust_estimateCSizeMonotonic(
const ZSTD_rust_estimateCSizeState* state);
typedef char ZSTD_rust_estimate_csize_state_layout[
(offsetof(ZSTD_rust_estimateCSizeState, callbackContext) == 0
&& offsetof(ZSTD_rust_estimateCSizeState, compressionLevel)
== sizeof(void*)
&& offsetof(ZSTD_rust_estimateCSizeState, estimate)
== 2 * sizeof(void*)
&& sizeof(ZSTD_rust_estimateCSizeState) == 3 * sizeof(void*))
? 1 : -1];
typedef ZSTD_compressionParameters (*ZSTD_rust_estimateCCtxSizeGetCParams_f)(
void* context, int compressionLevel, U64 srcSizeHint);
typedef size_t (*ZSTD_rust_estimateCCtxSizeEstimate_f)(
@@ -4391,21 +4425,24 @@ size_t ZSTD_estimateCCtxSize_usingCCtxParams(const ZSTD_CCtx_params* params)
&cParams, &params->ldmParams, 1, useRowMatchFinder, 0, 0, ZSTD_CONTENTSIZE_UNKNOWN, ZSTD_hasExtSeqProd(params), params->maxBlockSize);
}
static size_t ZSTD_rust_estimateCCtxSize_usingCParams_estimate(
void* context, int useRowMatchFinder)
{
ZSTD_compressionParameters const* const cParams =
(const ZSTD_compressionParameters*)context;
ZSTD_CCtx_params initialParams = ZSTD_makeCCtxParamsFromCParams(*cParams);
initialParams.useRowMatchFinder = (ZSTD_ParamSwitch_e)useRowMatchFinder;
return ZSTD_estimateCCtxSize_usingCCtxParams(&initialParams);
}
size_t ZSTD_estimateCCtxSize_usingCParams(ZSTD_compressionParameters cParams)
{
ZSTD_CCtx_params initialParams = ZSTD_makeCCtxParamsFromCParams(cParams);
if (ZSTD_rowMatchFinderSupported((int)cParams.strategy)) {
/* Pick bigger of not using and using row-based matchfinder for greedy and lazy strategies */
size_t noRowCCtxSize;
size_t rowCCtxSize;
initialParams.useRowMatchFinder = ZSTD_ps_disable;
noRowCCtxSize = ZSTD_estimateCCtxSize_usingCCtxParams(&initialParams);
initialParams.useRowMatchFinder = ZSTD_ps_enable;
rowCCtxSize = ZSTD_estimateCCtxSize_usingCCtxParams(&initialParams);
return MAX(noRowCCtxSize, rowCCtxSize);
} else {
return ZSTD_estimateCCtxSize_usingCCtxParams(&initialParams);
}
ZSTD_rust_estimateCSizeUsingCParamsState const state = {
&cParams,
(int)cParams.strategy,
ZSTD_rust_estimateCCtxSize_usingCParams_estimate
};
return ZSTD_rust_estimateCSizeUsingCParams(&state);
}
static ZSTD_compressionParameters ZSTD_rust_estimateCCtxSize_getCParams(
@@ -4434,16 +4471,21 @@ static size_t ZSTD_estimateCCtxSize_internal(int compressionLevel)
return ZSTD_rust_estimateCCtxSizeInternal(&state);
}
static size_t ZSTD_rust_estimateCCtxSize_estimateLevel(
void* context, int compressionLevel)
{
(void)context;
return ZSTD_estimateCCtxSize_internal(compressionLevel);
}
size_t ZSTD_estimateCCtxSize(int compressionLevel)
{
int level;
size_t memBudget = 0;
for (level=MIN(compressionLevel, 1); level<=compressionLevel; level++) {
/* Ensure monotonically increasing memory usage as compression level increases */
size_t const newMB = ZSTD_estimateCCtxSize_internal(level);
if (newMB > memBudget) memBudget = newMB;
}
return memBudget;
ZSTD_rust_estimateCSizeState const state = {
NULL,
compressionLevel,
ZSTD_rust_estimateCCtxSize_estimateLevel
};
return ZSTD_rust_estimateCSizeMonotonic(&state);
}
size_t ZSTD_estimateCStreamSize_usingCCtxParams(const ZSTD_CCtx_params* params)
@@ -4468,21 +4510,24 @@ size_t ZSTD_estimateCStreamSize_usingCCtxParams(const ZSTD_CCtx_params* params)
}
}
static size_t ZSTD_rust_estimateCStreamSize_usingCParams_estimate(
void* context, int useRowMatchFinder)
{
ZSTD_compressionParameters const* const cParams =
(const ZSTD_compressionParameters*)context;
ZSTD_CCtx_params initialParams = ZSTD_makeCCtxParamsFromCParams(*cParams);
initialParams.useRowMatchFinder = (ZSTD_ParamSwitch_e)useRowMatchFinder;
return ZSTD_estimateCStreamSize_usingCCtxParams(&initialParams);
}
size_t ZSTD_estimateCStreamSize_usingCParams(ZSTD_compressionParameters cParams)
{
ZSTD_CCtx_params initialParams = ZSTD_makeCCtxParamsFromCParams(cParams);
if (ZSTD_rowMatchFinderSupported((int)cParams.strategy)) {
/* Pick bigger of not using and using row-based matchfinder for greedy and lazy strategies */
size_t noRowCCtxSize;
size_t rowCCtxSize;
initialParams.useRowMatchFinder = ZSTD_ps_disable;
noRowCCtxSize = ZSTD_estimateCStreamSize_usingCCtxParams(&initialParams);
initialParams.useRowMatchFinder = ZSTD_ps_enable;
rowCCtxSize = ZSTD_estimateCStreamSize_usingCCtxParams(&initialParams);
return MAX(noRowCCtxSize, rowCCtxSize);
} else {
return ZSTD_estimateCStreamSize_usingCCtxParams(&initialParams);
}
ZSTD_rust_estimateCSizeUsingCParamsState const state = {
&cParams,
(int)cParams.strategy,
ZSTD_rust_estimateCStreamSize_usingCParams_estimate
};
return ZSTD_rust_estimateCSizeUsingCParams(&state);
}
static size_t ZSTD_estimateCStreamSize_internal(int compressionLevel)
@@ -4491,15 +4536,21 @@ static size_t ZSTD_estimateCStreamSize_internal(int compressionLevel)
return ZSTD_estimateCStreamSize_usingCParams(cParams);
}
static size_t ZSTD_rust_estimateCStreamSize_estimateLevel(
void* context, int compressionLevel)
{
(void)context;
return ZSTD_estimateCStreamSize_internal(compressionLevel);
}
size_t ZSTD_estimateCStreamSize(int compressionLevel)
{
int level;
size_t memBudget = 0;
for (level=MIN(compressionLevel, 1); level<=compressionLevel; level++) {
size_t const newMB = ZSTD_estimateCStreamSize_internal(level);
if (newMB > memBudget) memBudget = newMB;
}
return memBudget;
ZSTD_rust_estimateCSizeState const state = {
NULL,
compressionLevel,
ZSTD_rust_estimateCStreamSize_estimateLevel
};
return ZSTD_rust_estimateCSizeMonotonic(&state);
}
/* ZSTD_getFrameProgression():
rust/src/zstd_compress.rs
+209 -6
View File
@@ -30,12 +30,13 @@ use crate::zstd_compress_frame::{
use crate::zstd_compress_literals::min_gain;
use crate::zstd_compress_params::{
ZSTD_compressionParameters, ZSTD_frameParameters, ZSTD_parameters, ZSTD_resolveMaxBlockSize,
ZSTD_rustMatchStateSizing, ZSTD_rust_params_adjustCParams, ZSTD_rust_params_allocateChainTable,
ZSTD_rust_params_checkCParams, ZSTD_rust_params_defaultCLevel,
ZSTD_rust_params_estimateMatchStateSize, ZSTD_rust_params_getParamsInternal,
ZSTD_rust_params_maxNbSeq, ZSTD_rust_params_rowMatchFinderUsed,
ZSTD_rust_params_selectBlockCompressor, ZSTD_rust_params_selectCParams,
ZSTD_RUST_CPM_NO_ATTACH_DICT, ZSTD_RUST_PS_AUTO, ZSTD_RUST_PS_DISABLE, ZSTD_RUST_PS_ENABLE,
ZSTD_rowMatchFinderSupported, ZSTD_rustMatchStateSizing, ZSTD_rust_params_adjustCParams,
ZSTD_rust_params_allocateChainTable, ZSTD_rust_params_checkCParams,
ZSTD_rust_params_defaultCLevel, ZSTD_rust_params_estimateMatchStateSize,
ZSTD_rust_params_getParamsInternal, ZSTD_rust_params_maxNbSeq,
ZSTD_rust_params_rowMatchFinderUsed, ZSTD_rust_params_selectBlockCompressor,
ZSTD_rust_params_selectCParams, ZSTD_RUST_CPM_NO_ATTACH_DICT, ZSTD_RUST_PS_AUTO,
ZSTD_RUST_PS_DISABLE, ZSTD_RUST_PS_ENABLE,
};
use crate::zstd_compress_params_api::{
ZSTD_CCtxParams_setParameter, ZSTD_CCtx_params, ZSTD_customMem, ZSTD_rust_isUpdateAuthorized,
@@ -9329,6 +9330,93 @@ pub extern "C" fn ZSTD_rust_maxEstimateCCtxSize(
max_estimate_cctx_size(estimate0, estimate1, estimate2, estimate3)
}
type EstimateCSizeUsingCParamsFn = unsafe extern "C" fn(*mut c_void, c_int) -> usize;
/// C keeps private `ZSTD_CCtx_params` construction and the formula-specific
/// estimator behind this scalar callback. Rust owns only the row-mode policy.
#[repr(C)]
pub struct ZSTD_rust_estimateCSizeUsingCParamsState {
callback_context: *mut c_void,
strategy: c_int,
estimate: EstimateCSizeUsingCParamsFn,
}
const _: () = {
assert!(offset_of!(ZSTD_rust_estimateCSizeUsingCParamsState, callback_context) == 0);
assert!(offset_of!(ZSTD_rust_estimateCSizeUsingCParamsState, strategy) == size_of::<usize>());
assert!(
offset_of!(ZSTD_rust_estimateCSizeUsingCParamsState, estimate) == 2 * size_of::<usize>()
);
assert!(size_of::<ZSTD_rust_estimateCSizeUsingCParamsState>() == 3 * size_of::<usize>());
};
#[inline]
unsafe fn estimate_csize_using_cparams(state: &ZSTD_rust_estimateCSizeUsingCParamsState) -> usize {
if ZSTD_rowMatchFinderSupported(state.strategy) != 0 {
let no_row = unsafe { (state.estimate)(state.callback_context, ZSTD_RUST_PS_DISABLE) };
let row = unsafe { (state.estimate)(state.callback_context, ZSTD_RUST_PS_ENABLE) };
no_row.max(row)
} else {
unsafe { (state.estimate)(state.callback_context, ZSTD_RUST_PS_AUTO) }
}
}
/// Evaluate a CParams estimate while retaining the raw `size_t` max policy.
#[no_mangle]
pub unsafe extern "C" fn ZSTD_rust_estimateCSizeUsingCParams(
state: *const ZSTD_rust_estimateCSizeUsingCParamsState,
) -> usize {
let Some(state) = (unsafe { state.as_ref() }) else {
return ERROR(ZstdErrorCode::Generic);
};
unsafe { estimate_csize_using_cparams(state) }
}
type EstimateCSizeFn = unsafe extern "C" fn(*mut c_void, c_int) -> usize;
/// C retains each estimator's formula; Rust owns the increasing-level budget
/// policy shared by `ZSTD_estimateCCtxSize()` and `ZSTD_estimateCStreamSize()`.
#[repr(C)]
pub struct ZSTD_rust_estimateCSizeState {
callback_context: *mut c_void,
compression_level: c_int,
estimate: EstimateCSizeFn,
}
const _: () = {
assert!(offset_of!(ZSTD_rust_estimateCSizeState, callback_context) == 0);
assert!(offset_of!(ZSTD_rust_estimateCSizeState, compression_level) == size_of::<usize>());
assert!(offset_of!(ZSTD_rust_estimateCSizeState, estimate) == 2 * size_of::<usize>());
assert!(size_of::<ZSTD_rust_estimateCSizeState>() == 3 * size_of::<usize>());
};
#[inline]
unsafe fn estimate_csize_monotonic(state: &ZSTD_rust_estimateCSizeState) -> usize {
let mut level = state.compression_level.min(1);
let mut mem_budget = 0;
loop {
let new_memory = unsafe { (state.estimate)(state.callback_context, level) };
mem_budget = mem_budget.max(new_memory);
if level == state.compression_level {
break;
}
level += 1;
}
mem_budget
}
/// Evaluate every level from signed `MIN(compressionLevel, 1)` through the
/// requested level and return the raw `size_t` maximum.
#[no_mangle]
pub unsafe extern "C" fn ZSTD_rust_estimateCSizeMonotonic(
state: *const ZSTD_rust_estimateCSizeState,
) -> usize {
let Some(state) = (unsafe { state.as_ref() }) else {
return ERROR(ZstdErrorCode::Generic);
};
unsafe { estimate_csize_monotonic(state) }
}
#[inline]
fn reduce_table_internal(table: &mut [u32], reducer_value: u32, preserve_mark: bool) {
debug_assert_eq!(table.len() % ZSTD_ROWSIZE, 0);
@@ -10914,6 +11002,36 @@ mod tests {
parameter_tiers: Vec<u32>,
}
struct EstimateCSizeUsingCParamsTestContext {
modes: Vec<c_int>,
estimates: Vec<usize>,
}
unsafe extern "C" fn estimate_csize_using_cparams_test_estimate(
context: *mut c_void,
use_row_match_finder: c_int,
) -> usize {
let context = unsafe { &mut *(context.cast::<EstimateCSizeUsingCParamsTestContext>()) };
let index = context.modes.len();
context.modes.push(use_row_match_finder);
context.estimates[index]
}
struct EstimateCSizeTestContext {
levels: Vec<c_int>,
estimates: Vec<usize>,
}
unsafe extern "C" fn estimate_csize_test_estimate(
context: *mut c_void,
compression_level: c_int,
) -> usize {
let context = unsafe { &mut *(context.cast::<EstimateCSizeTestContext>()) };
let index = context.levels.len();
context.levels.push(compression_level);
context.estimates[index]
}
unsafe extern "C" fn estimate_cctx_size_test_get_c_params(
context: *mut c_void,
compression_level: c_int,
@@ -21676,6 +21794,91 @@ mod tests {
);
}
#[test]
fn estimate_using_cparams_tries_both_row_modes_for_supported_strategy() {
let mut context = EstimateCSizeUsingCParamsTestContext {
modes: Vec::new(),
estimates: vec![17, 42],
};
let state = ZSTD_rust_estimateCSizeUsingCParamsState {
callback_context: (&mut context as *mut EstimateCSizeUsingCParamsTestContext).cast(),
strategy: ZSTD_GREEDY,
estimate: estimate_csize_using_cparams_test_estimate,
};
assert_eq!(unsafe { ZSTD_rust_estimateCSizeUsingCParams(&state) }, 42);
assert_eq!(
context.modes,
vec![ZSTD_RUST_PS_DISABLE, ZSTD_RUST_PS_ENABLE]
);
}
#[test]
fn estimate_using_cparams_keeps_auto_mode_for_unsupported_strategy() {
let mut context = EstimateCSizeUsingCParamsTestContext {
modes: Vec::new(),
estimates: vec![73],
};
let state = ZSTD_rust_estimateCSizeUsingCParamsState {
callback_context: (&mut context as *mut EstimateCSizeUsingCParamsTestContext).cast(),
strategy: ZSTD_FAST,
estimate: estimate_csize_using_cparams_test_estimate,
};
assert_eq!(unsafe { ZSTD_rust_estimateCSizeUsingCParams(&state) }, 73);
assert_eq!(context.modes, vec![ZSTD_RUST_PS_AUTO]);
}
#[test]
fn estimate_using_cparams_preserves_raw_size_t_max_for_errors() {
let smaller_error = ERROR(ZstdErrorCode::DstSizeTooSmall);
let larger_error = ERROR(ZstdErrorCode::ParameterUnsupported);
let mut context = EstimateCSizeUsingCParamsTestContext {
modes: Vec::new(),
estimates: vec![smaller_error, larger_error],
};
let state = ZSTD_rust_estimateCSizeUsingCParamsState {
callback_context: (&mut context as *mut EstimateCSizeUsingCParamsTestContext).cast(),
strategy: ZSTD_GREEDY,
estimate: estimate_csize_using_cparams_test_estimate,
};
assert_eq!(
unsafe { ZSTD_rust_estimateCSizeUsingCParams(&state) },
larger_error
);
}
#[test]
fn estimate_csize_monotonic_uses_signed_min_start_and_raw_max() {
let larger_error = ERROR(ZstdErrorCode::ParameterUnsupported);
let mut context = EstimateCSizeTestContext {
levels: Vec::new(),
estimates: vec![11, larger_error, 23],
};
let state = ZSTD_rust_estimateCSizeState {
callback_context: (&mut context as *mut EstimateCSizeTestContext).cast(),
compression_level: 3,
estimate: estimate_csize_test_estimate,
};
assert_eq!(
unsafe { ZSTD_rust_estimateCSizeMonotonic(&state) },
larger_error
);
assert_eq!(context.levels, vec![1, 2, 3]);
context.levels.clear();
context.estimates = vec![77];
let state = ZSTD_rust_estimateCSizeState {
callback_context: (&mut context as *mut EstimateCSizeTestContext).cast(),
compression_level: -3,
estimate: estimate_csize_test_estimate,
};
assert_eq!(unsafe { ZSTD_rust_estimateCSizeMonotonic(&state) }, 77);
assert_eq!(context.levels, vec![-3]);
}
#[test]
fn estimate_cctx_size_internal_preserves_tier_order_and_max_policy() {
let mut context = EstimateCCtxSizeTestContext {