diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index 73d4fa0ac..81d8ae330 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -2639,7 +2639,8 @@ typedef char ZSTD_rust_external_sequence_store_state_layout[ typedef void* (*ZSTD_rust_initCDictReserveContent_f)( void* context, size_t dictSize); -typedef void* (*ZSTD_rust_initCDictReserveEntropy_f)(void* context); +typedef void* (*ZSTD_rust_initCDictReserveEntropy_f)( + void* context, size_t workspaceSize); typedef size_t (*ZSTD_rust_initCDictResetMatchState_f)( void* context, const ZSTD_compressionParameters* cParams, int useRowMatchFinder); @@ -6420,17 +6421,18 @@ ZSTD_compress_insertDictionary(ZSTD_compressedBlockState_t* bs, ZSTD_loadDictionaryContent_callback); } -static void* ZSTD_rust_initCDict_reserveContent(void* context, size_t dictSize) +static void* ZSTD_rust_initCDict_reserveContent( + void* context, size_t reservedSize) { ZSTD_CDict* const cdict = (ZSTD_CDict*)context; - return ZSTD_cwksp_reserve_object( - &cdict->workspace, ZSTD_cwksp_align(dictSize, sizeof(void*))); + return ZSTD_cwksp_reserve_object(&cdict->workspace, reservedSize); } -static void* ZSTD_rust_initCDict_reserveEntropy(void* context) +static void* ZSTD_rust_initCDict_reserveEntropy( + void* context, size_t workspaceSize) { ZSTD_CDict* const cdict = (ZSTD_CDict*)context; - return ZSTD_cwksp_reserve_object(&cdict->workspace, HUF_WORKSPACE_SIZE); + return ZSTD_cwksp_reserve_object(&cdict->workspace, workspaceSize); } static size_t ZSTD_rust_initCDict_resetMatchState( diff --git a/rust/src/zstd_compress_dictionary.rs b/rust/src/zstd_compress_dictionary.rs index e4e1602cd..e71bd9951 100644 --- a/rust/src/zstd_compress_dictionary.rs +++ b/rust/src/zstd_compress_dictionary.rs @@ -47,6 +47,13 @@ const ZSTD_DCT_RAW_CONTENT: c_int = 1; const ZSTD_DCT_FULL_DICT: c_int = 2; const ZSTD_DTL_FULL: c_int = 1; const ZSTD_STATIC_WORKSPACE_ALIGNMENT: usize = 8; +const CDICT_CONTENT_ALIGNMENT: usize = size_of::<*const c_void>(); + +#[inline] +const fn dictionary_content_reservation_size(dict_size: usize) -> usize { + let mask = CDICT_CONTENT_ALIGNMENT - 1; + dict_size.wrapping_add(mask) & !mask +} /// C keeps match-state/content insertion private because it depends on the /// configuration-sensitive `ZSTD_MatchState_t`, `ldmState_t`, workspace, and @@ -1241,7 +1248,7 @@ pub unsafe extern "C" fn ZSTD_rust_compressBeginUsingDict( } type InitCDictReserveContentFn = unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void; -type InitCDictReserveEntropyFn = unsafe extern "C" fn(*mut c_void) -> *mut c_void; +type InitCDictReserveEntropyFn = unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void; type InitCDictResetMatchStateFn = unsafe extern "C" fn(*mut c_void, *const c_void, c_int) -> usize; type InitCDictInsertDictionaryFn = unsafe extern "C" fn( *mut c_void, @@ -1361,7 +1368,12 @@ pub unsafe extern "C" fn ZSTD_rust_initCDict( if dict_load_method == ZSTD_DLM_BY_REF || dict.is_null() || dict_size == 0 { unsafe { *state.dict_content = dict }; } else { - let internal_buffer = unsafe { reserve_content(state.callback_context, dict_size) }; + let internal_buffer = unsafe { + reserve_content( + state.callback_context, + dictionary_content_reservation_size(dict_size), + ) + }; if internal_buffer.is_null() { return ERROR(ZstdErrorCode::MemoryAllocation); } @@ -1378,7 +1390,7 @@ pub unsafe extern "C" fn ZSTD_rust_initCDict( *state.dict_content_type = dict_content_type; } - let entropy_workspace = unsafe { reserve_entropy(state.callback_context) }; + let entropy_workspace = unsafe { reserve_entropy(state.callback_context, HUF_WORKSPACE_SIZE) }; if entropy_workspace.is_null() { return ERROR(ZstdErrorCode::MemoryAllocation); } @@ -2626,6 +2638,7 @@ mod tests { inserted_content_type: c_int, inserted_dtlm: c_int, inserted_tfp: c_int, + reserved_entropy_size: usize, entropy_workspace: *mut c_void, reset_match_result: usize, insert_result: usize, @@ -2646,6 +2659,7 @@ mod tests { inserted_content_type: 0, inserted_dtlm: 0, inserted_tfp: 0, + reserved_entropy_size: 0, entropy_workspace: ptr::null_mut(), reset_match_result: 0, insert_result: 0, @@ -2671,9 +2685,13 @@ mod tests { } } - unsafe extern "C" fn init_cdict_reserve_entropy(context: *mut c_void) -> *mut c_void { + unsafe extern "C" fn init_cdict_reserve_entropy( + context: *mut c_void, + workspace_size: usize, + ) -> *mut c_void { let probe = unsafe { init_cdict_probe(context) }; probe.events.push("reserve"); + probe.reserved_entropy_size = workspace_size; probe.entropy_workspace } @@ -2784,6 +2802,7 @@ mod tests { assert_eq!(probe.inserted_content_type, ZSTD_DCT_RAW_CONTENT); assert_eq!(probe.inserted_dtlm, ZSTD_DTL_FULL); assert_eq!(probe.inserted_tfp, ZSTD_TFP_FOR_CDICT as c_int); + assert_eq!(probe.reserved_entropy_size, HUF_WORKSPACE_SIZE); assert_eq!(match_state_c_params, c_params); assert_eq!(dedicated_dict_search, enable_dedicated_dict_search); assert_eq!(dict_content_size, dictionary.len()); @@ -2865,13 +2884,17 @@ mod tests { probe.events, ["reserve-content", "reserve", "reset-match", "insert"] ); - assert_eq!(probe.reserved_content_size, dictionary.len()); + assert_eq!( + probe.reserved_content_size, + dictionary_content_reservation_size(dictionary.len()) + ); assert_eq!(&probe.content_storage[..dictionary.len()], &dictionary); assert_eq!(dict_content, probe.content_storage.as_ptr().cast()); assert_eq!(probe.inserted_dict, dict_content); assert_eq!(dict_content_size, dictionary.len()); assert_eq!(probe.inserted_dtlm, ZSTD_DTL_FULL); assert_eq!(probe.inserted_tfp, ZSTD_TFP_FOR_CDICT as c_int); + assert_eq!(probe.reserved_entropy_size, HUF_WORKSPACE_SIZE); } struct CompressBeginProbe { @@ -5907,6 +5930,27 @@ mod tests { assert_eq!(dict_n_count_repeat(&normalized, 2, 2), FSE_REPEAT_CHECK); } + #[test] + fn dictionary_content_reservation_matches_pointer_alignment() { + let alignment = CDICT_CONTENT_ALIGNMENT; + assert!(alignment.is_power_of_two()); + assert_eq!(dictionary_content_reservation_size(0), 0); + assert_eq!(dictionary_content_reservation_size(1), alignment); + assert_eq!( + dictionary_content_reservation_size(alignment - 1), + alignment + ); + assert_eq!(dictionary_content_reservation_size(alignment), alignment); + assert_eq!( + dictionary_content_reservation_size(alignment + 1), + alignment * 2 + ); + assert_eq!( + dictionary_content_reservation_size(usize::MAX), + usize::MAX.wrapping_add(alignment - 1) & !(alignment - 1) + ); + } + #[test] fn dictionary_content_policy_preserves_both_suffix_limits() { assert_eq!(dictionary_suffix(100, 20), (80, 20));