ddidderr/zstd-rs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(rust): port FSE entropy decoding

Browse Source 
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
This commit is contained in:
ddidderr
2026-07-10 20:04:26 +02:00
parent 554a30da41
commit a0f2b2a14c
6 changed files with 1264 additions and 652 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/entropy_common.rs
+597
View File
@@ -0,0 +1,597 @@
#![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);
}
}
}
rust/src/fse_decompress.rs
+653
View File
@@ -0,0 +1,653 @@
#![allow(non_snake_case)]
use crate::bits::ZSTD_highbit32;
use crate::bitstream::*;
use crate::errors::{ERR_isError, ZstdErrorCode, ERROR};
use crate::mem::MEM_write64;
use std::os::raw::{c_int, c_uint, c_void};
pub const FSE_MAX_SYMBOL_VALUE: u32 = 255;
pub const FSE_MAX_TABLELOG: u32 = 12;
#[repr(C)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct FSE_DTableHeader {
pub tableLog: u16,
pub fastMode: u16,
}
#[repr(C)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct FSE_decode_t {
pub newState: u16,
pub symbol: u8,
pub nbBits: u8,
}
/// `FSE_DTable` is an opaque `unsigned` in the C API. Its first word stores
/// `FSE_DTableHeader`, followed by `1 << tableLog` `FSE_decode_t` entries.
#[repr(C, align(4))]
pub struct FSE_DTable {
pub header: FSE_DTableHeader,
}
#[repr(C)]
pub struct FSE_DState_t {
pub state: usize,
pub table: *const c_void,
}
#[repr(C)]
pub struct FSE_DecompressWksp {
pub ncount: [i16; (FSE_MAX_SYMBOL_VALUE + 1) as usize],
}
#[inline]
fn FSE_TABLESTEP(table_size: usize) -> usize {
(table_size >> 1) + (table_size >> 3) + 3
}
#[inline]
fn fse_build_dtable_wksp_size(table_log: u32, max_symbol_value: u32) -> Option<usize> {
let table_size = 1usize.checked_shl(table_log)?;
let symbols_size = 2usize.checked_mul(max_symbol_value as usize + 1)?;
symbols_size.checked_add(table_size)?.checked_add(8)
}
#[inline]
fn fse_dtable_size(table_log: u32) -> Option<usize> {
let table_size = 1usize.checked_shl(table_log)?;
(table_size + 1).checked_mul(std::mem::size_of::<FSE_DTable>())
}
#[inline]
fn fse_decompress_wksp_size(table_log: u32, max_symbol_value: u32) -> Option<usize> {
let dtable_u32 = (1usize.checked_shl(table_log)?).checked_add(1)?;
let build_size = fse_build_dtable_wksp_size(table_log, max_symbol_value)?;
let build_u32 = build_size.checked_add(3)? / 4;
// Keep the extra trailing word from FSE_DECOMPRESS_WKSP_SIZE_U32.
let ncount_u32 = ((FSE_MAX_SYMBOL_VALUE as usize + 1) >> 1) + 1;
dtable_u32
.checked_add(1)?
.checked_add(build_u32)?
.checked_add(ncount_u32)?
.checked_mul(4)
}
unsafe fn FSE_buildDTable_internal(
dt: *mut FSE_DTable,
normalized_counter: *const i16,
max_symbol_value: u32,
table_log: u32,
work_space: *mut c_void,
wksp_size: usize,
) -> usize {
let required_wksp = match fse_build_dtable_wksp_size(table_log, max_symbol_value) {
Some(size) => size,
None => return ERROR(ZstdErrorCode::MaxSymbolValueTooLarge),
};
if required_wksp > wksp_size {
return ERROR(ZstdErrorCode::MaxSymbolValueTooLarge);
}
if max_symbol_value > FSE_MAX_SYMBOL_VALUE {
return ERROR(ZstdErrorCode::MaxSymbolValueTooLarge);
}
if table_log > FSE_MAX_TABLELOG {
return ERROR(ZstdErrorCode::TableLogTooLarge);
}
if table_log == 0 {
return ERROR(ZstdErrorCode::Generic);
}
let table_decode = (dt as *mut u8).add(std::mem::size_of::<FSE_DTable>()) as *mut FSE_decode_t;
let symbol_next = work_space as *mut u16;
let spread = symbol_next.add(max_symbol_value as usize + 1) as *mut u8;
let max_sv1 = max_symbol_value + 1;
let table_size = 1usize << table_log;
let mut high_threshold = table_size - 1;
let mut low_probability_count = 0usize;
let mut fast_mode = 1u16;
let large_limit = (1u32 << (table_log - 1)) as i16;
for symbol in 0..max_sv1 {
let count = *normalized_counter.add(symbol as usize);
if count == -1 {
if low_probability_count == table_size {
return ERROR(ZstdErrorCode::Generic);
}
let target = table_size - 1 - low_probability_count;
(*table_decode.add(target)).symbol = symbol as u8;
*symbol_next.add(symbol as usize) = 1;
low_probability_count += 1;
high_threshold = if low_probability_count == table_size {
usize::MAX
} else {
table_size - 1 - low_probability_count
};
} else {
if count >= large_limit {
fast_mode = 0;
}
*symbol_next.add(symbol as usize) = count as u16;
}
}
std::ptr::write(
dt,
FSE_DTable {
header: FSE_DTableHeader {
tableLog: table_log as u16,
fastMode: fast_mode,
},
},
);
if high_threshold == table_size - 1 {
let table_mask = table_size - 1;
let step = FSE_TABLESTEP(table_size);
let mut pos = 0usize;
let mut repeated_symbol = 0u64;
for symbol in 0..max_sv1 {
let count = *normalized_counter.add(symbol as usize);
if count < 0 {
return ERROR(ZstdErrorCode::Generic);
}
let count = count as usize;
MEM_write64(spread.add(pos) as *mut c_void, repeated_symbol);
for offset in (8..count).step_by(8) {
MEM_write64(spread.add(pos + offset) as *mut c_void, repeated_symbol);
}
pos += count;
repeated_symbol = repeated_symbol.wrapping_add(0x0101_0101_0101_0101);
}
let mut position = 0usize;
for symbol_index in (0..table_size).step_by(2) {
for unroll in 0..2 {
let target = (position + unroll * step) & table_mask;
(*table_decode.add(target)).symbol = *spread.add(symbol_index + unroll);
}
position = (position + 2 * step) & table_mask;
}
debug_assert_eq!(position, 0);
} else {
let table_mask = table_size - 1;
let step = FSE_TABLESTEP(table_size);
let mut position = 0usize;
for symbol in 0..max_sv1 {
let count = *normalized_counter.add(symbol as usize);
for _ in 0..count.max(0) as usize {
(*table_decode.add(position)).symbol = symbol as u8;
position = (position + step) & table_mask;
while position > high_threshold {
position = (position + step) & table_mask;
}
}
}
if position != 0 {
return ERROR(ZstdErrorCode::Generic);
}
}
for index in 0..table_size {
let symbol = (*table_decode.add(index)).symbol;
let next_state_ptr = symbol_next.add(symbol as usize);
let next_state = *next_state_ptr as u32;
*next_state_ptr = (*next_state_ptr).wrapping_add(1);
if next_state == 0 {
return ERROR(ZstdErrorCode::Generic);
}
let nb_bits = table_log - ZSTD_highbit32(next_state);
(*table_decode.add(index)).nbBits = nb_bits as u8;
(*table_decode.add(index)).newState =
((next_state << nb_bits).wrapping_sub(table_size as u32)) as u16;
}
0
}
#[no_mangle]
pub unsafe extern "C" fn FSE_buildDTable_wksp(
dt: *mut FSE_DTable,
normalized_counter: *const i16,
max_symbol_value: c_uint,
table_log: c_uint,
work_space: *mut c_void,
wksp_size: usize,
) -> usize {
FSE_buildDTable_internal(
dt,
normalized_counter,
max_symbol_value,
table_log,
work_space,
wksp_size,
)
}
unsafe fn FSE_initDState(
state: *mut FSE_DState_t,
bit_stream: *mut BIT_DStream_t,
dt: *const FSE_DTable,
) {
(*state).state = BIT_readBits(bit_stream, (*dt).header.tableLog as u32);
BIT_reloadDStream(bit_stream);
(*state).table = dt.add(1) as *const c_void;
}
unsafe fn FSE_decodeSymbol(state: *mut FSE_DState_t, bit_stream: *mut BIT_DStream_t) -> u8 {
let entry = &*((*state).table as *const FSE_decode_t).add((*state).state);
let low_bits = BIT_readBits(bit_stream, entry.nbBits as u32);
(*state).state = entry.newState as usize + low_bits;
entry.symbol
}
unsafe fn FSE_decodeSymbolFast(state: *mut FSE_DState_t, bit_stream: *mut BIT_DStream_t) -> u8 {
let entry = &*((*state).table as *const FSE_decode_t).add((*state).state);
let low_bits = BIT_readBitsFast(bit_stream, entry.nbBits as u32);
(*state).state = entry.newState as usize + low_bits;
entry.symbol
}
#[inline]
unsafe fn fse_decode_symbol(
state: *mut FSE_DState_t,
bit_stream: *mut BIT_DStream_t,
fast: bool,
) -> u8 {
if fast {
FSE_decodeSymbolFast(state, bit_stream)
} else {
FSE_decodeSymbol(state, bit_stream)
}
}
unsafe fn FSE_decompress_usingDTable_generic(
dst: *mut c_void,
max_dst_size: usize,
c_src: *const c_void,
c_src_size: usize,
dt: *const FSE_DTable,
fast: bool,
) -> usize {
let mut bit_stream = std::mem::zeroed::<BIT_DStream_t>();
let init_result = BIT_initDStream(&mut bit_stream, c_src, c_src_size);
if ERR_isError(init_result) {
return init_result;
}
let mut state1 = FSE_DState_t {
state: 0,
table: std::ptr::null(),
};
let mut state2 = FSE_DState_t {
state: 0,
table: std::ptr::null(),
};
FSE_initDState(&mut state1, &mut bit_stream, dt);
FSE_initDState(&mut state2, &mut bit_stream, dt);
if BIT_reloadDStream(&mut bit_stream) == BIT_DStream_status::Overflow {
return ERROR(ZstdErrorCode::CorruptionDetected);
}
let output = dst as *mut u8;
let mut output_pos = 0usize;
while BIT_reloadDStream(&mut bit_stream) == BIT_DStream_status::Unfinished
&& output_pos < max_dst_size.saturating_sub(3)
{
*output.add(output_pos) = fse_decode_symbol(&mut state1, &mut bit_stream, fast);
output_pos += 1;
if FSE_MAX_TABLELOG * 2 + 7 > usize::BITS {
BIT_reloadDStream(&mut bit_stream);
}
*output.add(output_pos) = fse_decode_symbol(&mut state2, &mut bit_stream, fast);
output_pos += 1;
if FSE_MAX_TABLELOG * 4 + 7 > usize::BITS
&& BIT_reloadDStream(&mut bit_stream) != BIT_DStream_status::Unfinished
{
break;
}
*output.add(output_pos) = fse_decode_symbol(&mut state1, &mut bit_stream, fast);
output_pos += 1;
if FSE_MAX_TABLELOG * 2 + 7 > usize::BITS {
BIT_reloadDStream(&mut bit_stream);
}
*output.add(output_pos) = fse_decode_symbol(&mut state2, &mut bit_stream, fast);
output_pos += 1;
}
loop {
if max_dst_size.saturating_sub(output_pos) < 2 {
return ERROR(ZstdErrorCode::DstSizeTooSmall);
}
*output.add(output_pos) = fse_decode_symbol(&mut state1, &mut bit_stream, fast);
output_pos += 1;
if BIT_reloadDStream(&mut bit_stream) == BIT_DStream_status::Overflow {
*output.add(output_pos) = fse_decode_symbol(&mut state2, &mut bit_stream, fast);
output_pos += 1;
break;
}
if max_dst_size.saturating_sub(output_pos) < 2 {
return ERROR(ZstdErrorCode::DstSizeTooSmall);
}
*output.add(output_pos) = fse_decode_symbol(&mut state2, &mut bit_stream, fast);
output_pos += 1;
if BIT_reloadDStream(&mut bit_stream) == BIT_DStream_status::Overflow {
*output.add(output_pos) = fse_decode_symbol(&mut state1, &mut bit_stream, fast);
output_pos += 1;
break;
}
}
output_pos
}
#[allow(clippy::too_many_arguments)]
unsafe fn FSE_decompress_wksp_body(
dst: *mut c_void,
dst_capacity: usize,
c_src: *const c_void,
mut c_src_size: usize,
max_log: u32,
work_space: *mut c_void,
wksp_size: usize,
bmi2: c_int,
) -> usize {
if wksp_size < std::mem::size_of::<FSE_DecompressWksp>() {
return ERROR(ZstdErrorCode::Generic);
}
let input_start = c_src as *const u8;
let wksp = work_space as *mut FSE_DecompressWksp;
let dtable =
(work_space as *mut u8).add(std::mem::size_of::<FSE_DecompressWksp>()) as *mut FSE_DTable;
let mut table_log = 0u32;
let mut max_symbol_value = FSE_MAX_SYMBOL_VALUE;
let ncount_length = crate::entropy_common::FSE_readNCount_bmi2(
(*wksp).ncount.as_mut_ptr(),
&mut max_symbol_value,
&mut table_log,
c_src,
c_src_size,
bmi2,
);
if ERR_isError(ncount_length) {
return ncount_length;
}
if table_log > max_log {
return ERROR(ZstdErrorCode::TableLogTooLarge);
}
if ncount_length > c_src_size {
return ERROR(ZstdErrorCode::CorruptionDetected);
}
let required_size = match fse_decompress_wksp_size(table_log, max_symbol_value) {
Some(size) => size,
None => return ERROR(ZstdErrorCode::TableLogTooLarge),
};
if required_size > wksp_size {
return ERROR(ZstdErrorCode::TableLogTooLarge);
}
let dtable_size = fse_dtable_size(table_log).expect("validated FSE table dimensions fit usize");
let build_offset = std::mem::size_of::<FSE_DecompressWksp>() + dtable_size;
let build_workspace = (work_space as *mut u8).add(build_offset) as *mut c_void;
let build_workspace_size = wksp_size - build_offset;
let build_result = FSE_buildDTable_internal(
dtable,
(*wksp).ncount.as_ptr(),
max_symbol_value,
table_log,
build_workspace,
build_workspace_size,
);
if ERR_isError(build_result) {
return build_result;
}
let payload = input_start.add(ncount_length);
c_src_size -= ncount_length;
FSE_decompress_usingDTable_generic(
dst,
dst_capacity,
payload as *const c_void,
c_src_size,
dtable,
(*dtable).header.fastMode != 0,
)
}
#[no_mangle]
pub unsafe extern "C" fn FSE_decompress_wksp_bmi2(
dst: *mut c_void,
dst_capacity: usize,
c_src: *const c_void,
c_src_size: usize,
max_log: c_uint,
work_space: *mut c_void,
wksp_size: usize,
bmi2: c_int,
) -> usize {
FSE_decompress_wksp_body(
dst,
dst_capacity,
c_src,
c_src_size,
max_log,
work_space,
wksp_size,
bmi2,
)
}
#[cfg(test)]
mod tests {
use super::*;
fn dtable_storage(table_log: u32) -> Vec<u32> {
vec![0; 1 + (1usize << table_log)]
}
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()
}
fn decompress(compressed: &[u8], capacity: usize) -> (usize, Vec<u8>) {
let mut output = vec![0u8; capacity];
let mut workspace = vec![0u32; 6000];
let result = unsafe {
FSE_decompress_wksp_bmi2(
output.as_mut_ptr() as *mut c_void,
output.len(),
compressed.as_ptr() as *const c_void,
compressed.len(),
FSE_MAX_TABLELOG,
workspace.as_mut_ptr() as *mut c_void,
workspace.len() * 4,
0,
)
};
(result, output)
}
#[test]
fn abi_layout_matches_fse_h() {
assert_eq!(std::mem::size_of::<FSE_DTable>(), 4);
assert_eq!(std::mem::align_of::<FSE_DTable>(), 4);
assert_eq!(std::mem::size_of::<FSE_decode_t>(), 4);
assert_eq!(std::mem::size_of::<FSE_DecompressWksp>(), 512);
}
#[test]
fn build_dtable_advances_each_symbol_state() {
let normalized = [8i16, 8, 8, 8];
let mut table = dtable_storage(5);
let mut workspace = vec![0u32; 64];
let result = unsafe {
FSE_buildDTable_wksp(
table.as_mut_ptr() as *mut FSE_DTable,
normalized.as_ptr(),
3,
5,
workspace.as_mut_ptr() as *mut c_void,
workspace.len() * 4,
)
};
assert_eq!(result, 0);
let header = unsafe { &*(table.as_ptr() as *const FSE_DTableHeader) };
assert_eq!(
*header,
FSE_DTableHeader {
tableLog: 5,
fastMode: 1
}
);
let entries = unsafe {
std::slice::from_raw_parts(
table.as_ptr().add(1) as *const FSE_decode_t,
1 << header.tableLog,
)
};
for symbol in 0..4u8 {
let mut states: Vec<u16> = entries
.iter()
.filter(|entry| entry.symbol == symbol)
.map(|entry| {
assert_eq!(entry.nbBits, 2);
entry.newState
})
.collect();
states.sort_unstable();
assert_eq!(states, [0, 4, 8, 12, 16, 20, 24, 28]);
}
}
#[test]
fn build_dtable_places_low_probability_symbols_at_end() {
let normalized = [-1i16, -1, 15, 15];
let mut table = dtable_storage(5);
let mut workspace = vec![0u32; 64];
let result = unsafe {
FSE_buildDTable_wksp(
table.as_mut_ptr() as *mut FSE_DTable,
normalized.as_ptr(),
3,
5,
workspace.as_mut_ptr() as *mut c_void,
workspace.len() * 4,
)
};
assert_eq!(result, 0);
let entries =
unsafe { std::slice::from_raw_parts(table.as_ptr().add(1) as *const FSE_decode_t, 32) };
assert_eq!(entries[31].symbol, 0);
assert_eq!(entries[30].symbol, 1);
assert_eq!(entries[31].nbBits, 5);
assert_eq!(entries[31].newState, 0);
assert_eq!(entries[30].nbBits, 5);
assert_eq!(entries[30].newState, 0);
}
#[test]
fn build_dtable_reports_workspace_and_dimension_errors() {
let normalized = [32i16];
let mut table = dtable_storage(5);
let mut workspace = [0u32; 1];
let too_small = unsafe {
FSE_buildDTable_wksp(
table.as_mut_ptr() as *mut FSE_DTable,
normalized.as_ptr(),
0,
5,
workspace.as_mut_ptr() as *mut c_void,
workspace.len() * 4,
)
};
assert_eq!(too_small, ERROR(ZstdErrorCode::MaxSymbolValueTooLarge));
let mut large_workspace = vec![0u32; 4096];
let table_log_too_large = unsafe {
FSE_buildDTable_wksp(
table.as_mut_ptr() as *mut FSE_DTable,
normalized.as_ptr(),
0,
FSE_MAX_TABLELOG + 1,
large_workspace.as_mut_ptr() as *mut c_void,
large_workspace.len() * 4,
)
};
assert_eq!(table_log_too_large, ERROR(ZstdErrorCode::TableLogTooLarge));
}
#[test]
fn decompression_workspace_size_matches_c_macro() {
assert_eq!(fse_build_dtable_wksp_size(6, 11), Some(96));
assert_eq!(fse_decompress_wksp_size(6, 11), Some(876));
assert_eq!(fse_dtable_size(6), Some(260));
}
#[test]
fn decompresses_balanced_reference_c_stream() {
// Generated by the pristine HEAD implementation using a 257-byte,
// seven-symbol input and FSE_compress_usingCTable().
let compressed = bytes(
"d1284aa97cd5fce4cd4fdefce4cd4fdefce4cd4fdefce4cd4fdefce4cd4fdefc\
e4cd4fdefce4cd4fdefce4cd4fdefce4cd4fdefce4cd4fdefce4cd4fdefce4cd\
4fdefce4cd4fdefce4cd4fdefce4cd4fdefce4cd4fdefce4cd4fdefce4cd4fde\
c415",
);
let expected: Vec<u8> = (0..257).map(|index| (index % 7) as u8).collect();
let (result, output) = decompress(&compressed, expected.len());
assert_eq!(result, expected.len());
assert_eq!(output, expected);
}
#[test]
fn decompresses_skewed_reference_c_stream_and_reports_small_output() {
// Exercises low-probability symbols and the non-fast decode path.
let compressed = bytes(
"0100800c000000003dc2200147143cc008a2061471b400828c42a4a2871844e090\
820e18429481428e2a40a0518858d41103093852e001460835a0c8a1051034142\
20bf00f",
);
let expected: Vec<u8> = (0..511)
.map(|index| {
if index % 10 != 0 {
3
} else {
(index % 17) as u8
}
})
.collect();
let (result, output) = decompress(&compressed, expected.len());
assert_eq!(result, expected.len());
assert_eq!(output, expected);
let (too_small, _) = decompress(&compressed, expected.len() - 1);
assert_eq!(too_small, ERROR(ZstdErrorCode::DstSizeTooSmall));
}
}
rust/src/lib.rs
+2
View File
@@ -5,7 +5,9 @@ pub mod bitstream;
pub mod common;
pub mod cpu;
pub mod debug;
pub mod entropy_common;
pub mod errors;
pub mod fse_decompress;
pub mod mem;
pub mod xxhash;
pub mod zstd_common;