From 59b3c7021b324d9bbfc5cde44148dfe1ed8a4289 Mon Sep 17 00:00:00 2001 From: ddidderr Date: Mon, 20 Jul 2026 01:55:14 +0200 Subject: [PATCH] feat(compress): route split blocks through Rust Project frame-chunk split blocks into Rust alongside the existing internal and target paths. Rust now builds the sequence store, asks the C-owned block-split context to derive partitions, refreshes the per-block first-block state, and emits the split blocks through the Rust post-build path. Keep the C partition discovery callback and the existing callback fallback as compatibility seams. Preserve the no-compress path without requiring sequence-store cursors that it does not use. Test Plan: - ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo fmt --manifest-path rust/Cargo.toml --all -- --check - ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/Cargo.toml --lib 'zstd_compress::tests::frame_chunk' -- --nocapture - ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/Cargo.toml --lib - ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings - ulimit -v 41943040; make -j1 - ulimit -v 41943040; make -j1 -C tests test-zstream ZSTREAM_TESTTIME=-T2s - ulimit -v 41943040; make -j1 -C tests test-fuzzer FUZZERTEST=-T3s FUZZER_FLAGS=--no-big-tests --- lib/compress/zstd_compress.c | 42 ++++++++++++- rust/src/zstd_compress.rs | 111 +++++++++++++++++++++++++++++++---- 2 files changed, 140 insertions(+), 13 deletions(-) diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index 6cf2ec4e4..e88c4c0d0 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -819,12 +819,17 @@ typedef struct ZSTD_rust_compressContinueBlockState_s ZSTD_rust_compressContinueBlockState; typedef struct ZSTD_rust_targetCBlockSizeState_s ZSTD_rust_targetCBlockSizeState; +typedef struct ZSTD_rust_splitBlockState_s + ZSTD_rust_splitBlockState; typedef size_t (*ZSTD_rust_frameChunkCompress_f)(void* context, void* dst, size_t dstCapacity, const void* src, size_t srcSize, U32 lastBlock); +typedef size_t (*ZSTD_rust_frameChunkDeriveSplits_f)(void* context, + U32* partitions, + U32 nbSeq); typedef struct { void* callbackContext; void* tmpWorkspace; @@ -846,6 +851,8 @@ typedef struct { ZSTD_rust_frameChunkCompress_f compressInternal; const ZSTD_rust_compressContinueBlockState* compressInternalState; const ZSTD_rust_targetCBlockSizeState* compressTargetState; + const ZSTD_rust_splitBlockState* compressSplitState; + ZSTD_rust_frameChunkDeriveSplits_f deriveBlockSplits; } ZSTD_rust_frameChunkState; size_t ZSTD_rust_compressFrameChunk( const ZSTD_rust_frameChunkState* state, @@ -879,8 +886,12 @@ typedef char ZSTD_rust_frame_chunk_state_layout[ == 11 * sizeof(void*) + sizeof(S64) + 6 * sizeof(int) && offsetof(ZSTD_rust_frameChunkState, compressTargetState) == 12 * sizeof(void*) + sizeof(S64) + 6 * sizeof(int) + && offsetof(ZSTD_rust_frameChunkState, compressSplitState) + == 13 * sizeof(void*) + sizeof(S64) + 6 * sizeof(int) + && offsetof(ZSTD_rust_frameChunkState, deriveBlockSplits) + == 14 * sizeof(void*) + sizeof(S64) + 6 * sizeof(int) && sizeof(ZSTD_rust_frameChunkState) - == 13 * sizeof(void*) + sizeof(S64) + 6 * sizeof(int)) + == 15 * sizeof(void*) + sizeof(S64) + 6 * sizeof(int)) ? 1 : -1]; /* The high-level continue/block entry points are Rust-owned. This projection @@ -1295,7 +1306,7 @@ typedef char ZSTD_rust_target_cblock_state_layout[ /* The post-split partition loop only needs this projection of ZSTD_CCtx. * Split discovery remains in the surrounding C function; the Rust body owns * partition accounting, repcode simulation, and per-partition emission. */ -typedef struct { +struct ZSTD_rust_splitBlockState_s { const SeqStore_t* seqStore; const U32* partitions; SeqStore_t* nextSeqStore; @@ -1310,7 +1321,7 @@ typedef struct { int disableLiteralCompression; int bmi2; int isFirstBlock; -} ZSTD_rust_splitBlockState; +}; size_t ZSTD_rust_compressBlockSplit( const ZSTD_rust_splitBlockState* state, void* dst, size_t dstCapacity, @@ -5432,6 +5443,7 @@ typedef struct { ZSTD_rust_blockInternalState blockInternalState; ZSTD_rust_compressContinueBlockState blockState; ZSTD_rust_targetCBlockSizeState targetBlockState; + ZSTD_rust_splitBlockState splitBlockState; } ZSTD_rust_compressContinueContext; static int ZSTD_rust_overflowCorrect_need( @@ -5561,6 +5573,13 @@ static size_t ZSTD_rust_frameChunk_compressSplit( (ZSTD_CCtx*)context, dst, dstCapacity, src, srcSize, lastBlock); } +static size_t ZSTD_rust_frameChunk_deriveBlockSplits( + void* context, U32* partitions, U32 nbSeq) +{ + return ZSTD_deriveBlockSplits( + (ZSTD_CCtx*)context, partitions, nbSeq); +} + static void ZSTD_compressContinue_prepare( ZSTD_CCtx* cctx, size_t blockSizeMax, int checkBlockSize, ZSTD_rust_compressContinueContext* context) @@ -5611,6 +5630,8 @@ static void ZSTD_compressContinue_prepare( context->frameChunkState.compressSplit = ZSTD_rust_frameChunk_compressSplit; context->frameChunkState.compressInternal = NULL; context->frameChunkState.compressTargetState = &context->targetBlockState; + context->frameChunkState.compressSplitState = &context->splitBlockState; + context->frameChunkState.deriveBlockSplits = ZSTD_rust_frameChunk_deriveBlockSplits; ZSTD_initBuildSeqStoreState(cctx, &context->buildSeqStoreState); context->blockInternalState.seqStore = &cctx->seqStore; @@ -5639,6 +5660,21 @@ static void ZSTD_compressContinue_prepare( context->targetBlockState.windowLog = cctx->appliedParams.cParams.windowLog; context->targetBlockState.targetCBlockSize = cctx->appliedParams.targetCBlockSize; context->targetBlockState.isFirstBlock = cctx->isFirstBlock; + context->splitBlockState.seqStore = &cctx->seqStore; + context->splitBlockState.partitions = cctx->blockSplitCtx.partitions; + context->splitBlockState.nextSeqStore = &cctx->blockSplitCtx.nextSeqStore; + context->splitBlockState.currSeqStore = &cctx->blockSplitCtx.currSeqStore; + context->splitBlockState.prevCBlock = &cctx->blockState.prevCBlock; + context->splitBlockState.nextCBlock = &cctx->blockState.nextCBlock; + context->splitBlockState.tmpWorkspace = cctx->tmpWorkspace; + context->splitBlockState.tmpWkspSize = cctx->tmpWkspSize; + context->splitBlockState.seqCollector = &cctx->seqCollector; + context->splitBlockState.blockSizeMax = cctx->blockSizeMax; + context->splitBlockState.strategy = (int)cctx->appliedParams.cParams.strategy; + context->splitBlockState.disableLiteralCompression = + ZSTD_literalsCompressionIsDisabled(&cctx->appliedParams); + context->splitBlockState.bmi2 = cctx->bmi2; + context->splitBlockState.isFirstBlock = cctx->isFirstBlock; context->overflowContext.matchState = ms; context->overflowContext.workspace = &cctx->workspace; diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs index 2e2a23616..197714ee3 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -180,6 +180,7 @@ type FrameChunkPrepareOverflowFn = unsafe extern "C" fn(*mut c_void, *const c_vo type FrameChunkPrepareWindowFn = unsafe extern "C" fn(*mut c_void, *const c_void, usize, c_uint); type FrameChunkCompressFn = unsafe extern "C" fn(*mut c_void, *mut c_void, usize, *const c_void, usize, c_uint) -> usize; +type FrameChunkDeriveSplitsFn = unsafe extern "C" fn(*mut c_void, *mut u32, c_uint) -> usize; type BuildSeqStoreSkipFn = unsafe extern "C" fn(*mut c_void, usize); type BuildSeqStorePrepareFn = unsafe extern "C" fn(*mut c_void, *const c_void, usize); @@ -587,6 +588,8 @@ pub struct ZSTD_rust_frameChunkState { compress_internal: Option, compress_internal_state: *const ZSTD_rust_compressContinueBlockState, compress_target_state: *const ZSTD_rust_targetCBlockSizeState, + compress_split_state: *const ZSTD_rust_splitBlockState, + derive_block_splits: Option, } const _: () = { @@ -635,9 +638,17 @@ const _: () = { == 12 * size_of::() + size_of::() + 6 * size_of::() ); assert!( - size_of::() + offset_of!(ZSTD_rust_frameChunkState, compress_split_state) == 13 * size_of::() + size_of::() + 6 * size_of::() ); + assert!( + offset_of!(ZSTD_rust_frameChunkState, derive_block_splits) + == 14 * size_of::() + size_of::() + 6 * size_of::() + ); + assert!( + size_of::() + == 15 * size_of::() + size_of::() + 6 * size_of::() + ); }; /// Rust implementation of `ZSTD_compress_frameChunk`. @@ -772,15 +783,85 @@ unsafe fn compress_frame_chunk_body_with( } } } else if state.block_splitter_enabled != 0 { - unsafe { - (state.compress_split)( - state.callback_context, - op.cast(), - remaining_capacity, - ip.cast(), - block_size, - last_block, - ) + if !state.compress_split_state.is_null() { + let block_state = if state.compress_internal_state.is_null() { + return ERROR(ZstdErrorCode::Generic); + } else { + unsafe { &*state.compress_internal_state } + }; + let split_state = unsafe { &*state.compress_split_state }; + if block_state.build_seq_store_state.is_null() { + return ERROR(ZstdErrorCode::Generic); + } + let bss = unsafe { + ZSTD_rust_buildSeqStore( + block_state.build_seq_store_state, + ip.cast(), + block_size, + ) + }; + if ERR_isError(bss) { + return bss; + } + let num_splits = if bss == ZSTD_BSS_COMPRESS as usize { + if split_state.seq_store.is_null() || split_state.partitions.is_null() { + return ERROR(ZstdErrorCode::Generic); + } + let sequence_store = unsafe { &*split_state.seq_store }; + if sequence_store.sequencesStart.is_null() || sequence_store.sequences.is_null() + { + return ERROR(ZstdErrorCode::Generic); + } + let sequence_count = unsafe { + sequence_store + .sequences + .offset_from(sequence_store.sequencesStart) + }; + if sequence_count < 0 { + return ERROR(ZstdErrorCode::Generic); + } + let Some(derive_block_splits) = state.derive_block_splits else { + return ERROR(ZstdErrorCode::Generic); + }; + let result = unsafe { + derive_block_splits( + state.callback_context, + split_state.partitions.cast_mut(), + sequence_count as c_uint, + ) + }; + if ERR_isError(result) { + return result; + } + result + } else { + 0 + }; + let mut split_state = unsafe { ptr::read(state.compress_split_state) }; + split_state.is_first_block = unsafe { *state.is_first_block }; + unsafe { + ZSTD_rust_compressBlockSplitAfterBuild( + &split_state, + op.cast(), + remaining_capacity, + ip.cast(), + block_size, + last_block, + num_splits, + bss as c_int, + ) + } + } else { + unsafe { + (state.compress_split)( + state.callback_context, + op.cast(), + remaining_capacity, + ip.cast(), + block_size, + last_block, + ) + } } } else { let compressed_size = if !state.compress_internal_state.is_null() { @@ -10453,6 +10534,8 @@ mod tests { compress_internal: Some(compress_end_test_frame_target), compress_internal_state: ptr::null(), compress_target_state: ptr::null(), + compress_split_state: ptr::null(), + derive_block_splits: None, }); let frame_chunk_state = context .frame_chunk_state @@ -11330,6 +11413,8 @@ mod tests { compress_internal: Some(frame_chunk_test_internal), compress_internal_state: ptr::null(), compress_target_state: ptr::null(), + compress_split_state: ptr::null(), + derive_block_splits: None, } } @@ -11671,6 +11756,8 @@ mod tests { compress_internal: Some(compress_continue_test_frame), compress_internal_state: ptr::null(), compress_target_state: ptr::null(), + compress_split_state: ptr::null(), + derive_block_splits: None, }); let frame_chunk_state = context .frame_chunk_state @@ -12158,6 +12245,8 @@ mod tests { compress_internal: Some(compress_stream_test_block), compress_internal_state: ptr::null(), compress_target_state: ptr::null(), + compress_split_state: ptr::null(), + derive_block_splits: None, }); let continue_frame_chunk_state = context .continue_frame_chunk_state @@ -12184,6 +12273,8 @@ mod tests { compress_internal: Some(compress_stream_test_end), compress_internal_state: ptr::null(), compress_target_state: ptr::null(), + compress_split_state: ptr::null(), + derive_block_splits: None, }); let end_frame_chunk_state = context .end_frame_chunk_state