diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index d187d709e..46bce0e88 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -1390,6 +1390,30 @@ typedef char ZSTD_rust_external_sequence_producer_state_layout[ == 13 * sizeof(void*)) ? 1 : -1]; +typedef struct { + rawSeq** seq; + size_t* pos; + size_t* posInSequence; + size_t* size; + size_t* capacity; +} ZSTD_rust_externalSequenceStoreState; +void ZSTD_rust_referenceExternalSequences( + const ZSTD_rust_externalSequenceStoreState* state, + rawSeq* seq, size_t nbSeq); +typedef char ZSTD_rust_external_sequence_store_state_layout[ + (offsetof(ZSTD_rust_externalSequenceStoreState, seq) == 0 + && offsetof(ZSTD_rust_externalSequenceStoreState, pos) + == sizeof(void*) + && offsetof(ZSTD_rust_externalSequenceStoreState, posInSequence) + == 2 * sizeof(void*) + && offsetof(ZSTD_rust_externalSequenceStoreState, size) + == 3 * sizeof(void*) + && offsetof(ZSTD_rust_externalSequenceStoreState, capacity) + == 4 * sizeof(void*) + && sizeof(ZSTD_rust_externalSequenceStoreState) + == 5 * sizeof(void*)) + ? 1 : -1]; + /* The sequence-compression loop receives only the state it actually reads or * updates. In particular, neither ZSTD_CCtx nor a C function pointer crosses * the Rust ABI. */ @@ -4066,11 +4090,15 @@ void ZSTD_referenceExternalSequences(ZSTD_CCtx* cctx, rawSeq* seq, size_t nbSeq) { assert(cctx->stage == ZSTDcs_init); assert(nbSeq == 0 || cctx->appliedParams.ldmParams.enableLdm != ZSTD_ps_enable); - cctx->externSeqStore.seq = seq; - cctx->externSeqStore.size = nbSeq; - cctx->externSeqStore.capacity = nbSeq; - cctx->externSeqStore.pos = 0; - cctx->externSeqStore.posInSequence = 0; + { + ZSTD_rust_externalSequenceStoreState state; + state.seq = &cctx->externSeqStore.seq; + state.pos = &cctx->externSeqStore.pos; + state.posInSequence = &cctx->externSeqStore.posInSequence; + state.size = &cctx->externSeqStore.size; + state.capacity = &cctx->externSeqStore.capacity; + ZSTD_rust_referenceExternalSequences(&state, seq, nbSeq); + } } diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs index 27e3bb516..fe7405ca1 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -403,6 +403,55 @@ pub unsafe extern "C" fn ZSTD_rust_tryExternalSequenceProducer( 0 } +/// Reset the externally referenced raw-sequence store without exposing its C +/// representation to Rust. +#[repr(C)] +pub struct ZSTD_rust_externalSequenceStoreState { + seq: *mut *mut c_void, + pos: *mut usize, + pos_in_sequence: *mut usize, + size: *mut usize, + capacity: *mut usize, +} + +const _: () = { + assert!(offset_of!(ZSTD_rust_externalSequenceStoreState, seq) == 0); + assert!(offset_of!(ZSTD_rust_externalSequenceStoreState, pos) == size_of::()); + assert!( + offset_of!(ZSTD_rust_externalSequenceStoreState, pos_in_sequence) == 2 * size_of::() + ); + assert!(offset_of!(ZSTD_rust_externalSequenceStoreState, size) == 3 * size_of::()); + assert!(offset_of!(ZSTD_rust_externalSequenceStoreState, capacity) == 4 * size_of::()); + assert!(size_of::() == 5 * size_of::()); +}; + +#[no_mangle] +pub unsafe extern "C" fn ZSTD_rust_referenceExternalSequences( + state: *const ZSTD_rust_externalSequenceStoreState, + seq: *mut c_void, + nb_seq: usize, +) { + let Some(state) = state.as_ref() else { + return; + }; + if state.seq.is_null() + || state.pos.is_null() + || state.pos_in_sequence.is_null() + || state.size.is_null() + || state.capacity.is_null() + { + return; + } + + unsafe { + *state.seq = seq; + *state.size = nb_seq; + *state.capacity = nb_seq; + *state.pos = 0; + *state.pos_in_sequence = 0; + } +} + /// Explicit projection of the state used by `ZSTD_compress_frameChunk`. /// /// The Rust side owns the per-frame block loop and its savings/dispatch @@ -6214,6 +6263,54 @@ mod tests { const ZSTD_BTOPT: c_int = 7; const ZSTD_BTULTRA2: c_int = 9; + #[test] + fn reference_external_sequences_resets_store_state() { + let mut source = [0u8; 3]; + let sequence = source.as_mut_ptr().cast::(); + let mut stored_sequence = ptr::null_mut(); + let mut pos = 7; + let mut pos_in_sequence = 8; + let mut size = 9; + let mut capacity = 10; + let state = ZSTD_rust_externalSequenceStoreState { + seq: &mut stored_sequence, + pos: &mut pos, + pos_in_sequence: &mut pos_in_sequence, + size: &mut size, + capacity: &mut capacity, + }; + + unsafe { ZSTD_rust_referenceExternalSequences(&state, sequence, 11) }; + + assert_eq!(stored_sequence, sequence); + assert_eq!(pos, 0); + assert_eq!(pos_in_sequence, 0); + assert_eq!(size, 11); + assert_eq!(capacity, 11); + } + + #[test] + fn reference_external_sequences_rejects_incomplete_state() { + let mut pos = 7; + let mut pos_in_sequence = 8; + let mut size = 9; + let mut capacity = 10; + let state = ZSTD_rust_externalSequenceStoreState { + seq: ptr::null_mut(), + pos: &mut pos, + pos_in_sequence: &mut pos_in_sequence, + size: &mut size, + capacity: &mut capacity, + }; + + unsafe { ZSTD_rust_referenceExternalSequences(&state, ptr::null_mut(), 11) }; + + assert_eq!(pos, 7); + assert_eq!(pos_in_sequence, 8); + assert_eq!(size, 9); + assert_eq!(capacity, 10); + } + #[derive(Default)] struct ExternalSequenceProducerProbe { events: Vec<&'static str>,