feat(compress): move frame-block preparation order into Rust

Move the frame-chunk match-state preparation policy into Rust while keeping
window and workspace mutations behind narrow C callbacks. Rust now guarantees
the original overflow-correction, dictionary-validity, maximum-distance, and
nextToUpdate-clamp order; the C side retains the private ZSTD_MatchState_t
layout and exact primitive operations.

Test Plan:
- cargo fmt --manifest-path rust/Cargo.toml -- --check
- ulimit -v 41943040 && CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/Cargo.toml frame_chunk --lib (3 passed)
- ulimit -v 41943040 && CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/Cargo.toml --lib (691 passed)
- ulimit -v 41943040 && CARGO_BUILD_JOBS=1 cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings
- ulimit -v 41943040 && MAKEFLAGS=-j1 make -B -C programs -j1 zstd
- ulimit -v 41943040 && MAKEFLAGS=-j1 make -C tests -j1 test-zstream ZSTREAM_TESTTIME=-T1s (84 tests and both short fuzz rounds passed)

lib/compress/zstd_compress.c

@@ -673,12 +673,38 @@ typedef char ZSTD_rust_reset_match_state_layout[
                     + sizeof(void*))
         ? 1 : -1];
 
-/* The frame-chunk loop is Rust-owned.  Its callbacks keep the private
+/* The frame-chunk loop and preparation ordering are Rust-owned.  These
- * ZSTD_CCtx and match-state layout in C: Rust only drives the block loop and
+ * callbacks keep the private ZSTD_CCtx and match-state layout in C: Rust
- * passes this context back to these C-owned state-preparation/dispatch seams. */
+ * passes the context back to the narrow window/workspace operations. */
-typedef void (*ZSTD_rust_frameChunkPrepare_f)(void* context,
+typedef void (*ZSTD_rust_frameChunkPrepareOverflow_f)(void* context,
-                                               const void* src,
+                                                       const void* src,
-                                               size_t blockSize);
+                                                       size_t blockSize);
+typedef void (*ZSTD_rust_frameChunkPrepareWindow_f)(void* context,
+                                                     const void* src,
+                                                     size_t blockSize,
+                                                     U32 maxDist);
+typedef void (*ZSTD_rust_frameChunkPrepareClamp_f)(void* context);
+typedef struct {
+    void* callbackContext;
+    U32 maxDist;
+    ZSTD_rust_frameChunkPrepareOverflow_f correctOverflow;
+    ZSTD_rust_frameChunkPrepareWindow_f checkDictValidity;
+    ZSTD_rust_frameChunkPrepareWindow_f enforceMaxDist;
+    ZSTD_rust_frameChunkPrepareClamp_f clampNextToUpdate;
+} ZSTD_rust_frameChunkPrepareState;
+typedef char ZSTD_rust_frame_chunk_prepare_state_layout[
+        (offsetof(ZSTD_rust_frameChunkPrepareState, callbackContext) == 0
+         && offsetof(ZSTD_rust_frameChunkPrepareState, maxDist) == sizeof(void*)
+         && offsetof(ZSTD_rust_frameChunkPrepareState, correctOverflow)
+                == 2 * sizeof(void*)
+         && offsetof(ZSTD_rust_frameChunkPrepareState, checkDictValidity)
+                == 3 * sizeof(void*)
+         && offsetof(ZSTD_rust_frameChunkPrepareState, enforceMaxDist)
+                == 4 * sizeof(void*)
+         && offsetof(ZSTD_rust_frameChunkPrepareState, clampNextToUpdate)
+                == 5 * sizeof(void*)
+         && sizeof(ZSTD_rust_frameChunkPrepareState) == 6 * sizeof(void*))
+        ? 1 : -1];
 typedef size_t (*ZSTD_rust_frameChunkCompress_f)(void* context,
                                                  void* dst,
                                                  size_t dstCapacity,
@@ -703,7 +729,7 @@ typedef struct {
     int blockSplitterEnabled;
     int checksumFlag;
     int endingStage;
-    ZSTD_rust_frameChunkPrepare_f prepareBlock;
+    const ZSTD_rust_frameChunkPrepareState* prepareState;
     ZSTD_rust_frameChunkCompress_f compressTarget;
     ZSTD_rust_frameChunkCompress_f compressSplit;
     ZSTD_rust_frameChunkCompress_f compressInternal;
@@ -735,7 +761,7 @@ typedef char ZSTD_rust_frame_chunk_state_layout[
                 == 7 * sizeof(void*) + sizeof(S64) + 4 * sizeof(int)
          && offsetof(ZSTD_rust_frameChunkState, endingStage)
                 == 7 * sizeof(void*) + sizeof(S64) + 5 * sizeof(int)
-         && offsetof(ZSTD_rust_frameChunkState, prepareBlock)
+         && offsetof(ZSTD_rust_frameChunkState, prepareState)
                 == 7 * sizeof(void*) + sizeof(S64) + 6 * sizeof(int)
          && sizeof(ZSTD_rust_frameChunkState)
                 == 12 * sizeof(void*) + sizeof(S64) + 6 * sizeof(int))
@@ -5053,24 +5079,46 @@ static void ZSTD_overflowCorrectIfNeeded(ZSTD_MatchState_t* ms,
 
 #include "zstd_preSplit.h"
 
-static void ZSTD_rust_frameChunk_prepareBlock(void* context,
+static void ZSTD_rust_frameChunk_correctOverflow(void* context,
-                                               const void* src,
+                                                  const void* src,
-                                               size_t blockSize)
+                                                  size_t blockSize)
 {
     ZSTD_CCtx* const cctx = (ZSTD_CCtx*)context;
     ZSTD_MatchState_t* const ms = &cctx->blockState.matchState;
-    U32 const maxDist = (U32)1 << cctx->appliedParams.cParams.windowLog;
-
     ZSTD_overflowCorrectIfNeeded(
             ms, &cctx->workspace, &cctx->appliedParams,
             src, (const BYTE*)src + blockSize);
+}
+
+static void ZSTD_rust_frameChunk_checkDictValidity(void* context,
+                                                    const void* src,
+                                                    size_t blockSize,
+                                                    U32 maxDist)
+{
+    ZSTD_CCtx* const cctx = (ZSTD_CCtx*)context;
+    ZSTD_MatchState_t* const ms = &cctx->blockState.matchState;
     ZSTD_checkDictValidity(
             &ms->window, (const BYTE*)src + blockSize, maxDist,
             &ms->loadedDictEnd, &ms->dictMatchState);
+}
+
+static void ZSTD_rust_frameChunk_enforceMaxDist(void* context,
+                                                 const void* src,
+                                                 size_t blockSize,
+                                                 U32 maxDist)
+{
+    ZSTD_CCtx* const cctx = (ZSTD_CCtx*)context;
+    ZSTD_MatchState_t* const ms = &cctx->blockState.matchState;
     ZSTD_window_enforceMaxDist(
             &ms->window, src, maxDist,
             &ms->loadedDictEnd, &ms->dictMatchState);
+    (void)blockSize;
+}
 
+static void ZSTD_rust_frameChunk_clampNextToUpdate(void* context)
+{
+    ZSTD_CCtx* const cctx = (ZSTD_CCtx*)context;
+    ZSTD_MatchState_t* const ms = &cctx->blockState.matchState;
     /* Ensure hash/chain table insertion resumes no sooner than lowlimit. */
     if (ms->nextToUpdate < ms->window.lowLimit)
         ms->nextToUpdate = ms->window.lowLimit;
@@ -5121,6 +5169,13 @@ static size_t ZSTD_compress_frameChunk(ZSTD_CCtx* cctx,
                                      U32 lastFrameChunk)
 {
     ZSTD_rust_frameChunkState state;
+    ZSTD_rust_frameChunkPrepareState prepareState;
+    prepareState.callbackContext = cctx;
+    prepareState.maxDist = (U32)1 << cctx->appliedParams.cParams.windowLog;
+    prepareState.correctOverflow = ZSTD_rust_frameChunk_correctOverflow;
+    prepareState.checkDictValidity = ZSTD_rust_frameChunk_checkDictValidity;
+    prepareState.enforceMaxDist = ZSTD_rust_frameChunk_enforceMaxDist;
+    prepareState.clampNextToUpdate = ZSTD_rust_frameChunk_clampNextToUpdate;
     state.callbackContext = cctx;
     state.tmpWorkspace = cctx->tmpWorkspace;
     state.checksumState = &cctx->xxhState;
@@ -5135,7 +5190,7 @@ static size_t ZSTD_compress_frameChunk(ZSTD_CCtx* cctx,
     state.blockSplitterEnabled = ZSTD_blockSplitterEnabled(&cctx->appliedParams);
     state.checksumFlag = cctx->appliedParams.fParams.checksumFlag;
     state.endingStage = (int)ZSTDcs_ending;
-    state.prepareBlock = ZSTD_rust_frameChunk_prepareBlock;
+    state.prepareState = &prepareState;
     state.compressTarget = ZSTD_rust_frameChunk_compressTarget;
     state.compressSplit = ZSTD_rust_frameChunk_compressSplit;
     state.compressInternal = ZSTD_rust_frameChunk_compressInternal;

rust/README.md

@@ -97,7 +97,8 @@ zstd ABI:
     job-creation decisions, compression-job stage sequencing and error flow,
     pending-output decisions through scalar job projections, and frame
     progression's job-ring scan, error normalization, and active-worker
-    accounting.
+    accounting.  Rust also owns frame-block preparation ordering; C callbacks
+    retain the private match-state window and workspace operations.
 - Dictionary support
   - `zstd_ddict` owns, loads, copies, and references decode dictionaries.
 - Legacy decoding
@@ -148,7 +149,7 @@ shared- and separate-destination decompression scheduling,
 single-threaded stream initialization and the buffered/stable stream state
 machine, MT stream initialization, MT outer scheduling and flush policy, MT
 compression-job stage sequencing and error flow, MT frame-progression job
-aggregation, public sequence-API
+aggregation, frame-block preparation ordering, public sequence-API
 orchestration, sequence-store and block policy, external-producer invocation
 and success-path validation, external-sequence-store reset,
 external-sequence/literals block loop, optional-format decompression loops,

rust/src/zstd_compress.rs

@@ -170,7 +170,9 @@ const ZSTD_BSS_NO_COMPRESS: c_int = 1;
 const FSE_REPEAT_CHECK: c_int = 1;
 const FSE_REPEAT_VALID: c_int = 2;
 
-type FrameChunkPrepareFn = unsafe extern "C" fn(*mut c_void, *const c_void, usize);
+type FrameChunkPrepareOverflowFn = unsafe extern "C" fn(*mut c_void, *const c_void, usize);
+type FrameChunkPrepareWindowFn = unsafe extern "C" fn(*mut c_void, *const c_void, usize, c_uint);
+type FrameChunkPrepareClampFn = unsafe extern "C" fn(*mut c_void);
 type FrameChunkCompressFn =
     unsafe extern "C" fn(*mut c_void, *mut c_void, usize, *const c_void, usize, c_uint) -> usize;
 type FrameChunkChecksumFn = unsafe extern "C" fn(*mut c_void, *const c_void, usize);
@@ -493,6 +495,39 @@ pub unsafe extern "C" fn ZSTD_rust_referenceExternalSequences(
     }
 }
 
+/// C-owned operations used by Rust's frame-chunk block-preparation policy.
+///
+/// The callback context and private window/workspace layout remain in C; Rust
+/// controls the order in which these operations run.
+#[repr(C)]
+pub struct ZSTD_rust_frameChunkPrepareState {
+    callback_context: *mut c_void,
+    max_dist: c_uint,
+    correct_overflow: FrameChunkPrepareOverflowFn,
+    check_dict_validity: FrameChunkPrepareWindowFn,
+    enforce_max_dist: FrameChunkPrepareWindowFn,
+    clamp_next_to_update: FrameChunkPrepareClampFn,
+}
+
+const _: () = {
+    assert!(offset_of!(ZSTD_rust_frameChunkPrepareState, callback_context) == 0);
+    assert!(offset_of!(ZSTD_rust_frameChunkPrepareState, max_dist) == size_of::<usize>());
+    assert!(
+        offset_of!(ZSTD_rust_frameChunkPrepareState, correct_overflow) == 2 * size_of::<usize>()
+    );
+    assert!(
+        offset_of!(ZSTD_rust_frameChunkPrepareState, check_dict_validity) == 3 * size_of::<usize>()
+    );
+    assert!(
+        offset_of!(ZSTD_rust_frameChunkPrepareState, enforce_max_dist) == 4 * size_of::<usize>()
+    );
+    assert!(
+        offset_of!(ZSTD_rust_frameChunkPrepareState, clamp_next_to_update)
+            == 5 * size_of::<usize>()
+    );
+    assert!(size_of::<ZSTD_rust_frameChunkPrepareState>() == 6 * size_of::<usize>());
+};
+
 /// Explicit projection of the state used by `ZSTD_compress_frameChunk`.
 ///
 /// The Rust side owns the per-frame block loop and its savings/dispatch
@@ -515,7 +550,7 @@ pub struct ZSTD_rust_frameChunkState {
     block_splitter_enabled: c_int,
     checksum_flag: c_int,
     ending_stage: c_int,
-    prepare_block: FrameChunkPrepareFn,
+    prepare_state: *const ZSTD_rust_frameChunkPrepareState,
     compress_target: FrameChunkCompressFn,
     compress_split: FrameChunkCompressFn,
     compress_internal: FrameChunkCompressFn,
@@ -556,7 +591,7 @@ const _: () = {
             == 7 * size_of::<usize>() + size_of::<c_longlong>() + 5 * size_of::<c_int>()
     );
     assert!(
-        offset_of!(ZSTD_rust_frameChunkState, prepare_block)
+        offset_of!(ZSTD_rust_frameChunkState, prepare_state)
             == 7 * size_of::<usize>() + size_of::<c_longlong>() + 6 * size_of::<c_int>()
     );
     assert!(
@@ -567,9 +602,22 @@ const _: () = {
 
 /// Rust implementation of `ZSTD_compress_frameChunk`.
 ///
-/// C still prepares match-state windows and invokes the selected block body
+/// C still owns the private match-state/window operations and invokes the
-/// through callbacks. Rust owns the block-size heuristic call, output
+/// selected block body through callbacks. Rust owns preparation ordering, the
-/// framing, savings accounting, checksum sequencing, and frame-state update.
+/// block-size heuristic call, output framing, savings accounting, checksum
+/// sequencing, and frame-state update.
+#[inline]
+unsafe fn prepare_frame_chunk_block(
+    state: &ZSTD_rust_frameChunkPrepareState,
+    src: *const c_void,
+    block_size: usize,
+) {
+    unsafe { (state.correct_overflow)(state.callback_context, src, block_size) };
+    unsafe { (state.check_dict_validity)(state.callback_context, src, block_size, state.max_dist) };
+    unsafe { (state.enforce_max_dist)(state.callback_context, src, block_size, state.max_dist) };
+    unsafe { (state.clamp_next_to_update)(state.callback_context) };
+}
+
 #[allow(clippy::too_many_arguments)]
 unsafe fn compress_frame_chunk_body_with(
     state: &ZSTD_rust_frameChunkState,
@@ -579,7 +627,7 @@ unsafe fn compress_frame_chunk_body_with(
     src_size: usize,
     last_frame_chunk: c_uint,
 ) -> usize {
-    if state.is_first_block.is_null() || state.stage.is_null() {
+    if state.is_first_block.is_null() || state.stage.is_null() || state.prepare_state.is_null() {
         return ERROR(ZstdErrorCode::Generic);
     }
 
@@ -622,7 +670,7 @@ unsafe fn compress_frame_chunk_body_with(
             return ERROR(ZstdErrorCode::DstSizeTooSmall);
         }
 
-        unsafe { (state.prepare_block)(state.callback_context, ip.cast(), block_size) };
+        unsafe { prepare_frame_chunk_block(&*state.prepare_state, ip.cast(), block_size) };
 
         let c_size = if state.use_target_c_block_size != 0 {
             unsafe {
@@ -9484,6 +9532,10 @@ mod tests {
     struct FrameChunkTestContext {
         prepare_calls: usize,
         prepared_sizes: [usize; 4],
+        prepare_order: [u8; 16],
+        prepare_order_len: usize,
+        prepare_max_dist: c_uint,
+        prepare_state: Option<ZSTD_rust_frameChunkPrepareState>,
         target_calls: usize,
         split_calls: usize,
         internal_calls: usize,
@@ -9499,16 +9551,50 @@ mod tests {
         unsafe { &mut *context.cast::<FrameChunkTestContext>() }
     }
 
-    unsafe extern "C" fn frame_chunk_test_prepare(
+    unsafe fn frame_chunk_test_record_prepare(context: *mut c_void, step: u8, block_size: usize) {
+        let context = unsafe { frame_chunk_test_context(context) };
+        if context.prepare_order_len < context.prepare_order.len() {
+            context.prepare_order[context.prepare_order_len] = step;
+        }
+        context.prepare_order_len += 1;
+        if step == 1 {
+            if context.prepare_calls < context.prepared_sizes.len() {
+                context.prepared_sizes[context.prepare_calls] = block_size;
+            }
+            context.prepare_calls += 1;
+        }
+    }
+
+    unsafe extern "C" fn frame_chunk_test_correct_overflow(
         context: *mut c_void,
         _src: *const c_void,
         block_size: usize,
     ) {
-        let context = unsafe { frame_chunk_test_context(context) };
+        unsafe { frame_chunk_test_record_prepare(context, 1, block_size) };
-        if context.prepare_calls < context.prepared_sizes.len() {
+    }
-            context.prepared_sizes[context.prepare_calls] = block_size;
+
-        }
+    unsafe extern "C" fn frame_chunk_test_check_dict_validity(
-        context.prepare_calls += 1;
+        context: *mut c_void,
+        _src: *const c_void,
+        block_size: usize,
+        max_dist: c_uint,
+    ) {
+        unsafe { frame_chunk_test_record_prepare(context, 2, block_size) };
+        unsafe { frame_chunk_test_context(context) }.prepare_max_dist = max_dist;
+    }
+
+    unsafe extern "C" fn frame_chunk_test_enforce_max_dist(
+        context: *mut c_void,
+        _src: *const c_void,
+        block_size: usize,
+        max_dist: c_uint,
+    ) {
+        unsafe { frame_chunk_test_record_prepare(context, 3, block_size) };
+        unsafe { frame_chunk_test_context(context) }.prepare_max_dist = max_dist;
+    }
+
+    unsafe extern "C" fn frame_chunk_test_clamp_next_to_update(context: *mut c_void) {
+        unsafe { frame_chunk_test_record_prepare(context, 4, 0) };
     }
 
     unsafe extern "C" fn frame_chunk_test_target(
@@ -9577,11 +9663,28 @@ mod tests {
         block_splitter_enabled: c_int,
         checksum_flag: c_int,
     ) -> ZSTD_rust_frameChunkState {
-        let context = context as *mut FrameChunkTestContext as *mut c_void;
+        let context_ptr = context as *mut FrameChunkTestContext;
+        let callback_context = context_ptr.cast::<c_void>();
+        unsafe {
+            (*context_ptr).prepare_state = Some(ZSTD_rust_frameChunkPrepareState {
+                callback_context,
+                max_dist: 64,
+                correct_overflow: frame_chunk_test_correct_overflow,
+                check_dict_validity: frame_chunk_test_check_dict_validity,
+                enforce_max_dist: frame_chunk_test_enforce_max_dist,
+                clamp_next_to_update: frame_chunk_test_clamp_next_to_update,
+            });
+        }
+        let prepare_state = unsafe {
+            (*context_ptr)
+                .prepare_state
+                .as_ref()
+                .map_or(ptr::null(), |state| state as *const _)
+        };
         ZSTD_rust_frameChunkState {
-            callback_context: context,
+            callback_context,
             tmp_workspace: ptr::null_mut(),
-            checksum_state: context,
+            checksum_state: callback_context,
             is_first_block,
             stage,
             tmp_wksp_size: 0,
@@ -9593,7 +9696,7 @@ mod tests {
             block_splitter_enabled,
             checksum_flag,
             ending_stage: 77,
-            prepare_block: frame_chunk_test_prepare,
+            prepare_state,
             compress_target: frame_chunk_test_target,
             compress_split: frame_chunk_test_split,
             compress_internal: frame_chunk_test_internal,
@@ -9627,6 +9730,8 @@ mod tests {
         assert_eq!(result, 10);
         assert_eq!(context.prepare_calls, 2);
         assert_eq!(context.prepared_sizes[..2], [4, 4]);
+        assert_eq!(context.prepare_order[..8], [1, 2, 3, 4, 1, 2, 3, 4]);
+        assert_eq!(context.prepare_max_dist, 64);
         assert_eq!(context.internal_calls, 2);
         assert_eq!(context.last_blocks[..2], [0, 1]);
         assert_eq!(context.checksum_calls, 1);