feat(compress): move block size estimation to Rust

Block splitting still estimated literal and sequence section sizes through
three C helpers. Move the pure estimator into zstd_compress_stats.rs so its
histogram, HUF, and FSE cost calculations use the Rust implementations already
behind the compression ABI. Preserve the original literal headers, multi-stream
Huffman surcharge, sequence additional-bit costs, entropy-header accounting,
and fallback estimates; leave C responsible for split recursion and table
ownership.

Test Plan:
- `cargo test --manifest-path rust/Cargo.toml --no-default-features --features compression` -- 229 passed
- `cargo clippy --manifest-path rust/Cargo.toml --no-default-features --features compression` plus benches/tests -- passed before and after formatting
- `cargo +nightly fmt --manifest-path rust/Cargo.toml --all` -- passed
- `make -B -C lib -j2 lib` -- passed
- `make -C tests -j2 test-zstream` -- passed, including 84 deterministic and 15,375 randomized cases

lib/compress/zstd_compress.c

@@ -257,6 +257,14 @@ size_t ZSTD_rust_buildBlockEntropyStats(
         int strategy, int disableLiteralCompression,
         ZSTD_entropyCTablesMetadata_t* entropyMetadata,
         void* workspace, size_t wkspSize);
+size_t ZSTD_rust_estimateBlockSize(
+        const BYTE* literals, size_t litSize,
+        const BYTE* ofCodeTable, const BYTE* llCodeTable,
+        const BYTE* mlCodeTable, size_t nbSeq,
+        const ZSTD_entropyCTables_t* entropy,
+        const ZSTD_entropyCTablesMetadata_t* entropyMetadata,
+        void* workspace, size_t wkspSize,
+        int writeLitEntropy, int writeSeqEntropy);
 size_t ZSTD_rust_copyBlockSequences(
         SeqCollector* seqCollector, const SeqStore_t* seqStore,
         const U32 prevRepcodes[ZSTD_REP_NUM]);
@@ -2600,101 +2608,8 @@ size_t ZSTD_buildBlockEntropyStats(
                 workspace, wkspSize);
 }
 
