From a4a64f2e38312ddb776cc9d824ef0f87a2fc882d Mon Sep 17 00:00:00 2001
From: ddidderr <ddidderr@paul.network>
Date: Sat, 18 Jul 2026 17:45:30 +0200
Subject: [PATCH] feat(compress): move pledged source setter into Rust

Move the active-stage check and pledged-size update behind a scalar Rust ABI
leaf while keeping the CCtx layout and public C entry point in C. Preserve the
stage_wrong error encoding and the unsigned wrap behavior at U64::MAX.

Test Plan:
- cargo test --manifest-path rust/Cargo.toml --lib zstd_compress::tests (80 passed)
- cargo clippy --manifest-path rust/Cargo.toml --lib -- -D warnings
- rustfmt +nightly --edition 2021 rust/src/zstd_compress.rs --check
- make -C programs -j2 zstd
- git diff --check
---
 lib/compress/zstd_compress.c |  9 +++--
 rust/src/zstd_compress.rs    | 75 ++++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 4 deletions(-)

diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c
index 2623f027e..16fa57639 100644
--- a/lib/compress/zstd_compress.c
+++ b/lib/compress/zstd_compress.c
@@ -48,6 +48,9 @@ int ZSTD_rust_updateFrameProgression(unsigned long long* consumedSrcSize,
                                      unsigned long long pledgedSrcSizePlusOne,
                                      size_t srcSize, size_t cSize,
                                      size_t fhSize);
+size_t ZSTD_rust_setPledgedSrcSize(int streamStage,
+                                   unsigned long long pledgedSrcSize,
+                                   unsigned long long* pledgedSrcSizePlusOne);
 ZSTD_frameProgression ZSTD_rust_frameProgression(U64 consumedSrcSize,
                                                   size_t buffered,
                                                   U64 producedCSize);
@@ -974,10 +977,8 @@ size_t ZSTD_CCtx_setParams(ZSTD_CCtx* cctx, ZSTD_parameters params)
 size_t ZSTD_CCtx_setPledgedSrcSize(ZSTD_CCtx* cctx, unsigned long long pledgedSrcSize)
 {
     DEBUGLOG(4, "ZSTD_CCtx_setPledgedSrcSize to %llu bytes", pledgedSrcSize);
-    RETURN_ERROR_IF(cctx->streamStage != zcss_init, stage_wrong,
-                    "Can't set pledgedSrcSize when not in init stage.");
-    cctx->pledgedSrcSizePlusOne = pledgedSrcSize+1;
-    return 0;
+    return ZSTD_rust_setPledgedSrcSize((int)cctx->streamStage, pledgedSrcSize,
+                                       &cctx->pledgedSrcSizePlusOne);
 }
 
 static ZSTD_compressionParameters ZSTD_dedicatedDictSearch_getCParams(
diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs
index a53fd4d07..99ef1070d 100644
--- a/rust/src/zstd_compress.rs
+++ b/rust/src/zstd_compress.rs
@@ -88,6 +88,7 @@ const ZSTD_TARGET_CBLOCK_BSS_COMPRESS: c_int = 0;
 const ZSTD_BLOCK_HEADER_SIZE: usize = 3;
 const MIN_COMPRESSIBLE_BLOCK_SIZE: usize = 2 + ZSTD_BLOCK_HEADER_SIZE + 1 + 1;
 const ZSTD_ROWSIZE: usize = 16;
+const ZSTD_CSTREAM_STAGE_INIT: c_int = 0;
 const ZSTD_WINDOW_START_INDEX: u32 = 2;
 const ZSTD_DUBT_UNSORTED_MARK: u32 = 1;
 const ZSTD_INDEXOVERFLOW_MARGIN: usize = 16usize << 20;
@@ -877,6 +878,34 @@ pub unsafe extern "C" fn ZSTD_rust_windowClear(
     }
 }
 
+#[inline]
+fn set_pledged_src_size(
+    stream_stage: c_int,
+    pledged_src_size: u64,
+    pledged_src_size_plus_one: *mut u64,
+) -> usize {
+    if stream_stage != ZSTD_CSTREAM_STAGE_INIT {
+        return ERROR(ZstdErrorCode::StageWrong);
+    }
+    debug_assert!(!pledged_src_size_plus_one.is_null());
+    unsafe {
+        *pledged_src_size_plus_one = pledged_src_size.wrapping_add(1);
+    }
+    0
+}
+
+/// Set the pledged source size through scalar stage state and a C-owned field.
+/// The stage error remains the public `stage_wrong` encoding, while the
+/// unsigned increment preserves C's defined wrap at `U64::MAX`.
+#[no_mangle]
+pub unsafe extern "C" fn ZSTD_rust_setPledgedSrcSize(
+    stream_stage: c_int,
+    pledged_src_size: u64,
+    pledged_src_size_plus_one: *mut u64,
+) -> usize {
+    set_pledged_src_size(stream_stage, pledged_src_size, pledged_src_size_plus_one)
+}
+
 #[inline]
 fn zeroed_state() -> ZSTD_compressedBlockState_t {
     /* The state contains only integer arrays and enum fields. */
@@ -2188,6 +2217,52 @@ mod tests {
         assert_eq!(dict_limit, 0x89ab_cdef);
     }
 
+    #[test]
+    fn pledged_src_size_writes_the_init_stage_value_plus_one() {
+        let mut pledged_src_size_plus_one = 0;
+
+        let result = unsafe {
+            ZSTD_rust_setPledgedSrcSize(
+                ZSTD_CSTREAM_STAGE_INIT,
+                123,
+                &mut pledged_src_size_plus_one,
+            )
+        };
+
+        assert_eq!(result, 0);
+        assert_eq!(pledged_src_size_plus_one, 124);
+    }
+
+    #[test]
+    fn pledged_src_size_wraps_at_u64_max() {
+        let mut pledged_src_size_plus_one = 123;
+
+        let result = unsafe {
+            ZSTD_rust_setPledgedSrcSize(
+                ZSTD_CSTREAM_STAGE_INIT,
+                u64::MAX,
+                &mut pledged_src_size_plus_one,
+            )
+        };
+
+        assert_eq!(result, 0);
+        assert_eq!(pledged_src_size_plus_one, 0);
+    }
+
+    #[test]
+    fn pledged_src_size_rejects_non_init_stage_with_stage_wrong() {
+        let mut pledged_src_size_plus_one = 123;
+
+        let result = unsafe { ZSTD_rust_setPledgedSrcSize(1, 456, &mut pledged_src_size_plus_one) };
+
+        assert_eq!(result, ERROR(ZstdErrorCode::StageWrong));
+        assert_eq!(
+            crate::errors::ERR_getErrorCode(result),
+            ZstdErrorCode::StageWrong as i32
+        );
+        assert_eq!(pledged_src_size_plus_one, 123);
+    }
+
     #[test]
     fn sizeof_local_dict_ignores_size_without_a_buffer() {
         assert_eq!(sizeof_local_dict(0, 37, 11), 11);