feat(compress): move overflow correction policy into Rust

Move the overflow-correction branch and ordering into the Rust rewrite while
keeping the C-owned window correction, workspace table markers, index reducer,
and dictionary fields behind callbacks. Rust now preserves the original
need-correction fast path, dirty/reduce/clean ordering, saturating
nextToUpdate adjustment, and dictionary invalidation. Add focused callback
ordering and safe-window tests.

Test Plan:
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo fmt --manifest-path rust/Cargo.toml -- --check
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/Cargo.toml overflow_correction -- --nocapture
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/Cargo.toml
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings
- ulimit -v 41943040 make -B -C programs -j1 zstd
- ulimit -v 41943040 make -C tests -j1 test-zstream ZSTREAM_TESTTIME=-T1s

lib/compress/zstd_compress.c

@@ -609,6 +609,36 @@ void ZSTD_rust_reduceIndex(U32* hashTable, U32 hashSize,
                            U32* chainTable, U32 chainSize,
                            U32* hashTable3, U32 hashSize3,
                            U32 reducerValue, int preserveChainMark);
+typedef int (*ZSTD_rust_overflowNeedCorrection_f)(
+        void* context, const void* src, const void* srcEnd);
+typedef U32 (*ZSTD_rust_overflowCorrect_f)(
+        void* context, const void* src);
+typedef void (*ZSTD_rust_overflowCallback_f)(void* context);
+typedef void (*ZSTD_rust_overflowReduceIndex_f)(void* context, U32 correction);
+typedef struct {
+    void* callbackContext;
+    U32* nextToUpdate;
+    ZSTD_rust_overflowNeedCorrection_f needCorrection;
+    ZSTD_rust_overflowCorrect_f correctOverflow;
+    ZSTD_rust_overflowCallback_f markTablesDirty;
+    ZSTD_rust_overflowReduceIndex_f reduceIndex;
+    ZSTD_rust_overflowCallback_f markTablesClean;
+    ZSTD_rust_overflowCallback_f invalidateDictionary;
+} ZSTD_rust_overflowCorrectState;
+void ZSTD_rust_overflowCorrectIfNeeded(
+        const ZSTD_rust_overflowCorrectState* state,
+        const void* src, const void* srcEnd);
+typedef char ZSTD_rust_overflow_correct_state_layout[
+        (offsetof(ZSTD_rust_overflowCorrectState, callbackContext) == 0
+         && offsetof(ZSTD_rust_overflowCorrectState, nextToUpdate) == sizeof(void*)
+         && offsetof(ZSTD_rust_overflowCorrectState, needCorrection) == 2 * sizeof(void*)
+         && offsetof(ZSTD_rust_overflowCorrectState, correctOverflow) == 3 * sizeof(void*)
+         && offsetof(ZSTD_rust_overflowCorrectState, markTablesDirty) == 4 * sizeof(void*)
+         && offsetof(ZSTD_rust_overflowCorrectState, reduceIndex) == 5 * sizeof(void*)
+         && offsetof(ZSTD_rust_overflowCorrectState, markTablesClean) == 6 * sizeof(void*)
+         && offsetof(ZSTD_rust_overflowCorrectState, invalidateDictionary) == 7 * sizeof(void*)
+         && sizeof(ZSTD_rust_overflowCorrectState) == 8 * sizeof(void*))
+                ? 1 : -1];
 void ZSTD_rust_copyCDictTableIntoCCtx(U32* dst, U32 const* src,
                                       size_t tableSize, int tagged);
 U64 ZSTD_rust_advanceHashSalt(U64 hashSalt, U64 hashSaltEntropy);
@@ -5053,28 +5083,90 @@ static size_t ZSTD_compressBlock_targetCBlockSize(ZSTD_CCtx* zc,
     return cSize;
 }
 
