feat(compress): project MT frame headers into Rust

Pass the non-first multithreaded frame-header operation through an explicit
scalar projection so Rust writes the header, advances the compression stage,
and invalidates repcodes without calling the public C continue wrapper.
Retain job descriptors, pools, synchronization, and CCtx initialization in C.

Test Plan:
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo fmt --all -- --check
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo test --lib zstdmt_compress::tests::compression_job -- --nocapture
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo clippy --all-targets -- -D warnings
- ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo test
- 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/zstdmt_compress.c

@@ -121,10 +121,59 @@ typedef struct {
     size_t lastBlockSize;
 } ZSTDMT_chunkProcessResult;
 
+typedef struct {
+    void* dst;
+    size_t dstCapacity;
+    int* stage;
+    int noDictIDFlag;
+    int checksumFlag;
+    int contentSizeFlag;
+    int format;
+    U32 windowLog;
+    U64 pledgedSrcSizePlusOne;
+    U32 dictID;
+    U32* repCodes;
+} ZSTDMT_RustCompressionJobFrameHeaderState;
+typedef char ZSTDMT_compression_job_frame_header_state_layout[
+        (offsetof(ZSTDMT_RustCompressionJobFrameHeaderState, dst) == 0
+         && offsetof(ZSTDMT_RustCompressionJobFrameHeaderState, dstCapacity)
+                == sizeof(void*)
+         && offsetof(ZSTDMT_RustCompressionJobFrameHeaderState, stage)
+                == 2 * sizeof(void*)
+         && offsetof(ZSTDMT_RustCompressionJobFrameHeaderState, noDictIDFlag)
+                == 3 * sizeof(void*)
+         && offsetof(ZSTDMT_RustCompressionJobFrameHeaderState, checksumFlag)
+                == 3 * sizeof(void*) + sizeof(int)
+         && offsetof(ZSTDMT_RustCompressionJobFrameHeaderState, contentSizeFlag)
+                == 3 * sizeof(void*) + 2 * sizeof(int)
+         && offsetof(ZSTDMT_RustCompressionJobFrameHeaderState, format)
+                == 3 * sizeof(void*) + 3 * sizeof(int)
+         && offsetof(ZSTDMT_RustCompressionJobFrameHeaderState, windowLog)
+                == 3 * sizeof(void*) + 4 * sizeof(int)
+         && offsetof(ZSTDMT_RustCompressionJobFrameHeaderState, pledgedSrcSizePlusOne)
+                == (sizeof(void*) == 8 ? 6 * sizeof(void*) : 8 * sizeof(void*))
+         && offsetof(ZSTDMT_RustCompressionJobFrameHeaderState, dictID)
+                == (sizeof(void*) == 8 ? 6 * sizeof(void*) : 8 * sizeof(void*))
+                        + sizeof(U64)
+         && offsetof(ZSTDMT_RustCompressionJobFrameHeaderState, repCodes)
+                == (sizeof(void*) == 8 ? 8 * sizeof(void*) : 11 * sizeof(void*))
+         && sizeof(ZSTDMT_RustCompressionJobFrameHeaderState)
+                == (sizeof(void*) == 8 ? 72 : 48))
+        ? 1 : -1];
+
 typedef struct {
     unsigned firstJob;
     unsigned lastJob;
+    const ZSTDMT_RustCompressionJobFrameHeaderState* frameHeaderState;
 } ZSTDMT_RustCompressionJobProjection;
