diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index 5f606a8bb..4382b1043 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -66,6 +66,8 @@ int ZSTD_rust_simpleCompress2Level(const void* cctx); int ZSTD_rust_simpleCompressStream2Level(const void* cctx); void ZSTD_rust_reduceTable(U32* table, U32 size, U32 reducerValue, int preserveMark); +void ZSTD_rust_copyCDictTableIntoCCtx(U32* dst, U32 const* src, + size_t tableSize, int tagged); U64 ZSTD_rust_advanceHashSalt(U64 hashSalt, U64 hashSaltEntropy); int ZSTD_rust_indexTooCloseToMax(size_t nextSrcBaseOffset); int ZSTD_rust_dictTooBig(size_t loadedDictSize); @@ -1877,18 +1879,9 @@ ZSTD_resetCCtx_byAttachingCDict(ZSTD_CCtx* cctx, static void ZSTD_copyCDictTableIntoCCtx(U32* dst, U32 const* src, size_t tableSize, ZSTD_compressionParameters const* cParams) { - if (ZSTD_CDictIndicesAreTagged(cParams)){ - /* Remove tags from the CDict table if they are present. - * See docs on "short cache" in zstd_compress_internal.h for context. */ - size_t i; - for (i = 0; i < tableSize; i++) { - U32 const taggedIndex = src[i]; - U32 const index = taggedIndex >> ZSTD_SHORT_CACHE_TAG_BITS; - dst[i] = index; - } - } else { - ZSTD_memcpy(dst, src, tableSize * sizeof(U32)); - } + ZSTD_STATIC_ASSERT(ZSTD_SHORT_CACHE_TAG_BITS == 8); + ZSTD_rust_copyCDictTableIntoCCtx(dst, src, tableSize, + ZSTD_CDictIndicesAreTagged(cParams)); } static size_t ZSTD_resetCCtx_byCopyingCDict(ZSTD_CCtx* cctx, diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs index 7912f918d..49226d72e 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -83,6 +83,7 @@ const ZSTD_ROWSIZE: usize = 16; const ZSTD_WINDOW_START_INDEX: u32 = 2; const ZSTD_DUBT_UNSORTED_MARK: u32 = 1; const ZSTD_INDEXOVERFLOW_MARGIN: usize = 16usize << 20; +const ZSTD_SHORT_CACHE_TAG_BITS: u32 = 8; const ZSTD_CURRENT_MAX: usize = if size_of::() == 8 { 3500usize << 20 } else { @@ -345,6 +346,25 @@ pub unsafe extern "C" fn ZSTD_rust_reduceTable( reduce_table_internal(table, reducer_value, preserve_mark != 0); } +/// Copies a CDict match table into a CCtx, removing short-cache tags when the +/// C-owned compression parameters say the source table is tagged. +#[no_mangle] +pub unsafe extern "C" fn ZSTD_rust_copyCDictTableIntoCCtx( + dst: *mut u32, + src: *const u32, + table_size: usize, + tagged: c_int, +) { + if tagged != 0 { + for i in 0..table_size { + let index = unsafe { *src.add(i) } >> ZSTD_SHORT_CACHE_TAG_BITS; + unsafe { *dst.add(i) = index }; + } + } else if table_size != 0 { + unsafe { ptr::copy_nonoverlapping(src, dst, table_size) }; + } +} + #[inline] fn zeroed_state() -> ZSTD_compressedBlockState_t { /* The state contains only integer arrays and enum fields. */ @@ -856,6 +876,40 @@ mod tests { assert_eq!(table[ZSTD_ROWSIZE * 2 - 1], u32::MAX - 4); } + #[test] + fn copy_cdict_table_removes_short_cache_tags() { + let source = [0x1234_56ff, 0xdead_beef, 0x0000_0100, u32::MAX]; + let mut destination = [0u32; 4]; + + unsafe { + ZSTD_rust_copyCDictTableIntoCCtx( + destination.as_mut_ptr(), + source.as_ptr(), + source.len(), + 1, + ); + } + + assert_eq!(destination, [0x0012_3456, 0x00de_adbe, 1, 0x00ff_ffff]); + } + + #[test] + fn copy_cdict_table_preserves_untagged_entries() { + let source = [0, 2, 0x1234_5678, u32::MAX]; + let mut destination = [0xa5a5_a5a5; 4]; + + unsafe { + ZSTD_rust_copyCDictTableIntoCCtx( + destination.as_mut_ptr(), + source.as_ptr(), + source.len(), + 0, + ); + } + + assert_eq!(destination, source); + } + #[test] fn bitmix_and_hash_salt_match_the_c_arithmetic() { assert_eq!(bitmix(0x0123_4567_89ab_cdef, 8), 0xd498_d855_4e8d_d8cb);