+typedef struct {
+    ZSTD_MatchState_t* matchState;
+    ZSTD_cwksp* workspace;
+    const ZSTD_CCtx_params* params;
+} ZSTD_rust_overflowCorrectContext;
+
+static int ZSTD_rust_overflowCorrect_need(
+        void* context, const void* src, const void* srcEnd)
+{
+    ZSTD_rust_overflowCorrectContext const* const state =
+            (const ZSTD_rust_overflowCorrectContext*)context;
+    U32 const cycleLog = ZSTD_cycleLog(
+            state->params->cParams.chainLog,
+            state->params->cParams.strategy);
+    U32 const maxDist = (U32)1 << state->params->cParams.windowLog;
+    return (int)ZSTD_window_needOverflowCorrection(
+            state->matchState->window, cycleLog, maxDist,
+            state->matchState->loadedDictEnd, src, srcEnd);
+}
+
+static U32 ZSTD_rust_overflowCorrect_correct(void* context, const void* src)
+{
+    ZSTD_rust_overflowCorrectContext const* const state =
+            (const ZSTD_rust_overflowCorrectContext*)context;
+    U32 const cycleLog = ZSTD_cycleLog(
+            state->params->cParams.chainLog,
+            state->params->cParams.strategy);
+    U32 const maxDist = (U32)1 << state->params->cParams.windowLog;
+    return ZSTD_window_correctOverflow(
+            &state->matchState->window, cycleLog, maxDist, src);
+}
+
+static void ZSTD_rust_overflowCorrect_markTablesDirty(void* context)
+{
+    ZSTD_rust_overflowCorrectContext const* const state =
+            (const ZSTD_rust_overflowCorrectContext*)context;
+    ZSTD_cwksp_mark_tables_dirty(state->workspace);
+}
+
+static void ZSTD_rust_overflowCorrect_reduceIndex(void* context, U32 correction)
+{
+    ZSTD_rust_overflowCorrectContext const* const state =
+            (const ZSTD_rust_overflowCorrectContext*)context;
+    ZSTD_reduceIndex(state->matchState, state->params, correction);
+}
+
+static void ZSTD_rust_overflowCorrect_markTablesClean(void* context)
+{
+    ZSTD_rust_overflowCorrectContext const* const state =
+            (const ZSTD_rust_overflowCorrectContext*)context;
+    ZSTD_cwksp_mark_tables_clean(state->workspace);
+}
+
+static void ZSTD_rust_overflowCorrect_invalidateDictionary(void* context)
+{
+    ZSTD_rust_overflowCorrectContext const* const state =
+            (const ZSTD_rust_overflowCorrectContext*)context;
+    state->matchState->loadedDictEnd = 0;
+    state->matchState->dictMatchState = NULL;
+}
+
 static void ZSTD_overflowCorrectIfNeeded(ZSTD_MatchState_t* ms,
                                          ZSTD_cwksp* ws,
                                          ZSTD_CCtx_params const* params,
                                          void const* ip,
                                          void const* iend)
 {
-    U32 const cycleLog = ZSTD_cycleLog(params->cParams.chainLog, params->cParams.strategy);
+    ZSTD_rust_overflowCorrectContext context;
-    U32 const maxDist = (U32)1 << params->cParams.windowLog;
+    ZSTD_rust_overflowCorrectState state;
-    if (ZSTD_window_needOverflowCorrection(ms->window, cycleLog, maxDist, ms->loadedDictEnd, ip, iend)) {
+    ZSTD_STATIC_ASSERT(ZSTD_CHAINLOG_MAX <= 30);
-        U32 const correction = ZSTD_window_correctOverflow(&ms->window, cycleLog, maxDist, ip);
+    ZSTD_STATIC_ASSERT(ZSTD_WINDOWLOG_MAX_32 <= 30);
-        ZSTD_STATIC_ASSERT(ZSTD_CHAINLOG_MAX <= 30);
+    ZSTD_STATIC_ASSERT(ZSTD_WINDOWLOG_MAX <= 31);
-        ZSTD_STATIC_ASSERT(ZSTD_WINDOWLOG_MAX_32 <= 30);
+    context.matchState = ms;
-        ZSTD_STATIC_ASSERT(ZSTD_WINDOWLOG_MAX <= 31);
+    context.workspace = ws;
-        ZSTD_cwksp_mark_tables_dirty(ws);
+    context.params = params;
-        ZSTD_reduceIndex(ms, params, correction);
+    state.callbackContext = &context;
-        ZSTD_cwksp_mark_tables_clean(ws);
+    state.nextToUpdate = &ms->nextToUpdate;
-        if (ms->nextToUpdate < correction) ms->nextToUpdate = 0;
+    state.needCorrection = ZSTD_rust_overflowCorrect_need;
-        else ms->nextToUpdate -= correction;
+    state.correctOverflow = ZSTD_rust_overflowCorrect_correct;
-        /* invalidate dictionaries on overflow correction */
+    state.markTablesDirty = ZSTD_rust_overflowCorrect_markTablesDirty;
-        ms->loadedDictEnd = 0;
+    state.reduceIndex = ZSTD_rust_overflowCorrect_reduceIndex;
-        ms->dictMatchState = NULL;
+    state.markTablesClean = ZSTD_rust_overflowCorrect_markTablesClean;
-    }
+    state.invalidateDictionary = ZSTD_rust_overflowCorrect_invalidateDictionary;
+    ZSTD_rust_overflowCorrectIfNeeded(&state, ip, iend);
 }
 
 #include "zstd_preSplit.h"

rust/README.md

@@ -53,7 +53,9 @@ zstd ABI:
     owns the single-threaded buffered/stable stream state machine, including
     direct versus buffered output, pending-output draining, and frame reset
     policy. The external-sequence-and-literals block loop and public sequence