+typedef char ZSTDMT_compression_job_projection_layout[
+        (offsetof(ZSTDMT_RustCompressionJobProjection, firstJob) == 0
+         && offsetof(ZSTDMT_RustCompressionJobProjection, lastJob) == sizeof(unsigned)
+         && offsetof(ZSTDMT_RustCompressionJobProjection, frameHeaderState)
+                == 2 * sizeof(unsigned)
+         && sizeof(ZSTDMT_RustCompressionJobProjection)
+                == 2 * sizeof(unsigned) + sizeof(void*))
+        ? 1 : -1];
 
 typedef struct {
     int status;
@@ -149,7 +198,6 @@ void ZSTDMT_rust_compressionJob(
         ZSTDMT_compressionJobVoidFn generateSequences,
         ZSTDMT_compressionJobStepFn beginJob,
         ZSTDMT_compressionJobVoidFn applySequences,
-        ZSTDMT_compressionJobStepFn writeFrameHeader,
         ZSTDMT_compressionJobCompressFn compressJob,
         ZSTDMT_compressionJobVoidFn traceJob,
         ZSTDMT_compressionJobErrorFn setError,
@@ -1355,6 +1403,7 @@ typedef struct {
     ZSTD_CCtx* cctx;
     RawSeqStore_t rawSeqStore;
     Buffer dstBuff;
+    ZSTDMT_RustCompressionJobFrameHeaderState frameHeaderState;
 } ZSTDMT_compressionJobState;
 
 static size_t ZSTDMT_compressionJobAcquireResources(void* opaque)
@@ -1373,6 +1422,8 @@ static size_t ZSTDMT_compressionJobAcquireResources(void* opaque)
         if (state->dstBuff.start == NULL) return ERROR(memory_allocation);
         job->dstBuff = state->dstBuff;
     }
+    state->frameHeaderState.dst = state->dstBuff.start;
+    state->frameHeaderState.dstCapacity = state->dstBuff.capacity;
     if (state->jobParams.ldmParams.enableLdm == ZSTD_ps_enable &&
         state->rawSeqStore.seq == NULL)
         return ERROR(memory_allocation);
@@ -1418,18 +1469,17 @@ static size_t ZSTDMT_compressionJobBegin(void* opaque)
     ZSTDMT_compressionJobState* const state =
         (ZSTDMT_compressionJobState*)opaque;
     ZSTDMT_jobDescription* const job = state->job;
+    size_t initError;
 
     if (job->cdict) {
-        size_t const initError = ZSTD_compressBegin_advanced_internal(
+        initError = ZSTD_compressBegin_advanced_internal(
             state->cctx, NULL, 0, ZSTD_dct_auto, ZSTD_dtlm_fast, job->cdict,
             &state->jobParams, job->fullFrameSize);
         assert(job->firstJob);  /* only allowed for first job */
-        return initError;
+    } else {
-    }
+        U64 const pledgedSrcSize = job->firstJob ? job->fullFrameSize : job->src.size;
-
-    {   U64 const pledgedSrcSize = job->firstJob ? job->fullFrameSize : job->src.size;
         size_t const forceWindowError = ZSTD_CCtxParams_setParameter(
-            &state->jobParams, ZSTD_c_forceMaxWindow, !job->firstJob);
+                &state->jobParams, ZSTD_c_forceMaxWindow, !job->firstJob);
         if (ZSTD_isError(forceWindowError)) return forceWindowError;
         if (!job->firstJob) {
             size_t const err = ZSTD_CCtxParams_setParameter(
@@ -1437,11 +1487,29 @@ static size_t ZSTDMT_compressionJobBegin(void* opaque)
             if (ZSTD_isError(err)) return err;
         }
         DEBUGLOG(6, "ZSTDMT_compressionJob: job %u: loading prefix of size %zu", job->jobID, job->prefix.size);
-        return ZSTD_compressBegin_advanced_internal(
+        initError = ZSTD_compressBegin_advanced_internal(
             state->cctx, job->prefix.start, job->prefix.size,
             ZSTD_dct_rawContent, ZSTD_dtlm_fast, NULL, /*cdict*/
             &state->jobParams, pledgedSrcSize);
     }
+
+    if (ZSTD_isError(initError)) return initError;
+    state->frameHeaderState.stage = (int*)&state->cctx->stage;
+    state->frameHeaderState.noDictIDFlag =
+            state->cctx->appliedParams.fParams.noDictIDFlag;
+    state->frameHeaderState.checksumFlag =
+            state->cctx->appliedParams.fParams.checksumFlag;
+    state->frameHeaderState.contentSizeFlag =
+            state->cctx->appliedParams.fParams.contentSizeFlag;
+    state->frameHeaderState.format = (int)state->cctx->appliedParams.format;
+    state->frameHeaderState.windowLog =
+            state->cctx->appliedParams.cParams.windowLog;
+    state->frameHeaderState.pledgedSrcSizePlusOne =
+            state->cctx->pledgedSrcSizePlusOne;
+    state->frameHeaderState.dictID = state->cctx->dictID;
+    state->frameHeaderState.repCodes = state->cctx->blockState.prevCBlock == NULL
+            ? NULL : state->cctx->blockState.prevCBlock->rep;
+    return initError;
 }
 
 static void ZSTDMT_compressionJobApplySequences(void* opaque)
@@ -1455,20 +1523,6 @@ static void ZSTDMT_compressionJobApplySequences(void* opaque)
                                       &state->rawSeqStore);
 }
 
