ddidderr/zstd-rs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
0caec24fd43eca6af5bb12152387e242c0fedc82
zstd-rs/rust/src/entropy_common.rs
T
ddidderr a0f2b2a14c feat(rust): port FSE entropy decoding 
Move FSE normalized-count parsing, Huffman statistics decoding, FSE table
construction, and FSE stream decompression into Rust. The C source files now
only retain the headers needed by the existing build configuration.

The implementation keeps the public C ABI and verifies C-generated balanced,
skewed, and Huffman-statistics streams. Compression and the higher-level frame
decoder remain C for now.

Test Plan:
- cargo fmt --check
- cargo test --all-targets
- cargo clippy --all-targets -- -D warnings
- cargo build --release
- compile both C shims with -Werror and -Wredundant-decls

Refs: rust/README.md
2026-07-10 20:04:26 +02:00

598 lines
18 KiB
Rust
Raw Blame History

#![allow(non_snake_case)]
use crate::bits::{ZSTD_countTrailingZeros32, ZSTD_highbit32};
use crate::errors::{ERR_getErrorName, ERR_isError, ZstdErrorCode, ERROR};
use crate::mem::{MEM_readLE32, BYTE, U32};
use std::os::raw::{c_char, c_int, c_short, c_uint, c_void};
pub const FSE_VERSION_MAJOR: u32 = 0;
pub const FSE_VERSION_MINOR: u32 = 9;
pub const FSE_VERSION_RELEASE: u32 = 0;
pub const FSE_VERSION_NUMBER: u32 =
FSE_VERSION_MAJOR * 100 * 100 + FSE_VERSION_MINOR * 100 + FSE_VERSION_RELEASE;
pub const FSE_MIN_TABLELOG: u32 = 5;
pub const FSE_TABLELOG_ABSOLUTE_MAX: u32 = 15;
pub const HUF_TABLELOG_MAX: u32 = 12;
pub const HUF_FLAGS_BMI2: c_int = 1 << 0;
// FSE_DECOMPRESS_WKSP_SIZE_U32(6, HUF_TABLELOG_MAX-1) from fse.h / huf.h
const HUF_READ_STATS_WORKSPACE_SIZE_U32: usize = 219;
#[no_mangle]
pub extern "C" fn FSE_versionNumber() -> c_uint {
FSE_VERSION_NUMBER
}
#[no_mangle]
pub extern "C" fn FSE_isError(code: usize) -> c_uint {
ERR_isError(code) as c_uint
}
#[no_mangle]
pub extern "C" fn FSE_getErrorName(code: usize) -> *const c_char {
ERR_getErrorName(code)
}
#[no_mangle]
pub extern "C" fn HUF_isError(code: usize) -> c_uint {
ERR_isError(code) as c_uint
}
#[no_mangle]
pub extern "C" fn HUF_getErrorName(code: usize) -> *const c_char {
ERR_getErrorName(code)
}
fn fse_read_ncount_body(
normalized_counter: *mut c_short,
max_sv_ptr: *mut c_uint,
table_log_ptr: *mut c_uint,
header_buffer: *const c_void,
hb_size: usize,
) -> usize {
unsafe {
if hb_size < 8 {
let mut buffer = [0u8; 8];
if hb_size > 0 {
std::ptr::copy_nonoverlapping(
header_buffer as *const u8,
buffer.as_mut_ptr(),
hb_size,
);
}
let count_size = FSE_readNCount(
normalized_counter,
max_sv_ptr,
table_log_ptr,
buffer.as_ptr() as *const c_void,
buffer.len(),
);
if ERR_isError(count_size) {
return count_size;
}
if count_size > hb_size {
return ERROR(ZstdErrorCode::CorruptionDetected);
}
return count_size;
}
let istart = header_buffer as *const BYTE;
let iend = istart.add(hb_size);
let mut ip = istart;
let max_sv1 = *max_sv_ptr + 1;
// Zero frequency table for all symbols up to maxSVPtr.
std::ptr::write_bytes(normalized_counter, 0, (*max_sv_ptr as usize) + 1);
let mut bit_stream = MEM_readLE32(ip as *const c_void);
let mut nb_bits = ((bit_stream & 0xF) + FSE_MIN_TABLELOG) as i32;
if nb_bits as u32 > FSE_TABLELOG_ABSOLUTE_MAX {
return ERROR(ZstdErrorCode::TableLogTooLarge);
}
bit_stream >>= 4;
let mut bit_count: i32 = 4;
*table_log_ptr = nb_bits as c_uint;
let mut remaining = (1 << nb_bits) + 1;
let mut threshold = 1 << nb_bits;
nb_bits += 1;
let mut charnum: u32 = 0;
let mut previous0 = false;
loop {
if previous0 {
let mut repeats =
(ZSTD_countTrailingZeros32((!bit_stream) | 0x8000_0000) >> 1) as i32;
while repeats >= 12 {
charnum += 3 * 12;
if ip <= iend.sub(7) {
ip = ip.add(3);
} else {
bit_count -= 8 * (iend.offset_from(ip) as i32 - 7);
bit_count &= 31;
ip = iend.sub(4);
}
bit_stream = MEM_readLE32(ip as *const c_void) >> bit_count;
repeats = (ZSTD_countTrailingZeros32((!bit_stream) | 0x8000_0000) >> 1) as i32;
}
charnum += 3 * repeats as u32;
bit_stream >>= 2 * repeats;
bit_count += 2 * repeats;
debug_assert!((bit_stream & 3) < 3);
charnum += bit_stream & 3;
bit_count += 2;
if charnum >= max_sv1 {
break;
}
if ip <= iend.sub(7) || ip.add((bit_count >> 3) as usize) <= iend.sub(4) {
debug_assert!((bit_count >> 3) <= 3);
ip = ip.add((bit_count >> 3) as usize);
bit_count &= 7;
} else {
bit_count -= 8 * (iend.offset_from(ip) as i32 - 4);
bit_count &= 31;
ip = iend.sub(4);
}
bit_stream = MEM_readLE32(ip as *const c_void) >> bit_count;
}
{
let max = (2 * threshold - 1) - remaining;
let mut count: i32;
if (bit_stream & ((threshold as u32) - 1)) < max as u32 {
count = (bit_stream & ((threshold as u32) - 1)) as i32;
bit_count += nb_bits - 1;
} else {
count = (bit_stream & ((2 * threshold as u32) - 1)) as i32;
if count >= threshold {
count -= max;
}
bit_count += nb_bits;
}
count -= 1;
if count >= 0 {
remaining -= count;
} else {
debug_assert!(count == -1);
remaining += count;
}
*normalized_counter.add(charnum as usize) = count as c_short;
charnum += 1;
previous0 = count == 0;
debug_assert!(threshold > 1);
if remaining < threshold {
if remaining <= 1 {
break;
}
nb_bits = ZSTD_highbit32(remaining as u32) as i32 + 1;
threshold = 1 << (nb_bits - 1);
}
if charnum >= max_sv1 {
break;
}
if ip <= iend.sub(7) || ip.add((bit_count >> 3) as usize) <= iend.sub(4) {
ip = ip.add((bit_count >> 3) as usize);
bit_count &= 7;
} else {
bit_count -= 8 * (iend.offset_from(ip) as i32 - 4);
bit_count &= 31;
ip = iend.sub(4);
}
bit_stream = MEM_readLE32(ip as *const c_void) >> bit_count;
}
}
if remaining != 1 {
return ERROR(ZstdErrorCode::CorruptionDetected);
}
if charnum > max_sv1 {
return ERROR(ZstdErrorCode::MaxSymbolValueTooSmall);
}
if bit_count > 32 {
return ERROR(ZstdErrorCode::CorruptionDetected);
}
*max_sv_ptr = charnum - 1;
ip = ip.add(((bit_count + 7) >> 3) as usize);
ip.offset_from(istart) as usize
}
}
#[no_mangle]
pub unsafe extern "C" fn FSE_readNCount_bmi2(
normalized_counter: *mut c_short,
max_sv_ptr: *mut c_uint,
table_log_ptr: *mut c_uint,
header_buffer: *const c_void,
hb_size: usize,
_bmi2: c_int,
) -> usize {
fse_read_ncount_body(
normalized_counter,
max_sv_ptr,
table_log_ptr,
header_buffer,
hb_size,
)
}
#[no_mangle]
pub unsafe extern "C" fn FSE_readNCount(
normalized_counter: *mut c_short,
max_sv_ptr: *mut c_uint,
table_log_ptr: *mut c_uint,
header_buffer: *const c_void,
hb_size: usize,
) -> usize {
FSE_readNCount_bmi2(
normalized_counter,
max_sv_ptr,
table_log_ptr,
header_buffer,
hb_size,
0,
)
}
#[allow(clippy::too_many_arguments)]
fn huf_read_stats_body(
huff_weight: *mut BYTE,
hw_size: usize,
rank_stats: *mut U32,
nb_symbols_ptr: *mut U32,
table_log_ptr: *mut U32,
src: *const c_void,
src_size: usize,
work_space: *mut c_void,
wksp_size: usize,
bmi2: c_int,
) -> usize {
unsafe {
if src_size == 0 {
return ERROR(ZstdErrorCode::SrcSizeWrong);
}
let mut ip = src as *const BYTE;
let mut i_size = *ip as usize;
let o_size: usize;
if i_size >= 128 {
o_size = i_size - 127;
i_size = o_size.div_ceil(2);
if i_size + 1 > src_size {
return ERROR(ZstdErrorCode::SrcSizeWrong);
}
if o_size >= hw_size {
return ERROR(ZstdErrorCode::CorruptionDetected);
}
ip = ip.add(1);
let mut n = 0usize;
while n < o_size {
*huff_weight.add(n) = *ip.add(n / 2) >> 4;
*huff_weight.add(n + 1) = *ip.add(n / 2) & 15;
n += 2;
}
} else {
if i_size + 1 > src_size {
return ERROR(ZstdErrorCode::SrcSizeWrong);
}
let dec = crate::fse_decompress::FSE_decompress_wksp_bmi2(
huff_weight as *mut c_void,
hw_size - 1,
ip.add(1) as *const c_void,
i_size,
6,
work_space,
wksp_size,
bmi2,
);
if ERR_isError(dec) {
return dec;
}
o_size = dec;
}
std::ptr::write_bytes(rank_stats, 0, (HUF_TABLELOG_MAX as usize) + 1);
let mut weight_total: U32 = 0;
for n in 0..o_size {
let w = *huff_weight.add(n) as U32;
if w > HUF_TABLELOG_MAX {
return ERROR(ZstdErrorCode::CorruptionDetected);
}
*rank_stats.add(w as usize) += 1;
weight_total += (1u32 << w) >> 1;
}
if weight_total == 0 {
return ERROR(ZstdErrorCode::CorruptionDetected);
}
let table_log = ZSTD_highbit32(weight_total) + 1;
if table_log > HUF_TABLELOG_MAX {
return ERROR(ZstdErrorCode::CorruptionDetected);
}
*table_log_ptr = table_log;
{
let total = 1u32 << table_log;
let rest = total - weight_total;
let verif = 1u32 << ZSTD_highbit32(rest);
let last_weight = ZSTD_highbit32(rest) + 1;
if verif != rest {
return ERROR(ZstdErrorCode::CorruptionDetected);
}
*huff_weight.add(o_size) = last_weight as BYTE;
*rank_stats.add(last_weight as usize) += 1;
}
let r1 = *rank_stats.add(1);
if r1 < 2 || (r1 & 1) != 0 {
return ERROR(ZstdErrorCode::CorruptionDetected);
}
*nb_symbols_ptr = (o_size + 1) as U32;
i_size + 1
}
}
#[no_mangle]
pub unsafe extern "C" fn HUF_readStats_wksp(
huff_weight: *mut BYTE,
hw_size: usize,
rank_stats: *mut U32,
nb_symbols_ptr: *mut U32,
table_log_ptr: *mut U32,
src: *const c_void,
src_size: usize,
work_space: *mut c_void,
wksp_size: usize,
flags: c_int,
) -> usize {
let bmi2 = if (flags & HUF_FLAGS_BMI2) != 0 { 1 } else { 0 };
huf_read_stats_body(
huff_weight,
hw_size,
rank_stats,
nb_symbols_ptr,
table_log_ptr,
src,
src_size,
work_space,
wksp_size,
bmi2,
)
}
#[no_mangle]
pub unsafe extern "C" fn HUF_readStats(
huff_weight: *mut BYTE,
hw_size: usize,
rank_stats: *mut U32,
nb_symbols_ptr: *mut U32,
table_log_ptr: *mut U32,
src: *const c_void,
src_size: usize,
) -> usize {
let mut wksp = [0u32; HUF_READ_STATS_WORKSPACE_SIZE_U32];
HUF_readStats_wksp(
huff_weight,
hw_size,
rank_stats,
nb_symbols_ptr,
table_log_ptr,
src,
src_size,
wksp.as_mut_ptr() as *mut c_void,
std::mem::size_of_val(&wksp),
0,
)
}
#[cfg(test)]
mod tests {
use super::*;
fn bytes(hex: &str) -> Vec<u8> {
assert_eq!(hex.len() % 2, 0);
hex.as_bytes()
.chunks_exact(2)
.map(|pair| {
let digit = |byte: u8| match byte {
b'0'..=b'9' => byte - b'0',
b'a'..=b'f' => byte - b'a' + 10,
_ => panic!("invalid hex digit"),
};
(digit(pair[0]) << 4) | digit(pair[1])
})
.collect()
}
#[test]
fn fse_version_matches_public_header() {
assert_eq!(FSE_versionNumber(), 900);
}
#[test]
fn fse_read_ncount_matches_reference_header_and_short_prefix_errors() {
// Written by FSE_writeNCount() in the pristine C implementation.
let header = [0xd1u8, 0x28, 0x4a, 0xa9, 0x7c];
let mut normalized = [0x7fffi16; 7];
let mut max_symbol = 6u32;
let mut table_log = 0u32;
let result = unsafe {
FSE_readNCount(
normalized.as_mut_ptr(),
&mut max_symbol,
&mut table_log,
header.as_ptr() as *const c_void,
header.len(),
)
};
assert_eq!(result, header.len());
assert_eq!(max_symbol, 6);
assert_eq!(table_log, 6);
assert_eq!(normalized, [12, 9, 9, 9, 9, 8, 8]);
for prefix_size in 0..header.len() {
normalized.fill(0x7fff);
max_symbol = 6;
table_log = 0;
let result = unsafe {
FSE_readNCount(
normalized.as_mut_ptr(),
&mut max_symbol,
&mut table_log,
header.as_ptr() as *const c_void,
prefix_size,
)
};
assert_eq!(result, ERROR(ZstdErrorCode::CorruptionDetected));
}
}
#[test]
fn huf_read_stats_decodes_direct_even_weight_header() {
let source = [129u8, 0x11];
let mut weights = [0u8; 4];
let mut ranks = [0u32; HUF_TABLELOG_MAX as usize + 1];
let mut symbols = 0u32;
let mut table_log = 0u32;
let result = unsafe {
HUF_readStats(
weights.as_mut_ptr(),
weights.len(),
ranks.as_mut_ptr(),
&mut symbols,
&mut table_log,
source.as_ptr() as *const c_void,
source.len(),
)
};
assert_eq!(result, source.len());
assert_eq!(symbols, 3);
assert_eq!(table_log, 2);
assert_eq!(&weights[..symbols as usize], &[1, 1, 2]);
assert_eq!(ranks[1], 2);
assert_eq!(ranks[2], 1);
}
#[test]
fn huf_read_stats_decodes_direct_odd_weight_header() {
let source = [128u8, 0x1f];
let mut weights = [0u8; 2];
let mut ranks = [0u32; HUF_TABLELOG_MAX as usize + 1];
let mut symbols = 0u32;
let mut table_log = 0u32;
let result = unsafe {
HUF_readStats(
weights.as_mut_ptr(),
weights.len(),
ranks.as_mut_ptr(),
&mut symbols,
&mut table_log,
source.as_ptr() as *const c_void,
source.len(),
)
};
assert_eq!(result, source.len());
assert_eq!(symbols, 2);
assert_eq!(table_log, 1);
assert_eq!(weights, [1, 1]);
assert_eq!(ranks[1], 2);
}
#[test]
fn huf_read_stats_rejects_empty_truncated_and_invalid_trees() {
let mut weights = [0u8; 4];
let mut ranks = [0u32; HUF_TABLELOG_MAX as usize + 1];
let mut symbols = 0u32;
let mut table_log = 0u32;
let empty = unsafe {
HUF_readStats(
weights.as_mut_ptr(),
weights.len(),
ranks.as_mut_ptr(),
&mut symbols,
&mut table_log,
std::ptr::null(),
0,
)
};
assert_eq!(empty, ERROR(ZstdErrorCode::SrcSizeWrong));
let truncated_source = [130u8, 0x11];
let truncated = unsafe {
HUF_readStats(
weights.as_mut_ptr(),
weights.len(),
ranks.as_mut_ptr(),
&mut symbols,
&mut table_log,
truncated_source.as_ptr() as *const c_void,
truncated_source.len(),
)
};
assert_eq!(truncated, ERROR(ZstdErrorCode::SrcSizeWrong));
let invalid_tree_source = [128u8, 0x20];
let invalid_tree = unsafe {
HUF_readStats(
weights.as_mut_ptr(),
weights.len(),
ranks.as_mut_ptr(),
&mut symbols,
&mut table_log,
invalid_tree_source.as_ptr() as *const c_void,
invalid_tree_source.len(),
)
};
assert_eq!(invalid_tree, ERROR(ZstdErrorCode::CorruptionDetected));
}
#[test]
fn huf_read_stats_decodes_fse_compressed_reference_c_header() {
// Produced by HUF_writeCTable_wksp() from the pristine C implementation.
let source = bytes("181010c4a87a61c89e4674d3cb1ee70281b784d34794aa614a");
let expected_weights = bytes(
"0204060306050404040506030604020606020406030605040404050603060402\
0606020406030605040404050603060401060601040603060504030405060306",
);
let expected_ranks = [0u32, 2, 5, 9, 18, 8, 22, 0, 0, 0, 0, 0, 0];
for flags in [0, HUF_FLAGS_BMI2] {
let mut weights = [0u8; 256];
let mut ranks = [0u32; HUF_TABLELOG_MAX as usize + 1];
let mut symbols = 0u32;
let mut table_log = 0u32;
let mut workspace = [0u32; HUF_READ_STATS_WORKSPACE_SIZE_U32];
let result = unsafe {
HUF_readStats_wksp(
weights.as_mut_ptr(),
weights.len(),
ranks.as_mut_ptr(),
&mut symbols,
&mut table_log,
source.as_ptr() as *const c_void,
source.len(),
workspace.as_mut_ptr() as *mut c_void,
std::mem::size_of_val(&workspace),
flags,
)
};
assert_eq!(result, source.len());
assert_eq!(symbols, expected_weights.len() as u32);
assert_eq!(table_log, 10);
assert_eq!(&weights[..symbols as usize], expected_weights);
assert_eq!(ranks, expected_ranks);
}
}
}
Reference in New Issue View Git Blame Copy Permalink