-    conversion are Rust-owned as well; C retains only CCtx-facing adapters.
+    conversion are Rust-owned as well; Rust also owns overflow-correction
+    branch/order while C retains the private window, workspace, and index
+    callbacks and only CCtx-facing adapters.
   - `zstd_compress_frame` serializes frame headers, skippable frames, and the
     last empty block; it takes scalar frame parameters so the C-owned
     `ZSTD_CCtx_params` layout never crosses the language boundary.
@@ -152,7 +154,7 @@ 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, frame-block preparation ordering, MT serial turn/skip and
-LDM/checksum sequencing, public sequence-API
+LDM/checksum sequencing, overflow-correction policy/order, 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

@@ -5238,6 +5238,96 @@ pub extern "C" fn ZSTD_rust_windowNeedOverflowCorrection(
     ) as u32
 }
 
+type OverflowNeedCorrectionFn =
+    unsafe extern "C" fn(*mut c_void, *const c_void, *const c_void) -> c_int;
+type OverflowCorrectFn = unsafe extern "C" fn(*mut c_void, *const c_void) -> c_uint;
+type OverflowCallbackFn = unsafe extern "C" fn(*mut c_void);
+type OverflowReduceIndexFn = unsafe extern "C" fn(*mut c_void, c_uint);
+
+/// Projection for the overflow-correction ordering around C-owned match state.
+///
+/// Rust owns the correction branch and callback order. C retains the window,
+/// workspace, match tables, and dictionary pointers behind these callbacks.
+#[repr(C)]
+pub struct ZSTD_rust_overflowCorrectState {
+    callback_context: *mut c_void,
+    next_to_update: *mut c_uint,
+    need_correction: Option<OverflowNeedCorrectionFn>,
+    correct_overflow: Option<OverflowCorrectFn>,
+    mark_tables_dirty: Option<OverflowCallbackFn>,
+    reduce_index: Option<OverflowReduceIndexFn>,
+    mark_tables_clean: Option<OverflowCallbackFn>,
+    invalidate_dictionary: Option<OverflowCallbackFn>,
+}
+
+const _: () = {
+    assert!(size_of::<OverflowNeedCorrectionFn>() == size_of::<usize>());
+    assert!(size_of::<OverflowCorrectFn>() == size_of::<usize>());
+    assert!(size_of::<OverflowCallbackFn>() == size_of::<usize>());
+    assert!(size_of::<OverflowReduceIndexFn>() == size_of::<usize>());
+    assert!(offset_of!(ZSTD_rust_overflowCorrectState, callback_context) == 0);
+    assert!(offset_of!(ZSTD_rust_overflowCorrectState, next_to_update) == size_of::<usize>());
+    assert!(offset_of!(ZSTD_rust_overflowCorrectState, need_correction) == 2 * size_of::<usize>());
+    assert!(offset_of!(ZSTD_rust_overflowCorrectState, correct_overflow) == 3 * size_of::<usize>());
+    assert!(
+        offset_of!(ZSTD_rust_overflowCorrectState, mark_tables_dirty) == 4 * size_of::<usize>()
+    );
+    assert!(offset_of!(ZSTD_rust_overflowCorrectState, reduce_index) == 5 * size_of::<usize>());
+    assert!(
+        offset_of!(ZSTD_rust_overflowCorrectState, mark_tables_clean) == 6 * size_of::<usize>()
+    );
+    assert!(
+        offset_of!(ZSTD_rust_overflowCorrectState, invalidate_dictionary) == 7 * size_of::<usize>()
+    );
+    assert!(size_of::<ZSTD_rust_overflowCorrectState>() == size_of::<[usize; 8]>());
+};
+
+/// Apply one overflow correction while keeping all private codec state in C.
+#[no_mangle]
+pub unsafe extern "C" fn ZSTD_rust_overflowCorrectIfNeeded(
+    state: *const ZSTD_rust_overflowCorrectState,
+    src: *const c_void,
+    src_end: *const c_void,
+) {
+    if state.is_null() {
+        return;
+    }
+    let state = unsafe { &*state };
+    if state.callback_context.is_null() || state.next_to_update.is_null() {
+        return;
+    }
+    let (
+        Some(need_correction),
+        Some(correct_overflow),
+        Some(mark_tables_dirty),
+        Some(reduce_index),
+        Some(mark_tables_clean),
+        Some(invalidate_dictionary),
+    ) = (
+        state.need_correction,
+        state.correct_overflow,
+        state.mark_tables_dirty,
+        state.reduce_index,
+        state.mark_tables_clean,
+        state.invalidate_dictionary,
+    )
+    else {
+        return;
+    };
+
+    unsafe {
+        if need_correction(state.callback_context, src, src_end) == 0 {
+            return;
+        }
+        let correction = correct_overflow(state.callback_context, src);
+        mark_tables_dirty(state.callback_context);
+        reduce_index(state.callback_context, correction);
+        mark_tables_clean(state.callback_context);
+        *state.next_to_update = (*state.next_to_update).saturating_sub(correction);
+        invalidate_dictionary(state.callback_context);
+    }
+}
+
 #[inline]
 fn limit_next_to_update(curr: u32, next_to_update: u32) -> u32 {
     let const_gap = 384u32;
@@ -10760,6 +10850,116 @@ mod tests {
         assert_eq!(ZSTD_rust_indexTooCloseToMax(threshold + 1), 1);
     }
 
+    #[derive(Default)]
+    struct OverflowCorrectionTestContext {
+        events: Vec<&'static str>,
+        should_correct: c_int,
+        correction: c_uint,
+    }
+
+    unsafe fn overflow_correction_test_context(
+        context: *mut c_void,
+    ) -> &'static mut OverflowCorrectionTestContext {
+        unsafe { &mut *context.cast::<OverflowCorrectionTestContext>() }
+    }
+
+    unsafe extern "C" fn overflow_correction_test_need(
+        context: *mut c_void,
+        _src: *const c_void,
+        _src_end: *const c_void,
+    ) -> c_int {
+        let context = unsafe { overflow_correction_test_context(context) };
+        context.events.push("need");
+        context.should_correct
+    }
+
+    unsafe extern "C" fn overflow_correction_test_correct(
+        context: *mut c_void,
+        _src: *const c_void,
+    ) -> c_uint {
+        let context = unsafe { overflow_correction_test_context(context) };
+        context.events.push("correct");
+        context.correction
+    }
+
+    unsafe extern "C" fn overflow_correction_test_mark_dirty(context: *mut c_void) {
+        unsafe { overflow_correction_test_context(context) }
+            .events
+            .push("dirty");
+    }
+
+    unsafe extern "C" fn overflow_correction_test_reduce(
+        context: *mut c_void,
+        _correction: c_uint,
+    ) {
+        unsafe { overflow_correction_test_context(context) }
+            .events
+            .push("reduce");
+    }
+
+    unsafe extern "C" fn overflow_correction_test_mark_clean(context: *mut c_void) {
+        unsafe { overflow_correction_test_context(context) }
+            .events
+            .push("clean");
+    }
+
+    unsafe extern "C" fn overflow_correction_test_invalidate(context: *mut c_void) {
+        unsafe { overflow_correction_test_context(context) }
+            .events
+            .push("invalidate");
+    }
+
+    fn overflow_correction_test_state(
+        context: &mut OverflowCorrectionTestContext,
+        next_to_update: &mut c_uint,
+    ) -> ZSTD_rust_overflowCorrectState {
+        ZSTD_rust_overflowCorrectState {
+            callback_context: (context as *mut OverflowCorrectionTestContext).cast(),
+            next_to_update,
+            need_correction: Some(overflow_correction_test_need),
+            correct_overflow: Some(overflow_correction_test_correct),
+            mark_tables_dirty: Some(overflow_correction_test_mark_dirty),
+            reduce_index: Some(overflow_correction_test_reduce),
+            mark_tables_clean: Some(overflow_correction_test_mark_clean),
+            invalidate_dictionary: Some(overflow_correction_test_invalidate),
+        }
+    }
+
+    #[test]
+    fn overflow_correction_preserves_order_and_saturates_next_to_update() {
+        let mut context = OverflowCorrectionTestContext {
+            should_correct: 1,
+            correction: 10,
+            ..Default::default()
+        };
+        let mut next_to_update = 7;
+        let state = overflow_correction_test_state(&mut context, &mut next_to_update);
+
+        unsafe {
+            ZSTD_rust_overflowCorrectIfNeeded(&state, ptr::null(), ptr::null());
+        }
+
+        assert_eq!(next_to_update, 0);
+        assert_eq!(
+            context.events,
+            ["need", "correct", "dirty", "reduce", "clean", "invalidate"]
+        );
+    }
+
+    #[test]
+    fn overflow_correction_skips_callbacks_when_window_is_safe() {
+        let mut context = OverflowCorrectionTestContext::default();
+        let mut next_to_update = 123;
+        let state = overflow_correction_test_state(&mut context, &mut next_to_update);
+
+        unsafe {
+            ZSTD_rust_overflowCorrectIfNeeded(&state, ptr::null(), ptr::null());
+        }
+
+        assert_eq!(next_to_update, 123);
+        assert_eq!(context.events, ["need"]);
+    }
+
     #[test]
     fn window_correction_handles_current_cycle_start_boundary() {
         assert_eq!(window_correct_overflow(0x100, 3, 8), 0xf0);