From b0e90c39e223a9c1aec29348ee99697d98c1d886 Mon Sep 17 00:00:00 2001 From: ddidderr Date: Mon, 20 Jul 2026 19:47:51 +0200 Subject: [PATCH] refactor(mt): move stream dictionary branch policy to Rust The MT stream initializer and its later dictionary update both encoded the same three-way choice in C: copy a supplied dictionary, attach a borrowed CDict, or install a raw prefix. The old code also embedded the required release/clear-before-attach ordering in each callback wrapper. Project the presence and raw-content flags into Rust, where the branch and ordering are now explicit and tested. C retains only the private CDict allocation, prefix-storage, context publication, and destruction callbacks, so the configured C layouts and allocator behavior remain unchanged. Test Plan: - `ulimit -v 41943040; cargo +nightly fmt --manifest-path rust/Cargo.toml --all -- --check` -- passed - `ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings` -- passed - `git diff --check` and `rustfmt --edition 2021 --check` -- passed before commit - Capped native and original-test verification remains pending for the complete batch. --- lib/compress/zstdmt_compress.c | 119 +++++++++---- rust/src/zstdmt_compress.rs | 308 ++++++++++++++++++++++++++++++++- 2 files changed, 390 insertions(+), 37 deletions(-) diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index 124c9ed83..c2bf72058 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -742,6 +742,27 @@ typedef void (*ZSTDMT_initSetBufferSizeFn)(void* opaque, size_t size); typedef size_t (*ZSTDMT_initResizeRoundBufferFn)(void* opaque, size_t capacity); typedef void (*ZSTDMT_initResetStreamFn)(void* opaque); typedef size_t (*ZSTDMT_initSerialResetFn)(void* opaque, size_t targetSectionSize); +typedef struct { + unsigned hasDictionary; + unsigned rawContent; +} ZSTDMT_RustDictionaryProjection; +typedef char ZSTDMT_rust_dictionary_projection_layout[ + (offsetof(ZSTDMT_RustDictionaryProjection, hasDictionary) == 0 + && offsetof(ZSTDMT_RustDictionaryProjection, rawContent) + == sizeof(unsigned) + && sizeof(ZSTDMT_RustDictionaryProjection) == 2 * sizeof(unsigned)) + ? 1 : -1]; +size_t ZSTDMT_rust_prepareCStreamDictionary( + const ZSTDMT_RustDictionaryProjection* projection, void* opaque, + ZSTDMT_initResetStreamFn releaseLocal, + ZSTDMT_initDictionaryFn createCopied, + ZSTDMT_initResetStreamFn attachBorrowed); +size_t ZSTDMT_rust_updateCStreamDictionary( + const ZSTDMT_RustDictionaryProjection* projection, void* opaque, + ZSTDMT_initResetStreamFn clearLocal, + ZSTDMT_initResetStreamFn attachBorrowed, + ZSTDMT_initResetStreamFn setRawPrefix, + ZSTDMT_initDictionaryFn createReferenced); typedef void (*ZSTDMT_serialResetVoidFn)(void* opaque); typedef void (*ZSTDMT_serialResetSetNbSeqFn)(void* opaque, size_t nbSeq); typedef int (*ZSTDMT_serialResetResizeFn)(void* opaque); @@ -2634,23 +2655,42 @@ static void ZSTDMT_initCStreamApplyParameters(void* opaque, size_t jobSize) state->mtctx->frameContentSize = state->pledgedSrcSize; } -static size_t ZSTDMT_initCStreamPrepareDictionary(void* opaque) +static void ZSTDMT_initCStreamReleaseLocalCDict(void* opaque) +{ + ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque; + ZSTD_freeCDict(state->mtctx->cdictLocal); +} + +static size_t ZSTDMT_initCStreamCreateCopiedCDict(void* opaque) { ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque; ZSTDMT_CCtx* const mtctx = state->mtctx; + mtctx->cdictLocal = ZSTD_createCDict_advanced( + state->dict, state->dictSize, ZSTD_dlm_byCopy, + state->dictContentType, state->params.cParams, mtctx->cMem); + mtctx->cdict = mtctx->cdictLocal; + return mtctx->cdictLocal == NULL ? ERROR(memory_allocation) : 0; +} - ZSTD_freeCDict(mtctx->cdictLocal); - if (state->dict) { - mtctx->cdictLocal = ZSTD_createCDict_advanced( - state->dict, state->dictSize, ZSTD_dlm_byCopy, - state->dictContentType, state->params.cParams, mtctx->cMem); - mtctx->cdict = mtctx->cdictLocal; - if (mtctx->cdictLocal == NULL) return ERROR(memory_allocation); - } else { - mtctx->cdictLocal = NULL; - mtctx->cdict = state->cdict; - } - return 0; +static void ZSTDMT_initCStreamAttachBorrowedCDict(void* opaque) +{ + ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque; + state->mtctx->cdictLocal = NULL; + state->mtctx->cdict = state->cdict; +} + +static size_t ZSTDMT_initCStreamPrepareDictionary(void* opaque) +{ + ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque; + ZSTDMT_RustDictionaryProjection const projection = { + state->dict != NULL, + state->dictContentType == ZSTD_dct_rawContent + }; + return ZSTDMT_rust_prepareCStreamDictionary( + &projection, state, + ZSTDMT_initCStreamReleaseLocalCDict, + ZSTDMT_initCStreamCreateCopiedCDict, + ZSTDMT_initCStreamAttachBorrowedCDict); } static void ZSTDMT_initCStreamSetTargetPrefixSize(void* opaque, size_t size) @@ -2753,30 +2793,47 @@ static void ZSTDMT_initCStreamResetStream(void* opaque) ZSTDMT_rust_resetStream(&resetState); } -static size_t ZSTDMT_initCStreamUpdateDictionary(void* opaque) +static void ZSTDMT_initCStreamClearDictionary(void* opaque) { ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque; ZSTDMT_CCtx* const mtctx = state->mtctx; - ZSTD_freeCDict(mtctx->cdictLocal); mtctx->cdictLocal = NULL; mtctx->cdict = NULL; - if (state->dict) { - if (state->dictContentType == ZSTD_dct_rawContent) { - mtctx->inBuff.prefix.start = (const BYTE*)state->dict; - mtctx->inBuff.prefix.size = state->dictSize; - } else { - /* note : a loadPrefix becomes an internal CDict */ - mtctx->cdictLocal = ZSTD_createCDict_advanced( - state->dict, state->dictSize, ZSTD_dlm_byRef, - state->dictContentType, state->params.cParams, mtctx->cMem); - mtctx->cdict = mtctx->cdictLocal; - if (mtctx->cdictLocal == NULL) return ERROR(memory_allocation); - } - } else { - mtctx->cdict = state->cdict; - } - return 0; +} + +static void ZSTDMT_initCStreamSetRawPrefix(void* opaque) +{ + ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque; + state->mtctx->inBuff.prefix.start = (const BYTE*)state->dict; + state->mtctx->inBuff.prefix.size = state->dictSize; +} + +static size_t ZSTDMT_initCStreamCreateReferencedCDict(void* opaque) +{ + ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque; + ZSTDMT_CCtx* const mtctx = state->mtctx; + /* note : a loadPrefix becomes an internal CDict */ + mtctx->cdictLocal = ZSTD_createCDict_advanced( + state->dict, state->dictSize, ZSTD_dlm_byRef, + state->dictContentType, state->params.cParams, mtctx->cMem); + mtctx->cdict = mtctx->cdictLocal; + return mtctx->cdictLocal == NULL ? ERROR(memory_allocation) : 0; +} + +static size_t ZSTDMT_initCStreamUpdateDictionary(void* opaque) +{ + ZSTDMT_initCStreamState* const state = (ZSTDMT_initCStreamState*)opaque; + ZSTDMT_RustDictionaryProjection const projection = { + state->dict != NULL, + state->dictContentType == ZSTD_dct_rawContent + }; + return ZSTDMT_rust_updateCStreamDictionary( + &projection, state, + ZSTDMT_initCStreamClearDictionary, + ZSTDMT_initCStreamAttachBorrowedCDict, + ZSTDMT_initCStreamSetRawPrefix, + ZSTDMT_initCStreamCreateReferencedCDict); } static size_t ZSTDMT_initCStreamSerialReset(void* opaque, size_t targetSectionSize) diff --git a/rust/src/zstdmt_compress.rs b/rust/src/zstdmt_compress.rs index 34e21483e..1845839f2 100644 --- a/rust/src/zstdmt_compress.rs +++ b/rust/src/zstdmt_compress.rs @@ -7,12 +7,13 @@ //! //! The serial LDM state, job descriptor fields, worker callback, and streaming //! state still use private C layouts. `zstdmt_compress.c` therefore keeps -//! those operations and projects only allocation/lifecycle pieces and pure -//! sizing policy into this module. The entry points below are narrow C ABIs: -//! buffers, `ZSTD_CCtx *` values, and job descriptors remain opaque to Rust, -//! while allocation, reuse, table-capacity orchestration, and sizing policy -//! are Rust-owned. C retains ownership of private descriptor fields and -//! platform synchronization. +//! those operations and projects allocation/lifecycle pieces, dictionary +//! transition policy, and pure sizing policy into this module. The entry +//! points below are narrow C ABIs: buffers, `ZSTD_CCtx *` values, and job +//! descriptors remain opaque to Rust, while allocation, reuse, table-capacity +//! orchestration, dictionary branch/order policy, and sizing policy are +//! Rust-owned. C retains ownership of private descriptor fields and platform +//! synchronization. use std::mem::{self, offset_of, size_of, MaybeUninit}; use std::os::raw::{c_int, c_uint, c_void}; @@ -1055,6 +1056,175 @@ pub type ZSTDMT_initSetBufferSizeFn = unsafe extern "C" fn(*mut c_void, usize); pub type ZSTDMT_initResizeRoundBufferFn = unsafe extern "C" fn(*mut c_void, usize) -> usize; pub type ZSTDMT_initResetStreamFn = unsafe extern "C" fn(*mut c_void); pub type ZSTDMT_initSerialResetFn = unsafe extern "C" fn(*mut c_void, usize) -> usize; + +/// Scalar dictionary classification for the MT stream initializer. C keeps +/// dictionary handles, parameter layout, prefix storage, and CDict mutation +/// behind callbacks; Rust owns the copy/borrow/raw-prefix branch selection. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct ZSTDMT_RustDictionaryProjection { + pub hasDictionary: c_uint, + pub rawContent: c_uint, +} + +const _: () = { + assert!(offset_of!(ZSTDMT_RustDictionaryProjection, hasDictionary) == 0); + assert!(offset_of!(ZSTDMT_RustDictionaryProjection, rawContent) == size_of::()); + assert!(size_of::() == 2 * size_of::()); +}; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum PrepareDictionaryAction { + CreateCopied, + AttachBorrowed, +} + +#[inline] +fn classify_prepare_dictionary( + projection: ZSTDMT_RustDictionaryProjection, +) -> PrepareDictionaryAction { + if projection.hasDictionary != 0 { + PrepareDictionaryAction::CreateCopied + } else { + PrepareDictionaryAction::AttachBorrowed + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum UpdateDictionaryAction { + AttachBorrowed, + SetRawPrefix, + CreateReferenced, +} + +#[inline] +fn classify_update_dictionary( + projection: ZSTDMT_RustDictionaryProjection, +) -> UpdateDictionaryAction { + if projection.hasDictionary == 0 { + UpdateDictionaryAction::AttachBorrowed + } else if projection.rawContent != 0 { + UpdateDictionaryAction::SetRawPrefix + } else { + UpdateDictionaryAction::CreateReferenced + } +} + +#[inline] +fn prepare_c_stream_dictionary_with( + projection: ZSTDMT_RustDictionaryProjection, + mut release_local: R, + mut create_copied: C, + mut attach_borrowed: A, +) -> usize +where + R: FnMut(), + C: FnMut() -> usize, + A: FnMut(), +{ + release_local(); + match classify_prepare_dictionary(projection) { + PrepareDictionaryAction::CreateCopied => create_copied(), + PrepareDictionaryAction::AttachBorrowed => { + attach_borrowed(); + 0 + } + } +} + +#[inline] +fn update_c_stream_dictionary_with( + projection: ZSTDMT_RustDictionaryProjection, + mut clear_local: C, + mut attach_borrowed: A, + mut set_raw_prefix: P, + mut create_referenced: R, +) -> usize +where + C: FnMut(), + A: FnMut(), + P: FnMut(), + R: FnMut() -> usize, +{ + clear_local(); + match classify_update_dictionary(projection) { + UpdateDictionaryAction::AttachBorrowed => { + attach_borrowed(); + 0 + } + UpdateDictionaryAction::SetRawPrefix => { + set_raw_prefix(); + 0 + } + UpdateDictionaryAction::CreateReferenced => create_referenced(), + } +} + +/// Apply the MT stream dictionary transition policy while C retains all +/// private CDict allocation and context publication operations. +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_prepareCStreamDictionary( + projection: *const ZSTDMT_RustDictionaryProjection, + opaque: *mut c_void, + release_local: Option, + create_copied: Option, + attach_borrowed: Option, +) -> usize { + let Some(projection) = (unsafe { projection.as_ref() }).copied() else { + return ERROR(ZstdErrorCode::Generic); + }; + let (Some(release_local), Some(create_copied), Some(attach_borrowed)) = + (release_local, create_copied, attach_borrowed) + else { + return ERROR(ZstdErrorCode::Generic); + }; + if opaque.is_null() { + return ERROR(ZstdErrorCode::Generic); + } + + prepare_c_stream_dictionary_with( + projection, + || unsafe { release_local(opaque) }, + || unsafe { create_copied(opaque) }, + || unsafe { attach_borrowed(opaque) }, + ) +} + +/// Apply the MT stream dictionary update policy while C retains private raw +/// prefix and CDict state. The clear step always precedes the selected branch. +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_updateCStreamDictionary( + projection: *const ZSTDMT_RustDictionaryProjection, + opaque: *mut c_void, + clear_local: Option, + attach_borrowed: Option, + set_raw_prefix: Option, + create_referenced: Option, +) -> usize { + let Some(projection) = (unsafe { projection.as_ref() }).copied() else { + return ERROR(ZstdErrorCode::Generic); + }; + let (Some(clear_local), Some(attach_borrowed), Some(set_raw_prefix), Some(create_referenced)) = ( + clear_local, + attach_borrowed, + set_raw_prefix, + create_referenced, + ) else { + return ERROR(ZstdErrorCode::Generic); + }; + if opaque.is_null() { + return ERROR(ZstdErrorCode::Generic); + } + + update_c_stream_dictionary_with( + projection, + || unsafe { clear_local(opaque) }, + || unsafe { attach_borrowed(opaque) }, + || unsafe { set_raw_prefix(opaque) }, + || unsafe { create_referenced(opaque) }, + ) +} + type ZSTDMT_resetStreamCallbackFn = unsafe extern "C" fn(*mut c_void); pub type ZSTDMT_serialResetVoidFn = unsafe extern "C" fn(*mut c_void); pub type ZSTDMT_serialResetSetNbSeqFn = unsafe extern "C" fn(*mut c_void, usize); @@ -6667,6 +6837,132 @@ mod tests { } } + #[test] + fn dictionary_policy_classifies_copy_borrow_and_raw_prefix_branches() { + assert_eq!( + classify_prepare_dictionary(ZSTDMT_RustDictionaryProjection { + hasDictionary: 1, + rawContent: 0, + }), + PrepareDictionaryAction::CreateCopied + ); + assert_eq!( + classify_prepare_dictionary(ZSTDMT_RustDictionaryProjection { + hasDictionary: 0, + rawContent: 1, + }), + PrepareDictionaryAction::AttachBorrowed + ); + + for (projection, expected) in [ + ( + ZSTDMT_RustDictionaryProjection { + hasDictionary: 0, + rawContent: 0, + }, + UpdateDictionaryAction::AttachBorrowed, + ), + ( + ZSTDMT_RustDictionaryProjection { + hasDictionary: 1, + rawContent: 1, + }, + UpdateDictionaryAction::SetRawPrefix, + ), + ( + ZSTDMT_RustDictionaryProjection { + hasDictionary: 1, + rawContent: 0, + }, + UpdateDictionaryAction::CreateReferenced, + ), + ] { + assert_eq!(classify_update_dictionary(projection), expected); + } + } + + #[test] + fn prepare_dictionary_releases_before_copy_or_borrow_action() { + let expected_error = ERROR(ZstdErrorCode::MemoryAllocation); + let events = RefCell::new(Vec::new()); + let result = prepare_c_stream_dictionary_with( + ZSTDMT_RustDictionaryProjection { + hasDictionary: 1, + rawContent: 0, + }, + || events.borrow_mut().push("release"), + || { + events.borrow_mut().push("copy"); + expected_error + }, + || panic!("copy branch must not attach borrowed dictionary"), + ); + assert_eq!(result, expected_error); + assert_eq!(events.into_inner(), vec!["release", "copy"]); + + let events = RefCell::new(Vec::new()); + let result = prepare_c_stream_dictionary_with( + ZSTDMT_RustDictionaryProjection { + hasDictionary: 0, + rawContent: 0, + }, + || events.borrow_mut().push("release"), + || panic!("borrow branch must not create a copied dictionary"), + || events.borrow_mut().push("borrow"), + ); + assert_eq!(result, 0); + assert_eq!(events.into_inner(), vec!["release", "borrow"]); + } + + #[test] + fn update_dictionary_clears_before_borrow_raw_or_reference_action() { + let events = RefCell::new(Vec::new()); + let result = update_c_stream_dictionary_with( + ZSTDMT_RustDictionaryProjection { + hasDictionary: 0, + rawContent: 0, + }, + || events.borrow_mut().push("clear"), + || events.borrow_mut().push("borrow"), + || panic!("borrow branch must not set a raw prefix"), + || panic!("borrow branch must not create a referenced dictionary"), + ); + assert_eq!(result, 0); + assert_eq!(events.into_inner(), vec!["clear", "borrow"]); + + let events = RefCell::new(Vec::new()); + let result = update_c_stream_dictionary_with( + ZSTDMT_RustDictionaryProjection { + hasDictionary: 1, + rawContent: 1, + }, + || events.borrow_mut().push("clear"), + || panic!("raw-prefix branch must not attach a borrowed CDict"), + || events.borrow_mut().push("raw"), + || panic!("raw-prefix branch must not create a referenced dictionary"), + ); + assert_eq!(result, 0); + assert_eq!(events.into_inner(), vec!["clear", "raw"]); + + let expected_error = ERROR(ZstdErrorCode::MemoryAllocation); + let events = RefCell::new(Vec::new()); + let result = update_c_stream_dictionary_with( + ZSTDMT_RustDictionaryProjection { + hasDictionary: 1, + rawContent: 0, + }, + || events.borrow_mut().push("clear"), + || panic!("reference branch must not attach a borrowed CDict"), + || panic!("reference branch must not set a raw prefix"), + || { + events.borrow_mut().push("reference"); + expected_error + }, + ); + assert_eq!(result, expected_error); + assert_eq!(events.into_inner(), vec!["clear", "reference"]); + } + #[test] fn init_c_stream_preserves_success_order_and_normalization() { let projection = init_projection();