zstd-rs/rust/src/fse_decompress.rs

#![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));
    }
}