diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index e01afae33..bbe0e6d5f 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -66,6 +66,16 @@ void ZSTD_rust_reduceTable(U32* table, U32 size, U32 reducerValue, U64 ZSTD_rust_advanceHashSalt(U64 hashSalt, U64 hashSaltEntropy); int ZSTD_rust_indexTooCloseToMax(size_t nextSrcBaseOffset); int ZSTD_rust_dictTooBig(size_t loadedDictSize); +enum { + ZSTD_RUST_TARGET_CBLOCK_ACTION_RAW = 0, + ZSTD_RUST_TARGET_CBLOCK_ACTION_RLE = 1, + ZSTD_RUST_TARGET_CBLOCK_ACTION_COMPRESSED = 2, + ZSTD_RUST_TARGET_CBLOCK_ACTION_ERROR = 3 +}; +int ZSTD_rust_targetCBlockSizeAction(int bss, int isFirstBlock, + int maybeRLE, int isRLE, + size_t cSize, size_t srcSize, + int strategy); /* Context-free compression-parameter selection and sizing leaves live in * Rust (rust/src/zstd_compress_params.rs). This file retains @@ -3124,18 +3134,25 @@ static size_t ZSTD_compressBlock_targetCBlockSize_body(ZSTD_CCtx* zc, const void* src, size_t srcSize, const size_t bss, U32 lastBlock) { + int const isCompress = bss == ZSTDbss_compress; + int maybeRLE = 0; + int isRLE = 0; + int action; + DEBUGLOG(6, "Attempting ZSTD_compressSuperBlock()"); - if (bss == ZSTDbss_compress) { - if (/* We don't want to emit our first block as a RLE even if it qualifies because - * doing so will cause the decoder (cli only) to throw a "should consume all input error." - * This is only an issue for zstd <= v1.4.3 - */ - !zc->isFirstBlock && - ZSTD_maybeRLE(&zc->seqStore) && - ZSTD_isRLE((BYTE const*)src, srcSize)) - { - return ZSTD_rust_rleCompressBlock(dst, dstCapacity, *(BYTE const*)src, srcSize, lastBlock); - } + if (isCompress) { + maybeRLE = ZSTD_maybeRLE(&zc->seqStore); + isRLE = ZSTD_isRLE((BYTE const*)src, srcSize); + } + + action = ZSTD_rust_targetCBlockSizeAction( + (int)bss, (int)zc->isFirstBlock, maybeRLE, isRLE, + 0, srcSize, (int)zc->appliedParams.cParams.strategy); + if (action == ZSTD_RUST_TARGET_CBLOCK_ACTION_RLE) { + return ZSTD_rust_rleCompressBlock(dst, dstCapacity, *(BYTE const*)src, srcSize, lastBlock); + } + + if (isCompress) { /* Attempt superblock compression. * * Note that compressed size of ZSTD_compressSuperBlock() is not bound by the @@ -3156,14 +3173,15 @@ static size_t ZSTD_compressBlock_targetCBlockSize_body(ZSTD_CCtx* zc, */ { size_t const cSize = ZSTD_compressSuperBlock(zc, dst, dstCapacity, src, srcSize, lastBlock); - if (cSize != ERROR(dstSize_tooSmall)) { - size_t const maxCSize = - srcSize - ZSTD_minGain(srcSize, zc->appliedParams.cParams.strategy); + action = ZSTD_rust_targetCBlockSizeAction( + (int)bss, (int)zc->isFirstBlock, 0, 0, + cSize, srcSize, (int)zc->appliedParams.cParams.strategy); + if (action == ZSTD_RUST_TARGET_CBLOCK_ACTION_ERROR) { FORWARD_IF_ERROR(cSize, "ZSTD_compressSuperBlock failed"); - if (cSize != 0 && cSize < maxCSize + ZSTD_blockHeaderSize) { - ZSTD_blockState_confirmRepcodesAndEntropyTables(&zc->blockState); - return cSize; - } + } + if (action == ZSTD_RUST_TARGET_CBLOCK_ACTION_COMPRESSED) { + ZSTD_blockState_confirmRepcodesAndEntropyTables(&zc->blockState); + return cSize; } } } /* if (bss == ZSTDbss_compress)*/ diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs index c9f46f571..a2226214f 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -19,6 +19,7 @@ use crate::zstd_compress_api::ZSTD_compressBound; use crate::zstd_compress_frame::{ write_raw_block, ZSTD_rust_writeFrameHeader, ZSTD_writeLastEmptyBlock, }; +use crate::zstd_compress_literals::min_gain; use crate::zstd_compress_params::{ ZSTD_rust_params_adjustCParams, ZSTD_rust_params_maxNbSeq, ZSTD_rust_params_selectCParams, ZSTD_RUST_CPM_NO_ATTACH_DICT, ZSTD_RUST_PS_DISABLE, @@ -62,6 +63,8 @@ const ZSTD_FAST: c_int = 1; const ZSTD_DFAST: c_int = 2; const ZSTD_BLOCKSIZE_MAX: usize = 1 << 17; const ZSTD_CONTENTSIZE_UNKNOWN: u64 = u64::MAX; +const ZSTD_TARGET_CBLOCK_BSS_COMPRESS: c_int = 0; +const ZSTD_BLOCK_HEADER_SIZE: usize = 3; const ZSTD_ROWSIZE: usize = 16; const ZSTD_WINDOW_START_INDEX: u32 = 2; const ZSTD_DUBT_UNSORTED_MARK: u32 = 1; @@ -75,6 +78,71 @@ const ZSTD_CHUNKSIZE_MAX: usize = u32::MAX as usize - ZSTD_CURRENT_MAX; #[cfg(not(test))] const ZSTD_E_END: c_int = 2; +#[repr(i32)] +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +enum TargetCBlockAction { + Raw = 0, + Rle = 1, + Compressed = 2, + Error = 3, +} + +#[inline] +fn target_c_block_size_action( + bss: c_int, + is_first_block: c_int, + maybe_rle: c_int, + is_rle: c_int, + c_size: usize, + src_size: usize, + strategy: c_int, +) -> TargetCBlockAction { + if bss != ZSTD_TARGET_CBLOCK_BSS_COMPRESS { + return TargetCBlockAction::Raw; + } + if is_first_block == 0 && maybe_rle != 0 && is_rle != 0 { + return TargetCBlockAction::Rle; + } + if c_size == 0 || c_size == ERROR(ZstdErrorCode::DstSizeTooSmall) { + return TargetCBlockAction::Raw; + } + if ERR_isError(c_size) { + return TargetCBlockAction::Error; + } + + let max_c_size = src_size.wrapping_sub(min_gain(src_size, strategy)); + if c_size < max_c_size.wrapping_add(ZSTD_BLOCK_HEADER_SIZE) { + TargetCBlockAction::Compressed + } else { + TargetCBlockAction::Raw + } +} + +/// Classify the target-sized block policy without crossing the C context ABI. +/// +/// C calls this once before the superblock attempt to classify the RLE +/// precondition, then again with the superblock result and zero RLE flags. +#[no_mangle] +pub extern "C" fn ZSTD_rust_targetCBlockSizeAction( + bss: c_int, + is_first_block: c_int, + maybe_rle: c_int, + is_rle: c_int, + c_size: usize, + src_size: usize, + strategy: c_int, +) -> c_int { + target_c_block_size_action( + bss, + is_first_block, + maybe_rle, + is_rle, + c_size, + src_size, + strategy, + ) as c_int +} + #[cfg(not(test))] #[repr(C)] pub struct ZSTD_inBuffer { @@ -752,6 +820,111 @@ mod tests { assert_eq!(ZSTD_rust_dictTooBig(ZSTD_CHUNKSIZE_MAX + 1), 1); } + #[test] + fn target_block_policy_requires_compression_and_nonfirst_rle() { + assert_eq!( + target_c_block_size_action(0, 0, 1, 1, 0, 128, ZSTD_FAST), + TargetCBlockAction::Rle + ); + assert_eq!( + target_c_block_size_action(0, 1, 1, 1, 0, 128, ZSTD_FAST), + TargetCBlockAction::Raw + ); + assert_eq!( + target_c_block_size_action(0, 0, 0, 1, 0, 128, ZSTD_FAST), + TargetCBlockAction::Raw + ); + assert_eq!( + target_c_block_size_action(0, 0, 1, 0, 0, 128, ZSTD_FAST), + TargetCBlockAction::Raw + ); + assert_eq!( + ZSTD_rust_targetCBlockSizeAction(0, 0, 1, 1, 0, 128, ZSTD_FAST), + TargetCBlockAction::Rle as c_int + ); + } + + #[test] + fn target_block_policy_falls_back_for_uncompressed_and_small_results() { + assert_eq!( + target_c_block_size_action(1, 0, 0, 0, 0, 128, ZSTD_FAST), + TargetCBlockAction::Raw + ); + assert_eq!( + target_c_block_size_action( + ZSTD_TARGET_CBLOCK_BSS_COMPRESS, + 0, + 0, + 0, + ERROR(ZstdErrorCode::DstSizeTooSmall), + 128, + ZSTD_FAST, + ), + TargetCBlockAction::Raw + ); + assert_eq!( + target_c_block_size_action( + ZSTD_TARGET_CBLOCK_BSS_COMPRESS, + 0, + 0, + 0, + ERROR(ZstdErrorCode::Generic), + 128, + ZSTD_FAST, + ), + TargetCBlockAction::Error + ); + assert_eq!( + target_c_block_size_action(ZSTD_TARGET_CBLOCK_BSS_COMPRESS, 0, 0, 0, 0, 128, ZSTD_FAST,), + TargetCBlockAction::Raw + ); + } + + #[test] + fn target_block_policy_uses_a_strict_three_byte_header_boundary() { + let src_size = 128; + let max_c_size = src_size - min_gain(src_size, ZSTD_FAST); + let compressed = max_c_size + ZSTD_BLOCK_HEADER_SIZE - 1; + let raw = max_c_size + ZSTD_BLOCK_HEADER_SIZE; + + assert_eq!( + target_c_block_size_action( + ZSTD_TARGET_CBLOCK_BSS_COMPRESS, + 0, + 0, + 0, + compressed, + src_size, + ZSTD_FAST, + ), + TargetCBlockAction::Compressed + ); + assert_eq!( + target_c_block_size_action( + ZSTD_TARGET_CBLOCK_BSS_COMPRESS, + 0, + 0, + 0, + raw, + src_size, + ZSTD_FAST, + ), + TargetCBlockAction::Raw + ); + assert_eq!( + ZSTD_rust_targetCBlockSizeAction( + ZSTD_TARGET_CBLOCK_BSS_COMPRESS, + 0, + 0, + 0, + compressed, + src_size, + ZSTD_FAST, + ), + TargetCBlockAction::Compressed as c_int + ); + } + #[test] fn frame_progress_unknown_pledge_updates_counters() { let mut consumed = 7;