-static size_t ZSTDMT_compressionJobWriteFrameHeader(void* opaque)
-{
-    ZSTDMT_compressionJobState* const state =
-        (ZSTDMT_compressionJobState*)opaque;
-    ZSTDMT_jobDescription* const job = state->job;
-    size_t const hSize = ZSTD_compressContinue_public(
-        state->cctx, state->dstBuff.start, state->dstBuff.capacity,
-        job->src.start, 0);
-    if (ZSTD_isError(hSize)) return hSize;
-    DEBUGLOG(5, "ZSTDMT_compressionJob: flush and overwrite %u bytes of frame header (not first job)", (U32)hSize);
-    ZSTD_invalidateRepCodes(state->cctx);
-    return hSize;
-}
-
 static ZSTDMT_chunkProcessResult ZSTDMT_compressionJobCompress(
         void* opaque, unsigned lastJob)
 {
@@ -1545,7 +1599,8 @@ static void ZSTDMT_compressionJob(void* jobDescription)
     ZSTDMT_compressionJobState state;
     ZSTDMT_RustCompressionJobProjection const projection = {
         job->firstJob,
-        job->lastJob
+        job->lastJob,
+        &state.frameHeaderState
     };
 
     ZSTD_memset(&state, 0, sizeof(state));
@@ -1560,7 +1615,6 @@ static void ZSTDMT_compressionJob(void* jobDescription)
         ZSTDMT_compressionJobGenerateSequences,
         ZSTDMT_compressionJobBegin,
         ZSTDMT_compressionJobApplySequences,
-        ZSTDMT_compressionJobWriteFrameHeader,
         ZSTDMT_compressionJobCompress,
         ZSTDMT_compressionJobTrace,
         ZSTDMT_compressionJobSetError,

rust/src/zstdmt_compress.rs

@@ -21,7 +21,8 @@ use std::sync::Mutex;
 
 use crate::bits::ZSTD_highbit32;
 use crate::errors::{ERR_isError, ZstdErrorCode, ERROR};
-use crate::zstd_compress::ZSTD_frameProgression;
+use crate::zstd_compress::{ZSTD_frameProgression, ZSTD_rust_invalidateRepCodes};
+use crate::zstd_compress_frame::ZSTD_rust_writeFrameHeader;
 
 const ZSTDMT_JOBLOG_MAX: c_uint = if mem::size_of::<usize>() == 4 { 29 } else { 30 };
 const ZSTD_WINDOWLOG_MAX: c_uint = if mem::size_of::<usize>() == 4 { 30 } else { 31 };
@@ -68,13 +69,145 @@ pub struct ZSTDMT_chunkProcessResult {
 /// Scalar job state used by the Rust compression-job scheduler.
 ///
 /// The job descriptor, pools, synchronization, and codec state remain
-/// private to C. Rust uses only the frame-position flags to choose the
+/// private to C. Rust uses the frame-position flags and a scalar header
-/// sequencing and non-first-job header stages.
+/// projection to choose the non-first-job header stage.
 #[repr(C)]
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
 pub struct ZSTDMT_compressionJobProjection {
     pub firstJob: c_uint,
     pub lastJob: c_uint,
+    pub frameHeaderState: *const ZSTDMT_compressionJobFrameHeaderState,
+}
+
+/// Scalar projection for the non-first MT job frame header. The destination,
+/// stage, and repcodes point at the worker-owned C storage; all header
+/// serialization and repcode invalidation happen in Rust.
+#[repr(C)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub struct ZSTDMT_compressionJobFrameHeaderState {
+    dst: *mut c_void,
+    dst_capacity: usize,
+    stage: *mut c_int,
+    no_dict_id_flag: c_int,
+    checksum_flag: c_int,
+    content_size_flag: c_int,
+    format: c_int,
+    window_log: c_uint,
+    pledged_src_size_plus_one: u64,
+    dict_id: c_uint,
+    rep_codes: *mut c_uint,
+}
+
+const _: () = {
+    assert!(offset_of!(ZSTDMT_compressionJobProjection, firstJob) == 0);
+    assert!(offset_of!(ZSTDMT_compressionJobProjection, lastJob) == size_of::<c_uint>());
+    assert!(
+        offset_of!(ZSTDMT_compressionJobProjection, frameHeaderState) == size_of::<[c_uint; 2]>()
+    );
+    assert!(
+        size_of::<ZSTDMT_compressionJobProjection>()
+            == size_of::<[c_uint; 2]>() + size_of::<usize>()
+    );
+    assert!(offset_of!(ZSTDMT_compressionJobFrameHeaderState, dst) == 0);
+    assert!(offset_of!(ZSTDMT_compressionJobFrameHeaderState, dst_capacity) == size_of::<usize>());
+    assert!(offset_of!(ZSTDMT_compressionJobFrameHeaderState, stage) == 2 * size_of::<usize>());
+    assert!(
+        offset_of!(ZSTDMT_compressionJobFrameHeaderState, no_dict_id_flag)
+            == 3 * size_of::<usize>()
+    );
+    assert!(
+        offset_of!(ZSTDMT_compressionJobFrameHeaderState, checksum_flag)
+            == 3 * size_of::<usize>() + size_of::<c_int>()
+    );
+    assert!(
+        offset_of!(ZSTDMT_compressionJobFrameHeaderState, content_size_flag)
+            == 3 * size_of::<usize>() + 2 * size_of::<c_int>()
+    );
+    assert!(
+        offset_of!(ZSTDMT_compressionJobFrameHeaderState, format)
+            == 3 * size_of::<usize>() + 3 * size_of::<c_int>()
+    );
+    assert!(
+        offset_of!(ZSTDMT_compressionJobFrameHeaderState, window_log)
+            == 3 * size_of::<usize>() + 4 * size_of::<c_int>()
+    );
+    assert!(
+        offset_of!(
+            ZSTDMT_compressionJobFrameHeaderState,
+            pledged_src_size_plus_one
+        ) == ZSTDMT_HEADER_PLEDGED_OFFSET
+    );
+    assert!(
+        offset_of!(ZSTDMT_compressionJobFrameHeaderState, dict_id)
+            == ZSTDMT_HEADER_PLEDGED_OFFSET + size_of::<u64>()
+    );
+    assert!(
+        offset_of!(ZSTDMT_compressionJobFrameHeaderState, rep_codes)
+            == ZSTDMT_HEADER_REP_CODES_OFFSET
+    );
+    assert!(
+        size_of::<ZSTDMT_compressionJobFrameHeaderState>()
+            == if size_of::<usize>() == 8 { 72 } else { 48 }
+    );
+};
+
+const ZSTD_COMPRESSION_STAGE_CREATED: c_int = 0;
+const ZSTD_COMPRESSION_STAGE_INIT: c_int = 1;
+const ZSTD_COMPRESSION_STAGE_ONGOING: c_int = 2;
+const ZSTDMT_HEADER_PLEDGED_OFFSET: usize = if size_of::<usize>() == 8 {
+    size_of::<[usize; 6]>()
+} else {
+    size_of::<[usize; 8]>()
+};
+const ZSTDMT_HEADER_REP_CODES_OFFSET: usize = if size_of::<usize>() == 8 {
+    size_of::<[usize; 8]>()
+} else {
+    size_of::<[usize; 11]>()
+};
+
+#[no_mangle]
+pub unsafe extern "C" fn ZSTDMT_rust_writeFrameHeader(
+    state: *const ZSTDMT_compressionJobFrameHeaderState,
+) -> usize {
+    if state.is_null() {
+        return ERROR(ZstdErrorCode::Generic);
+    }
+    let state = unsafe { &*state };
+    if state.stage.is_null() || state.rep_codes.is_null() {
+        return ERROR(ZstdErrorCode::Generic);
+    }
+    if unsafe { *state.stage } == ZSTD_COMPRESSION_STAGE_CREATED {
+        return ERROR(ZstdErrorCode::StageWrong);
+    }
+
+    let header_size = if unsafe { *state.stage } == ZSTD_COMPRESSION_STAGE_INIT {
+        let header_size = unsafe {
+            ZSTD_rust_writeFrameHeader(
+                state.dst,
+                state.dst_capacity,
+                state.no_dict_id_flag,
+                state.checksum_flag,
+                state.content_size_flag,
+                state.format,
+                state.window_log,
+                state.pledged_src_size_plus_one.wrapping_sub(1),
+                state.dict_id,
+            )
+        };
+        if ERR_isError(header_size) {
+            return header_size;
+        }
+        if header_size > state.dst_capacity {
+            return ERROR(ZstdErrorCode::DstSizeTooSmall);
+        }
+        unsafe { *state.stage = ZSTD_COMPRESSION_STAGE_ONGOING };
+        header_size
+    } else {
+        0
+    };
+
+    unsafe { ZSTD_rust_invalidateRepCodes(state.rep_codes) };
+    header_size
 }
 
 pub type ZSTDMT_compressionJobStepFn = unsafe extern "C" fn(*mut c_void) -> usize;
@@ -1116,7 +1249,6 @@ pub unsafe extern "C" fn ZSTDMT_rust_compressionJob(
     generateSequences: Option<ZSTDMT_compressionJobVoidFn>,
     beginJob: Option<ZSTDMT_compressionJobStepFn>,
     applySequences: Option<ZSTDMT_compressionJobVoidFn>,
-    writeFrameHeader: Option<ZSTDMT_compressionJobStepFn>,
     compressJob: Option<ZSTDMT_compressionJobCompressFn>,
     traceJob: Option<ZSTDMT_compressionJobVoidFn>,
     setError: Option<ZSTDMT_compressionJobErrorFn>,
@@ -1131,7 +1263,6 @@ pub unsafe extern "C" fn ZSTDMT_rust_compressionJob(
         Some(generate_sequences),
         Some(begin_job),
         Some(apply_sequences),
-        Some(write_frame_header),
         Some(compress_job),
         Some(trace_job),
         Some(set_error),
@@ -1142,7 +1273,6 @@ pub unsafe extern "C" fn ZSTDMT_rust_compressionJob(
         generateSequences,
         beginJob,
         applySequences,
-        writeFrameHeader,
         compressJob,
         traceJob,
         setError,
@@ -1159,7 +1289,7 @@ pub unsafe extern "C" fn ZSTDMT_rust_compressionJob(
         || unsafe { generate_sequences(opaque) },
         || unsafe { begin_job(opaque) },
         || unsafe { apply_sequences(opaque) },
-        || unsafe { write_frame_header(opaque) },
+        || unsafe { ZSTDMT_rust_writeFrameHeader(projection.frameHeaderState) },
         |last_job| unsafe { compress_job(opaque, last_job) },
         || unsafe { trace_job(opaque) },
         |error| unsafe { set_error(opaque, error) },
@@ -3994,6 +4124,43 @@ mod tests {
         state.borrow_mut().events.push(event);
     }
 
+    #[test]
+    fn compression_job_frame_header_projection_writes_header_and_clears_repcodes() {
+        let mut stage = ZSTD_COMPRESSION_STAGE_INIT;
+        let mut rep_codes = [11u32, 22, 33];
+        let mut output = [0xa5u8; 18];
+        let state = ZSTDMT_compressionJobFrameHeaderState {
+            dst: output.as_mut_ptr().cast(),
+            dst_capacity: output.len(),
+            stage: &mut stage,
+            no_dict_id_flag: 0,
+            checksum_flag: 0,
+            content_size_flag: 1,
+            format: 0,
+            window_log: 20,
+            pledged_src_size_plus_one: 1,
+            dict_id: 0,
+            rep_codes: rep_codes.as_mut_ptr(),
+        };
+
+        let mut created_stage = ZSTD_COMPRESSION_STAGE_CREATED;
+        let created_state = ZSTDMT_compressionJobFrameHeaderState {
+            stage: &mut created_stage,
+            ..state
+        };
+        let created = unsafe { ZSTDMT_rust_writeFrameHeader(&created_state) };
+        assert_eq!(created, ERROR(ZstdErrorCode::StageWrong));
+        assert_eq!(created_stage, ZSTD_COMPRESSION_STAGE_CREATED);
+        assert_eq!(rep_codes, [11, 22, 33]);
+        let result = unsafe { ZSTDMT_rust_writeFrameHeader(&state) };
+
+        assert_eq!(result, 6);
+        assert_eq!(stage, ZSTD_COMPRESSION_STAGE_ONGOING);
+        assert_eq!(rep_codes, [0; 3]);
+        assert_eq!(&output[..result], &[0x28, 0xb5, 0x2f, 0xfd, 0x20, 0x00]);
+        assert_eq!(&output[result..], &[0xa5; 12]);
+    }
+
     #[test]
     fn compression_job_runs_first_job_stages_and_reports_final_block() {
         let state = Rc::new(RefCell::new(MockCompressionJob::default()));
@@ -4010,6 +4177,7 @@ mod tests {
             ZSTDMT_compressionJobProjection {
                 firstJob: 1,
                 lastJob: 1,
+                frameHeaderState: ptr::null(),
             },
             move || {
                 record_compression_job_event(&acquire_state, "acquire");
@@ -4250,6 +4418,7 @@ mod tests {
             ZSTDMT_compressionJobProjection {
                 firstJob: 0,
                 lastJob: 1,
+                frameHeaderState: ptr::null(),
             },
             move || {
                 record_compression_job_event(&acquire_state, "acquire");