-/* Returns the size estimate for the literals section (header + content) of a block */
+/* The block-size estimator is implemented in Rust; C retains the stateful
-static size_t
+ * block-splitting recursion and the entropy-table ownership. */
-ZSTD_estimateBlockSize_literal(const BYTE* literals, size_t litSize,
-                               const ZSTD_hufCTables_t* huf,
-                               const ZSTD_hufCTablesMetadata_t* hufMetadata,
-                               void* workspace, size_t wkspSize,
-                               int writeEntropy)
-{
-    unsigned* const countWksp = (unsigned*)workspace;
-    unsigned maxSymbolValue = HUF_SYMBOLVALUE_MAX;
-    size_t literalSectionHeaderSize = 3 + (litSize >= 1 KB) + (litSize >= 16 KB);
-    U32 singleStream = litSize < 256;
-
-    if (hufMetadata->hType == set_basic) return litSize;
-    else if (hufMetadata->hType == set_rle) return 1;
-    else if (hufMetadata->hType == set_compressed || hufMetadata->hType == set_repeat) {
-        size_t const largest = HIST_count_wksp (countWksp, &maxSymbolValue, (const BYTE*)literals, litSize, workspace, wkspSize);
-        if (ZSTD_isError(largest)) return litSize;
-        {   size_t cLitSizeEstimate = HUF_estimateCompressedSize((const HUF_CElt*)huf->CTable, countWksp, maxSymbolValue);
-            if (writeEntropy) cLitSizeEstimate += hufMetadata->hufDesSize;
-            if (!singleStream) cLitSizeEstimate += 6; /* multi-stream huffman uses 6-byte jump table */
-            return cLitSizeEstimate + literalSectionHeaderSize;
-    }   }
-    assert(0); /* impossible */
-    return 0;
-}
-
-/* Returns the size estimate for the FSE-compressed symbols (of, ml, ll) of a block */
-static size_t
-ZSTD_estimateBlockSize_symbolType(SymbolEncodingType_e type,
-                    const BYTE* codeTable, size_t nbSeq, unsigned maxCode,
-                    const FSE_CTable* fseCTable,
-                    const U8* additionalBits,
-                    short const* defaultNorm, U32 defaultNormLog, U32 defaultMax,
-                    void* workspace, size_t wkspSize)
-{
-    unsigned* const countWksp = (unsigned*)workspace;
-    const BYTE* ctp = codeTable;
-    const BYTE* const ctStart = ctp;
-    const BYTE* const ctEnd = ctStart + nbSeq;
-    size_t cSymbolTypeSizeEstimateInBits = 0;
-    unsigned max = maxCode;
-
-    HIST_countFast_wksp(countWksp, &max, codeTable, nbSeq, workspace, wkspSize);  /* can't fail */
-    if (type == set_basic) {
-        /* We selected this encoding type, so it must be valid. */
-        assert(max <= defaultMax);
-        (void)defaultMax;
-        cSymbolTypeSizeEstimateInBits = ZSTD_crossEntropyCost(defaultNorm, defaultNormLog, countWksp, max);
-    } else if (type == set_rle) {
-        cSymbolTypeSizeEstimateInBits = 0;
-    } else if (type == set_compressed || type == set_repeat) {
-        cSymbolTypeSizeEstimateInBits = ZSTD_fseBitCost(fseCTable, countWksp, max);
-    }
-    if (ZSTD_isError(cSymbolTypeSizeEstimateInBits)) {
-        return nbSeq * 10;
-    }
-    while (ctp < ctEnd) {
-        if (additionalBits) cSymbolTypeSizeEstimateInBits += additionalBits[*ctp];
-        else cSymbolTypeSizeEstimateInBits += *ctp; /* for offset, offset code is also the number of additional bits */
-        ctp++;
-    }
-    return cSymbolTypeSizeEstimateInBits >> 3;
-}
-
-/* Returns the size estimate for the sequences section (header + content) of a block */
-static size_t
-ZSTD_estimateBlockSize_sequences(const BYTE* ofCodeTable,
-                                 const BYTE* llCodeTable,
-                                 const BYTE* mlCodeTable,
-                                 size_t nbSeq,
-                                 const ZSTD_fseCTables_t* fseTables,
-                                 const ZSTD_fseCTablesMetadata_t* fseMetadata,
-                                 void* workspace, size_t wkspSize,
-                                 int writeEntropy)
-{
-    size_t sequencesSectionHeaderSize = 1 /* seqHead */ + 1 /* min seqSize size */ + (nbSeq >= 128) + (nbSeq >= LONGNBSEQ);
-    size_t cSeqSizeEstimate = 0;
-    cSeqSizeEstimate += ZSTD_estimateBlockSize_symbolType(fseMetadata->ofType, ofCodeTable, nbSeq, MaxOff,
-                                    fseTables->offcodeCTable, NULL,
-                                    OF_defaultNorm, OF_defaultNormLog, DefaultMaxOff,
-                                    workspace, wkspSize);
-    cSeqSizeEstimate += ZSTD_estimateBlockSize_symbolType(fseMetadata->llType, llCodeTable, nbSeq, MaxLL,
-                                    fseTables->litlengthCTable, LL_bits,
-                                    LL_defaultNorm, LL_defaultNormLog, MaxLL,
-                                    workspace, wkspSize);
-    cSeqSizeEstimate += ZSTD_estimateBlockSize_symbolType(fseMetadata->mlType, mlCodeTable, nbSeq, MaxML,
-                                    fseTables->matchlengthCTable, ML_bits,
-                                    ML_defaultNorm, ML_defaultNormLog, MaxML,
-                                    workspace, wkspSize);
-    if (writeEntropy) cSeqSizeEstimate += fseMetadata->fseTablesSize;
-    return cSeqSizeEstimate + sequencesSectionHeaderSize;
-}
-
-/* Returns the size estimate for a given stream of literals, of, ll, ml */
 static size_t
 ZSTD_estimateBlockSize(const BYTE* literals, size_t litSize,
                        const BYTE* ofCodeTable,
@@ -2706,13 +2621,12 @@ ZSTD_estimateBlockSize(const BYTE* literals, size_t litSize,
                        void* workspace, size_t wkspSize,
                        int writeLitEntropy, int writeSeqEntropy)
 {
-    size_t const literalsSize = ZSTD_estimateBlockSize_literal(literals, litSize,
+    return ZSTD_rust_estimateBlockSize(
-                                    &entropy->huf, &entropyMetadata->hufMetadata,
+            literals, litSize,
-                                    workspace, wkspSize, writeLitEntropy);
+            ofCodeTable, llCodeTable, mlCodeTable, nbSeq,
-    size_t const seqSize = ZSTD_estimateBlockSize_sequences(ofCodeTable, llCodeTable, mlCodeTable,
+            entropy, entropyMetadata,
-                                    nbSeq, &entropy->fse, &entropyMetadata->fseMetadata,
+            workspace, wkspSize,
-                                    workspace, wkspSize, writeSeqEntropy);
+            writeLitEntropy, writeSeqEntropy);
-    return seqSize + literalsSize + ZSTD_blockHeaderSize;
 }
 
 /* Builds entropy statistics and uses them for blocksize estimation.

rust/src/zstd_compress_stats.rs

@@ -9,15 +9,15 @@
 //! `ZSTD_copyBlockSequences()`.  The `ZSTD_CCtx` and `ZSTD_CCtx_params`
 //! layouts stay private to C: the C shims extract the sequence store, the
 //! entropy-table leaves, and the two parameter scalars these paths read
-//! (the compression strategy and the literals-compression switch).  Block
+//! (the compression strategy and the literals-compression switch). Block
-//! dispatch, block splitting, and the block-size estimation heuristics also
+//! dispatch and block-splitting recursion remain in C; the pure block-size
-//! remain in C for a later slice.
+//! estimator is owned here.
 
 use crate::bits::ZSTD_highbit32;
 use crate::common::{
-    DEFAULT_MAX_OFF, LL_DEFAULT_NORM, LL_DEFAULT_NORM_LOG, LL_FSE_LOG, LONGNBSEQ, MAX_LL, MAX_ML,
+    DEFAULT_MAX_OFF, LL_BITS, LL_DEFAULT_NORM, LL_DEFAULT_NORM_LOG, LL_FSE_LOG, LONGNBSEQ, MAX_LL,
-    MAX_OFF, MAX_SEQ, MINMATCH, ML_DEFAULT_NORM, ML_DEFAULT_NORM_LOG, ML_FSE_LOG, OFF_FSE_LOG,
+    MAX_ML, MAX_OFF, MAX_SEQ, MINMATCH, ML_BITS, ML_DEFAULT_NORM, ML_DEFAULT_NORM_LOG, ML_FSE_LOG,
-    OF_DEFAULT_NORM, OF_DEFAULT_NORM_LOG, REP_START_VALUE, ZSTD_MAX_FSE_HEADERS_SIZE,
+    OFF_FSE_LOG, OF_DEFAULT_NORM, OF_DEFAULT_NORM_LOG, REP_START_VALUE, ZSTD_MAX_FSE_HEADERS_SIZE,
     ZSTD_MAX_HUF_HEADER_SIZE, ZSTD_REP_NUM,
 };
 use crate::errors::{ERR_isError, ZstdErrorCode, ERROR};
@@ -31,7 +31,8 @@ use crate::zstd_compress_literals::{
     copy_huf_tables, min_gain, ZSTD_compressLiterals, ZSTD_hufCTables_t,
 };
 use crate::zstd_compress_sequences::{
-    SeqDef, ZSTD_buildCTable, ZSTD_encodeSequences, ZSTD_selectEncodingType,
+    SeqDef, ZSTD_buildCTable, ZSTD_crossEntropyCost, ZSTD_encodeSequences, ZSTD_fseBitCost,
+    ZSTD_selectEncodingType,
 };
 use std::ffi::c_void;
 use std::mem::{size_of, size_of_val};
@@ -246,6 +247,232 @@ pub unsafe extern "C" fn ZSTD_rust_get1BlockSummary(
     invalid()
 }
 
+/// Estimates the literal-section size used by the block splitter.
+unsafe fn estimate_block_size_literal(
+    literals: *const u8,
+    lit_size: usize,
+    huf: *const ZSTD_hufCTables_t,
+    huf_metadata: *const ZSTD_hufCTablesMetadata_t,
+    workspace: *mut c_void,
+    wksp_size: usize,
+    write_entropy: bool,
+) -> usize {
+    let metadata = unsafe { &*huf_metadata };
+    let literal_header_size =
+        3 + usize::from(lit_size >= 1024) + usize::from(lit_size >= 16 * 1024);
+    let single_stream = lit_size < 256;
+
+    match metadata.hType {
+        SET_BASIC => lit_size,
+        SET_RLE => 1,
+        SET_COMPRESSED | SET_REPEAT => {
+            let mut max_symbol_value = HUF_SYMBOLVALUE_MAX;
+            let largest = unsafe {
+                HIST_count_wksp(
+                    workspace.cast::<u32>(),
+                    &mut max_symbol_value,
+                    literals.cast::<c_void>(),
+                    lit_size,
+                    workspace,
+                    wksp_size,
+                )
+            };
+            if ERR_isError(largest) {
+                return lit_size;
+            }
+
+            let mut estimate = unsafe {
+                HUF_estimateCompressedSize(
+                    (*huf).CTable.as_ptr(),
+                    workspace.cast::<u32>(),
+                    max_symbol_value,
+                )
+            };
+            if write_entropy {
+                estimate = estimate.wrapping_add(metadata.hufDesSize);
+            }
+            if !single_stream {
+                estimate = estimate.wrapping_add(6);
+            }
+            estimate.wrapping_add(literal_header_size)
+        }
+        _ => {
+            debug_assert!(false, "invalid literal encoding type");
+            0
+        }
+    }
+}
+
+/// Estimates one FSE symbol stream used by the block splitter.
+#[allow(clippy::too_many_arguments)]
+unsafe fn estimate_block_size_symbol_type(
+    encoding_type: c_int,
+    code_table: *const u8,
+    nb_seq: usize,
+    max_code: u32,
+    fse_ctable: *const u32,
+    additional_bits: *const u8,
+    default_norm: *const i16,
+    default_norm_log: u32,
+    default_max: u32,
+    workspace: *mut c_void,
+    wksp_size: usize,
+) -> usize {
+    let mut max = max_code;
+    unsafe {
+        let _ = HIST_countFast_wksp(
+            workspace.cast::<u32>(),
+            &mut max,
+            code_table.cast::<c_void>(),
+            nb_seq,
+            workspace,
+            wksp_size,
+        );
+    }
+
+    let mut estimate_in_bits = match encoding_type {
+        SET_BASIC => {
+            debug_assert!(max <= default_max);
+            unsafe { ZSTD_crossEntropyCost(default_norm, default_norm_log, workspace.cast(), max) }
+        }
+        SET_RLE => 0,
+        SET_COMPRESSED | SET_REPEAT => unsafe {
+            ZSTD_fseBitCost(fse_ctable, workspace.cast(), max)
+        },
+        _ => 0,
+    };
+    if ERR_isError(estimate_in_bits) {
+        return nb_seq.wrapping_mul(10);
+    }
+
+    for index in 0..nb_seq {
+        let code = unsafe { *code_table.add(index) };
+        let additional = if additional_bits.is_null() {
+            code as usize
+        } else {
+            unsafe { *additional_bits.add(code as usize) as usize }
+        };
+        estimate_in_bits = estimate_in_bits.wrapping_add(additional);
+    }
+    estimate_in_bits >> 3
+}
+
+/// Estimates the sequence-section size used by the block splitter.
+#[allow(clippy::too_many_arguments)]
+unsafe fn estimate_block_size_sequences(
+    of_code_table: *const u8,
+    ll_code_table: *const u8,
+    ml_code_table: *const u8,
+    nb_seq: usize,
+    fse_tables: *const ZSTD_fseCTables_t,
+    fse_metadata: *const ZSTD_fseCTablesMetadata_t,
+    workspace: *mut c_void,
+    wksp_size: usize,
+    write_entropy: bool,
+) -> usize {
+    let metadata = unsafe { &*fse_metadata };
+    let tables = unsafe { &*fse_tables };
+    let header_size =
+        1 + 1 + usize::from(nb_seq >= 128) + usize::from(nb_seq >= LONGNBSEQ as usize);
+    let mut estimate = unsafe {
+        estimate_block_size_symbol_type(
+            metadata.ofType,
+            of_code_table,
+            nb_seq,
+            MAX_OFF as u32,
+            tables.offcodeCTable.as_ptr(),
+            ptr::null(),
+            OF_DEFAULT_NORM.as_ptr(),
+            OF_DEFAULT_NORM_LOG,
+            DEFAULT_MAX_OFF as u32,
+            workspace,
+            wksp_size,
+        )
+    };
+    estimate = estimate.wrapping_add(unsafe {
+        estimate_block_size_symbol_type(
+            metadata.llType,
+            ll_code_table,
+            nb_seq,
+            MAX_LL as u32,
+            tables.litlengthCTable.as_ptr(),
+            LL_BITS.as_ptr(),
+            LL_DEFAULT_NORM.as_ptr(),
+            LL_DEFAULT_NORM_LOG,
+            MAX_LL as u32,
+            workspace,
+            wksp_size,
+        )
+    });
+    estimate = estimate.wrapping_add(unsafe {
+        estimate_block_size_symbol_type(
+            metadata.mlType,
+            ml_code_table,
+            nb_seq,
+            MAX_ML as u32,
+            tables.matchlengthCTable.as_ptr(),
+            ML_BITS.as_ptr(),
+            ML_DEFAULT_NORM.as_ptr(),
+            ML_DEFAULT_NORM_LOG,
+            MAX_ML as u32,
+            workspace,
+            wksp_size,
+        )
+    });
+    if write_entropy {
+        estimate = estimate.wrapping_add(metadata.fseTablesSize);
+    }
+    estimate.wrapping_add(header_size)
+}
+
+/// Estimates a block's compressed size for block-split decisions.  This is
+/// the Rust leaf for C's `ZSTD_estimateBlockSize()` and keeps the workspace and
+/// entropy-table representations on the existing C ABI.
+#[allow(clippy::too_many_arguments)]
+#[no_mangle]
+pub unsafe extern "C" fn ZSTD_rust_estimateBlockSize(
+    literals: *const u8,
+    lit_size: usize,
+    of_code_table: *const u8,
+    ll_code_table: *const u8,
+    ml_code_table: *const u8,
+    nb_seq: usize,
+    entropy: *const ZSTD_entropyCTables_t,
+    entropy_metadata: *const ZSTD_entropyCTablesMetadata_t,
+    workspace: *mut c_void,
+    wksp_size: usize,
+    write_lit_entropy: c_int,
+    write_seq_entropy: c_int,
+) -> usize {
+    let entropy = unsafe { &*entropy };
+    let entropy_metadata = unsafe { &*entropy_metadata };
+    let literals_size = unsafe {
+        estimate_block_size_literal(
+            literals,
+            lit_size,
+            ptr::addr_of!(entropy.huf),
+            ptr::addr_of!(entropy_metadata.hufMetadata),
+            workspace,
+            wksp_size,
+            write_lit_entropy != 0,
+        )
+    };
+    let sequence_size = unsafe {
+        estimate_block_size_sequences(
+            of_code_table,
+            ll_code_table,
+            ml_code_table,
+            nb_seq,
+            ptr::addr_of!(entropy.fse),
+            ptr::addr_of!(entropy_metadata.fseMetadata),
+            workspace,
+            wksp_size,
+            write_seq_entropy != 0,
+        )
+    };
+    sequence_size.wrapping_add(literals_size).wrapping_add(3)
+}
+
 /// Converts a raw sequence offset to the stored offBase representation.
 ///
 /// This is the Rust leaf for C's `ZSTD_finalizeOffBase()`.  The repcode
@@ -1994,6 +2221,97 @@ mod tests {
         );
     }
 
+    #[test]
+    fn estimate_block_size_counts_basic_literals_and_rle_sequences() {
+        let entropy = ZSTD_entropyCTables_t {
+            huf: empty_huf_tables(HUF_REPEAT_NONE),
+            fse: empty_fse_tables(FSE_REPEAT_NONE),
+        };
+        let mut metadata = unsafe {
+            std::mem::MaybeUninit::<ZSTD_entropyCTablesMetadata_t>::zeroed().assume_init()
+        };
+        metadata.hufMetadata.hType = SET_BASIC;
+        metadata.fseMetadata.ofType = SET_RLE;
+        metadata.fseMetadata.llType = SET_RLE;
+        metadata.fseMetadata.mlType = SET_RLE;
+        metadata.fseMetadata.fseTablesSize = 7;
+        let literals = *b"abcd";
+        let codes = [0u8];
+        let mut workspace = vec![0u64; ENTROPY_WORKSPACE_SIZE / size_of::<u64>()];
+
+        let without_entropy = unsafe {
+            ZSTD_rust_estimateBlockSize(
+                literals.as_ptr(),
+                literals.len(),
+                codes.as_ptr(),
+                codes.as_ptr(),
+                codes.as_ptr(),
+                1,
+                &entropy,
+                &metadata,
+                workspace.as_mut_ptr().cast(),
+                ENTROPY_WORKSPACE_SIZE,
+                0,
+                0,
+            )
+        };
+        assert_eq!(without_entropy, 9);
+
+        let with_entropy = unsafe {
+            ZSTD_rust_estimateBlockSize(
+                literals.as_ptr(),
+                literals.len(),
+                codes.as_ptr(),
+                codes.as_ptr(),
+                codes.as_ptr(),
+                1,
+                &entropy,
+                &metadata,
+                workspace.as_mut_ptr().cast(),
+                ENTROPY_WORKSPACE_SIZE,
+                0,
+                1,
+            )
+        };
+        assert_eq!(with_entropy, 16);
+    }
+
+    #[test]
+    fn estimate_block_size_counts_rle_literals() {
+        let entropy = ZSTD_entropyCTables_t {
+            huf: empty_huf_tables(HUF_REPEAT_NONE),
+            fse: empty_fse_tables(FSE_REPEAT_NONE),
+        };
+        let mut metadata = unsafe {
+            std::mem::MaybeUninit::<ZSTD_entropyCTablesMetadata_t>::zeroed().assume_init()
+        };
+        metadata.hufMetadata.hType = SET_RLE;
+        metadata.fseMetadata.ofType = SET_RLE;
+        metadata.fseMetadata.llType = SET_RLE;
+        metadata.fseMetadata.mlType = SET_RLE;
+        let literals = [b'x'; 1024];
+        let codes = [0u8];
+        let mut workspace = vec![0u64; ENTROPY_WORKSPACE_SIZE / size_of::<u64>()];
+
+        let estimate = unsafe {
+            ZSTD_rust_estimateBlockSize(
+                literals.as_ptr(),
+                literals.len(),
+                codes.as_ptr(),
+                codes.as_ptr(),
+                codes.as_ptr(),
+                1,
+                &entropy,
+                &metadata,
+                workspace.as_mut_ptr().cast(),
+                ENTROPY_WORKSPACE_SIZE,
+                0,
+                0,
+            )
+        };
+        assert_eq!(estimate, 6);
+    }
+
     #[test]
     fn compressed_block_state_reset_restores_repcodes_and_repeat_modes() {
         let mut block_state =