diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index 095292b00..69a2d36f8 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -1064,7 +1064,8 @@ typedef struct { int inBufferMode; int outBufferMode; int format; - int nbWorkers; + /* Transparent init may resolve a different worker mode for unknown sizes. */ + const int* nbWorkers; } ZSTD_rust_compressStream2Fields; typedef struct { ZSTD_rust_compressStream2Init_f initCompressStream2; @@ -1112,10 +1113,11 @@ typedef char ZSTD_rust_compress_stream2_fields_layout[ && offsetof(ZSTD_rust_compressStream2Fields, format) == 10 * sizeof(void*) + sizeof(unsigned long long) + 2 * sizeof(int) && offsetof(ZSTD_rust_compressStream2Fields, nbWorkers) - == 10 * sizeof(void*) + sizeof(unsigned long long) + 3 * sizeof(int) + == ((10 * sizeof(void*) + sizeof(unsigned long long) + 3 * sizeof(int) + + sizeof(void*) - 1) / sizeof(void*)) * sizeof(void*) && sizeof(ZSTD_rust_compressStream2Fields) - == ((10 * sizeof(void*) + sizeof(unsigned long long) + 4 * sizeof(int) - + sizeof(void*) - 1) / sizeof(void*)) * sizeof(void*)) + == offsetof(ZSTD_rust_compressStream2Fields, nbWorkers) + + sizeof(const int*)) ? 1 : -1]; typedef char ZSTD_rust_compress_stream2_callbacks_layout[ (offsetof(ZSTD_rust_compressStream2Callbacks, initCompressStream2) == 0 @@ -8922,11 +8924,7 @@ size_t ZSTD_compressStream2_c( ZSTD_CCtx* cctx, (int)cctx->requestedParams.inBufferMode, (int)cctx->requestedParams.outBufferMode, (int)cctx->requestedParams.format, -#ifdef ZSTD_MULTITHREAD - (int)cctx->appliedParams.nbWorkers, -#else - 0, -#endif + &cctx->appliedParams.nbWorkers, }; ZSTD_rust_compressStream2Callbacks const callbacks = { ZSTD_rust_compressStream2_init, diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs index 554e6a42f..a3f20c724 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -922,7 +922,8 @@ pub struct ZSTD_rust_compressStream2Fields { in_buffer_mode: c_int, out_buffer_mode: c_int, format: c_int, - nb_workers: c_int, + // Transparent init may update the applied worker mode before dispatch. + nb_workers: *const c_int, } #[repr(C)] @@ -988,14 +989,8 @@ const _: () = { == 10 * size_of::() + size_of::() + 2 * size_of::() ); assert!( - offset_of!(ZSTD_rust_compressStream2Fields, nb_workers) - == 10 * size_of::() + size_of::() + 3 * size_of::() - ); - assert!( - size_of::() - == (10 * size_of::() + size_of::() + 4 * size_of::()) - .div_ceil(size_of::()) - * size_of::() + offset_of!(ZSTD_rust_compressStream2Fields, nb_workers) + size_of::<*const c_int>() + == size_of::() ); assert!(size_of::() == size_of::()); assert!(size_of::() == size_of::()); @@ -1069,7 +1064,7 @@ unsafe fn compress_stream2_fallback(state: &ZSTD_rust_compressStream2State) -> u let Some(set_buffer_expectations) = callbacks.set_buffer_expectations else { return ERROR(ZstdErrorCode::Generic); }; - if fields.callback_context.is_null() { + if fields.callback_context.is_null() || fields.nb_workers.is_null() { return ERROR(ZstdErrorCode::Generic); } if output.pos > output.size { @@ -1183,7 +1178,7 @@ unsafe fn compress_stream2_fallback(state: &ZSTD_rust_compressStream2State) -> u return stability_result; } - if fields.nb_workers > 0 { + if unsafe { *fields.nb_workers } > 0 { if !fields.c_params_changed.is_null() { let c_params_changed = unsafe { &mut *fields.c_params_changed }; if *c_params_changed != 0 { @@ -22928,6 +22923,8 @@ mod tests { struct Stream2FallbackProbe { events: Vec<&'static str>, compress_result: usize, + init_nb_workers: c_int, + nb_workers: c_int, } unsafe extern "C" fn stream2_test_init( @@ -22935,9 +22932,9 @@ mod tests { _end_op: c_int, _input_size: usize, ) -> usize { - unsafe { &mut *context.cast::() } - .events - .push("init"); + let probe = unsafe { &mut *context.cast::() }; + probe.events.push("init"); + probe.nb_workers = probe.init_nb_workers; 0 } @@ -22977,6 +22974,27 @@ mod tests { probe.compress_result } + unsafe extern "C" fn stream2_test_mt_step( + context: *mut c_void, + output: *mut ZSTD_outBuffer, + input: *mut ZSTD_inBuffer, + _end_op: c_int, + ) -> usize { + let probe = unsafe { &mut *context.cast::() }; + probe.events.push("mt"); + unsafe { + (*input).pos = (*input).size; + (*output).pos += 3; + } + 0 + } + + unsafe extern "C" fn stream2_test_mt_trace(_context: *mut c_void) {} + + unsafe extern "C" fn stream2_test_mt_reset(_context: *mut c_void) -> usize { + 0 + } + fn stream2_test_state<'a>( probe: &'a mut Stream2FallbackProbe, output: &'a mut ZSTD_outBuffer, @@ -23003,7 +23021,7 @@ mod tests { in_buffer_mode: ZSTD_BM_BUFFERED, out_buffer_mode: ZSTD_BM_BUFFERED, format: ZSTD_F_ZSTD1, - nb_workers: 0, + nb_workers: &probe.nb_workers, }); let callbacks = Box::new(ZSTD_rust_compressStream2Callbacks { init_compress_stream2: Some(stream2_test_init), @@ -23011,9 +23029,9 @@ mod tests { set_buffer_expectations: Some(stream2_test_set), compress_stream: Some(stream2_test_compress), update_mt_cparams: None, - mt_step: None, - mt_trace: None, - mt_reset: None, + mt_step: Some(stream2_test_mt_step), + mt_trace: Some(stream2_test_mt_trace), + mt_reset: Some(stream2_test_mt_reset), }); ZSTD_rust_compressStream2State { fields: Box::into_raw(fields), @@ -23027,6 +23045,8 @@ mod tests { let mut probe = Stream2FallbackProbe { events: Vec::new(), compress_result: 7, + init_nb_workers: 0, + nb_workers: 0, }; let mut output_bytes = [0u8; 32]; let input_bytes = [1u8; 8]; @@ -23079,6 +23099,8 @@ mod tests { let mut probe = Stream2FallbackProbe { events: Vec::new(), compress_result: 0, + init_nb_workers: 0, + nb_workers: 0, }; let input_bytes = [1u8; 16]; let mut output = ZSTD_outBuffer { @@ -23134,4 +23156,58 @@ mod tests { drop(Box::from_raw(state.callbacks.cast_mut())); } } + + #[test] + fn stream2_fallback_reads_worker_mode_after_initialization() { + let mut probe = Stream2FallbackProbe { + events: Vec::new(), + compress_result: 0, + init_nb_workers: 1, + nb_workers: 0, + }; + let mut output_bytes = [0u8; 32]; + let input_bytes = [1u8; 8]; + let mut output = ZSTD_outBuffer { + dst: output_bytes.as_mut_ptr().cast(), + size: output_bytes.len(), + pos: 0, + }; + let mut input = ZSTD_inBuffer { + src: input_bytes.as_ptr().cast(), + size: input_bytes.len(), + pos: 0, + }; + let mut stage = ZSTD_CSTREAM_STAGE_INIT; + let mut stable_in_not_consumed = 0; + let mut expected_in_buffer = ZSTD_inBuffer { + src: ptr::null(), + size: 0, + pos: 0, + }; + let mut simple_compress2_completed = 0; + let out_buff_content_size = 0; + let out_buff_flushed_size = 0; + let state = stream2_test_state( + &mut probe, + &mut output, + &mut input, + &mut stage, + &mut stable_in_not_consumed, + &mut expected_in_buffer, + &mut simple_compress2_completed, + &out_buff_content_size, + &out_buff_flushed_size, + ); + + let result = unsafe { ZSTD_rust_compressStream2(&state) }; + + assert_eq!(result, 0); + assert_eq!(probe.events, ["init", "check", "mt", "set"]); + assert_eq!(input.pos, input.size); + assert_eq!(output.pos, 3); + unsafe { + drop(Box::from_raw(state.fields.cast_mut())); + drop(Box::from_raw(state.callbacks.cast_mut())); + } + } }