feat(compress): move begin dictionary policy into Rust

Move shared compression-begin dictionary selection, CDict attach
thresholds, reset ordering, insertion dispatch, and dictionary result
publication into a Rust-owned ABI leaf. Keep CCtx reset, CDict attach,
and private dictionary insertion behind C callbacks.

Test Plan:
- cargo test --manifest-path rust/Cargo.toml --lib
- cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings
- make -B -C programs -j1 zstd
- make -C tests -j1 test-zstream ZSTREAM_TESTTIME=-T1s
- focused compress_begin unit tests
This commit is contained in:
2026-07-19 15:12:38 +02:00
parent cc8cc47da7
commit a3127511ce
2 changed files with 530 additions and 30 deletions
+400
View File
@@ -404,6 +404,172 @@ pub unsafe extern "C" fn ZSTD_rust_initCDict(
0
}
type CompressBeginResetInternalFn =
unsafe extern "C" fn(*mut c_void, *const c_void, u64, usize, c_int) -> usize;
type CompressBeginResetUsingCDictFn =
unsafe extern "C" fn(*mut c_void, *const c_void, *const c_void, u64, c_int) -> usize;
type CompressBeginInsertDictionaryFn =
unsafe extern "C" fn(*mut c_void, *const c_void, *const c_void, usize, c_int, c_int) -> usize;
/// Projection for the dictionary-selection portion of
/// `ZSTD_compressBegin_internal`.
///
/// Rust owns the CDict attach decision and dictionary-result publication. C
/// retains the context reset, CDict attach, and private insertion callbacks.
#[repr(C)]
pub struct ZSTD_rust_compressBeginState {
callback_context: *mut c_void,
params: *const c_void,
cdict: *const c_void,
cdict_content_size: *const usize,
cdict_compression_level: *const c_int,
attach_dict_pref: *const c_int,
dict: *const c_void,
dict_size: *const usize,
dict_content_type: *const c_int,
dtlm: *const c_int,
pledged_src_size: *const u64,
zbuff: *const c_int,
force_load: *const c_int,
dict_id: *mut c_uint,
dict_content_size: *mut usize,
reset_internal: Option<CompressBeginResetInternalFn>,
reset_using_cdict: Option<CompressBeginResetUsingCDictFn>,
insert_dictionary: Option<CompressBeginInsertDictionaryFn>,
}
const _: () = {
assert!(size_of::<CompressBeginResetInternalFn>() == size_of::<usize>());
assert!(size_of::<CompressBeginResetUsingCDictFn>() == size_of::<usize>());
assert!(size_of::<CompressBeginInsertDictionaryFn>() == size_of::<usize>());
assert!(offset_of!(ZSTD_rust_compressBeginState, callback_context) == 0);
assert!(offset_of!(ZSTD_rust_compressBeginState, params) == size_of::<usize>());
assert!(offset_of!(ZSTD_rust_compressBeginState, cdict) == 2 * size_of::<usize>());
assert!(offset_of!(ZSTD_rust_compressBeginState, cdict_content_size) == 3 * size_of::<usize>());
assert!(
offset_of!(ZSTD_rust_compressBeginState, cdict_compression_level) == 4 * size_of::<usize>()
);
assert!(offset_of!(ZSTD_rust_compressBeginState, attach_dict_pref) == 5 * size_of::<usize>());
assert!(offset_of!(ZSTD_rust_compressBeginState, dict) == 6 * size_of::<usize>());
assert!(offset_of!(ZSTD_rust_compressBeginState, dict_size) == 7 * size_of::<usize>());
assert!(offset_of!(ZSTD_rust_compressBeginState, dict_content_type) == size_of::<[usize; 8]>());
assert!(offset_of!(ZSTD_rust_compressBeginState, dtlm) == 9 * size_of::<usize>());
assert!(offset_of!(ZSTD_rust_compressBeginState, pledged_src_size) == 10 * size_of::<usize>());
assert!(offset_of!(ZSTD_rust_compressBeginState, zbuff) == 11 * size_of::<usize>());
assert!(offset_of!(ZSTD_rust_compressBeginState, force_load) == size_of::<[usize; 12]>());
assert!(offset_of!(ZSTD_rust_compressBeginState, dict_id) == 13 * size_of::<usize>());
assert!(
offset_of!(ZSTD_rust_compressBeginState, dict_content_size) == size_of::<[usize; 14]>()
);
assert!(offset_of!(ZSTD_rust_compressBeginState, reset_internal) == 15 * size_of::<usize>());
assert!(offset_of!(ZSTD_rust_compressBeginState, reset_using_cdict) == 16 * size_of::<usize>());
assert!(offset_of!(ZSTD_rust_compressBeginState, insert_dictionary) == 17 * size_of::<usize>());
assert!(size_of::<ZSTD_rust_compressBeginState>() == size_of::<[usize; 18]>());
};
/// Select and begin a dictionary-backed compression context.
#[no_mangle]
pub unsafe extern "C" fn ZSTD_rust_compressBegin(
state: *const ZSTD_rust_compressBeginState,
) -> usize {
if state.is_null() {
return ERROR(ZstdErrorCode::Generic);
}
let state = unsafe { &*state };
let Some(reset_internal) = state.reset_internal else {
return ERROR(ZstdErrorCode::Generic);
};
let Some(reset_using_cdict) = state.reset_using_cdict else {
return ERROR(ZstdErrorCode::Generic);
};
let Some(insert_dictionary) = state.insert_dictionary else {
return ERROR(ZstdErrorCode::Generic);
};
if state.callback_context.is_null()
|| state.params.is_null()
|| state.cdict_content_size.is_null() && !state.cdict.is_null()
|| state.cdict_compression_level.is_null() && !state.cdict.is_null()
|| state.attach_dict_pref.is_null()
|| state.dict_size.is_null()
|| state.dict_content_type.is_null()
|| state.dtlm.is_null()
|| state.pledged_src_size.is_null()
|| state.zbuff.is_null()
|| state.force_load.is_null()
|| state.dict_id.is_null()
|| state.dict_content_size.is_null()
{
return ERROR(ZstdErrorCode::Generic);
}
let cdict_present = !state.cdict.is_null();
let pledged_src_size = unsafe { *state.pledged_src_size };
let dict_content_size = if cdict_present {
unsafe { *state.cdict_content_size }
} else {
unsafe { *state.dict_size }
};
let cdict_compression_level = if cdict_present {
unsafe { *state.cdict_compression_level }
} else {
0
};
let should_attach = cdict_present
&& dict_content_size > 0
&& (pledged_src_size < CDICT_PARAMS_SRC_SIZE_CUTOFF
|| pledged_src_size
< (dict_content_size as u64).wrapping_mul(CDICT_PARAMS_DICT_SIZE_MULTIPLIER)
|| pledged_src_size == ZSTD_CONTENTSIZE_UNKNOWN
|| cdict_compression_level == 0)
&& unsafe { *state.attach_dict_pref } != unsafe { *state.force_load };
if should_attach {
return unsafe {
reset_using_cdict(
state.callback_context,
state.cdict,
state.params,
pledged_src_size,
*state.zbuff,
)
};
}
let reset_result = unsafe {
reset_internal(
state.callback_context,
state.params,
pledged_src_size,
dict_content_size,
*state.zbuff,
)
};
if ERR_isError(reset_result) {
return reset_result;
}
let dict_id = unsafe {
insert_dictionary(
state.callback_context,
state.cdict,
state.dict,
*state.dict_size,
*state.dict_content_type,
*state.dtlm,
)
};
if ERR_isError(dict_id) {
return dict_id;
}
if dict_id > c_uint::MAX as usize {
return ERROR(ZstdErrorCode::Generic);
}
unsafe {
*state.dict_id = dict_id as c_uint;
*state.dict_content_size = dict_content_size;
}
0
}
/// Scalar projections for the public CDict query helpers.
#[repr(C)]
pub struct ZSTD_rust_cdictQueryState {
@@ -1100,6 +1266,240 @@ mod tests {
assert_eq!(content_size_flag, 1);
}
struct CompressBeginProbe {
events: Vec<&'static str>,
reset_params: *const c_void,
reset_loaded_dict_size: usize,
reset_pledged_src_size: u64,
reset_zbuff: c_int,
reset_result: usize,
attach_cdict: *const c_void,
attach_params: *const c_void,
attach_pledged_src_size: u64,
attach_zbuff: c_int,
attach_result: usize,
insert_cdict: *const c_void,
insert_dict: *const c_void,
insert_size: usize,
insert_content_type: c_int,
insert_dtlm: c_int,
insert_result: usize,
}
impl Default for CompressBeginProbe {
fn default() -> Self {
Self {
events: Vec::new(),
reset_params: ptr::null(),
reset_loaded_dict_size: 0,
reset_pledged_src_size: 0,
reset_zbuff: 0,
reset_result: 0,
attach_cdict: ptr::null(),
attach_params: ptr::null(),
attach_pledged_src_size: 0,
attach_zbuff: 0,
attach_result: 0,
insert_cdict: ptr::null(),
insert_dict: ptr::null(),
insert_size: 0,
insert_content_type: 0,
insert_dtlm: 0,
insert_result: 0,
}
}
}
unsafe fn compress_begin_probe(context: *mut c_void) -> &'static mut CompressBeginProbe {
unsafe { &mut *context.cast::<CompressBeginProbe>() }
}
unsafe extern "C" fn compress_begin_reset_internal(
context: *mut c_void,
params: *const c_void,
pledged_src_size: u64,
loaded_dict_size: usize,
zbuff: c_int,
) -> usize {
let probe = unsafe { compress_begin_probe(context) };
probe.events.push("reset");
probe.reset_params = params;
probe.reset_pledged_src_size = pledged_src_size;
probe.reset_loaded_dict_size = loaded_dict_size;
probe.reset_zbuff = zbuff;
probe.reset_result
}
unsafe extern "C" fn compress_begin_reset_using_cdict(
context: *mut c_void,
cdict: *const c_void,
params: *const c_void,
pledged_src_size: u64,
zbuff: c_int,
) -> usize {
let probe = unsafe { compress_begin_probe(context) };
probe.events.push("attach");
probe.attach_cdict = cdict;
probe.attach_params = params;
probe.attach_pledged_src_size = pledged_src_size;
probe.attach_zbuff = zbuff;
probe.attach_result
}
unsafe extern "C" fn compress_begin_insert_dictionary(
context: *mut c_void,
cdict: *const c_void,
dict: *const c_void,
dict_size: usize,
dict_content_type: c_int,
dtlm: c_int,
) -> usize {
let probe = unsafe { compress_begin_probe(context) };
probe.events.push("insert");
probe.insert_cdict = cdict;
probe.insert_dict = dict;
probe.insert_size = dict_size;
probe.insert_content_type = dict_content_type;
probe.insert_dtlm = dtlm;
probe.insert_result
}
#[allow(clippy::too_many_arguments)]
fn compress_begin_test_state(
probe: &mut CompressBeginProbe,
params: *const c_void,
cdict: *const c_void,
cdict_content_size: *const usize,
cdict_compression_level: *const c_int,
attach_dict_pref: &c_int,
dict: *const c_void,
dict_size: &usize,
dict_content_type: &c_int,
dtlm: &c_int,
pledged_src_size: &u64,
zbuff: &c_int,
force_load: &c_int,
dict_id: &mut c_uint,
dict_content_size: &mut usize,
) -> ZSTD_rust_compressBeginState {
ZSTD_rust_compressBeginState {
callback_context: (probe as *mut CompressBeginProbe).cast(),
params,
cdict,
cdict_content_size,
cdict_compression_level,
attach_dict_pref,
dict,
dict_size,
dict_content_type,
dtlm,
pledged_src_size,
zbuff,
force_load,
dict_id,
dict_content_size,
reset_internal: Some(compress_begin_reset_internal),
reset_using_cdict: Some(compress_begin_reset_using_cdict),
insert_dictionary: Some(compress_begin_insert_dictionary),
}
}
#[test]
fn compress_begin_resets_then_inserts_for_loaded_dictionary() {
let dictionary = [1u8, 2, 3, 4];
let mut probe = CompressBeginProbe {
insert_result: 17,
..Default::default()
};
let params = 0x3000usize as *const c_void;
let attach_dict_pref = 0;
let dict_size = dictionary.len();
let dict_content_type = ZSTD_DCT_RAW_CONTENT;
let dtlm = 7;
let pledged_src_size = 1u64 << 20;
let zbuff = 9;
let force_load = 1;
let mut dict_id = 0;
let mut dict_content_size = 0;
let state = compress_begin_test_state(
&mut probe,
params,
ptr::null(),
ptr::null(),
ptr::null(),
&attach_dict_pref,
dictionary.as_ptr().cast(),
&dict_size,
&dict_content_type,
&dtlm,
&pledged_src_size,
&zbuff,
&force_load,
&mut dict_id,
&mut dict_content_size,
);
let result = unsafe { ZSTD_rust_compressBegin(&state) };
assert_eq!(result, 0);
assert_eq!(probe.events, ["reset", "insert"]);
assert_eq!(probe.reset_params, params);
assert_eq!(probe.reset_loaded_dict_size, dictionary.len());
assert_eq!(probe.reset_pledged_src_size, pledged_src_size);
assert_eq!(probe.reset_zbuff, zbuff);
assert!(probe.insert_cdict.is_null());
assert_eq!(probe.insert_dict, dictionary.as_ptr().cast());
assert_eq!(probe.insert_size, dictionary.len());
assert_eq!(probe.insert_content_type, dict_content_type);
assert_eq!(probe.insert_dtlm, dtlm);
assert_eq!(dict_id, probe.insert_result as c_uint);
assert_eq!(dict_content_size, dictionary.len());
}
#[test]
fn compress_begin_attaches_small_cdict_without_insertion() {
let mut probe = CompressBeginProbe::default();
let params = 0x3000usize as *const c_void;
let cdict = 0x4000usize as *const c_void;
let cdict_content_size = 16;
let cdict_compression_level = 0;
let attach_dict_pref = 0;
let dict_size = 0;
let dict_content_type = ZSTD_DCT_AUTO;
let dtlm = 7;
let pledged_src_size = ZSTD_CONTENTSIZE_UNKNOWN;
let zbuff = 9;
let force_load = 1;
let mut dict_id = 0;
let mut dict_content_size = 0;
let state = compress_begin_test_state(
&mut probe,
params,
cdict,
&cdict_content_size,
&cdict_compression_level,
&attach_dict_pref,
ptr::null(),
&dict_size,
&dict_content_type,
&dtlm,
&pledged_src_size,
&zbuff,
&force_load,
&mut dict_id,
&mut dict_content_size,
);
let result = unsafe { ZSTD_rust_compressBegin(&state) };
assert_eq!(result, 0);
assert_eq!(probe.events, ["attach"]);
assert_eq!(probe.attach_cdict, cdict);
assert_eq!(probe.attach_params, params);
assert_eq!(probe.attach_pledged_src_size, pledged_src_size);
assert_eq!(probe.attach_zbuff, zbuff);
}
unsafe fn cctx_policy_probe(context: *mut c_void) -> &'static mut CctxPolicyProbe {
unsafe { &mut *context.cast::<CctxPolicyProbe>() }
}