deduplicate some FSE functions and make them (more) safe

lib/common/fse.rs

@@ -1,7 +1,6 @@
 use libc::ptrdiff_t;
 
 use crate::lib::common::bitstream::{BIT_CStream_t, BIT_addBits, BIT_flushBits, BitContainerType};
-use crate::lib::common::mem::MEM_read16;
 
 pub(crate) type FSE_CTable = core::ffi::c_uint;
 
@@ -69,63 +68,97 @@ pub(crate) struct FSE_symbolCompressionTransform {
     pub(crate) deltaNbBits: u32,
 }
 
+/// Pack two consecutive `u16` values into one `u32`.
 #[inline]
-pub(crate) unsafe fn FSE_initCState(statePtr: *mut FSE_CState_t, ct: *const FSE_CTable) {
-    let ptr = ct as *const core::ffi::c_void;
-    let u16ptr = ptr as *const u16;
-    let tableLog = MEM_read16(ptr) as u32;
-    (*statePtr).value = (1) << tableLog;
-    (*statePtr).stateTable = u16ptr.add(2) as *const core::ffi::c_void;
-    (*statePtr).symbolTT = ct.add(1).offset(
-        (if tableLog != 0 {
-            (1) << tableLog.wrapping_sub(1)
-        } else {
-            1
-        }) as isize,
-    ) as *const core::ffi::c_void;
-    (*statePtr).stateLog = tableLog;
+pub(crate) fn FSE_writeU16Pair(first: u16, second: u16) -> FSE_CTable {
+    let [a, b] = first.to_ne_bytes();
+    let [c, d] = second.to_ne_bytes();
+
+    u32::from_ne_bytes([a, b, c, d])
+}
+
+/// Read the `index`th `u16` of a `&[u32]`.
+#[inline]
+fn FSE_readU16(ct: &[FSE_CTable], index: usize) -> u16 {
+    let bytes = ct[index / 2].to_ne_bytes();
+    if index.is_multiple_of(2) {
+        u16::from_ne_bytes([bytes[0], bytes[1]])
+    } else {
+        u16::from_ne_bytes([bytes[2], bytes[3]])
+    }
+}
+
+#[inline]
+pub(crate) fn FSE_initCState(statePtr: &mut FSE_CState_t, ct: &[FSE_CTable]) {
+    // the table header occupies the first two bytes of `ct`
+    let tableLog = FSE_readU16(ct, 0) as u32;
+
+    // the state table follows the header
+    let stateTable = &ct[1..];
+    let symbolTT = &ct[FSE_symbolTTIndex(tableLog)..];
+
+    statePtr.value = 1 << tableLog;
+    statePtr.stateTable = stateTable.as_ptr().cast::<core::ffi::c_void>();
+    statePtr.symbolTT = symbolTT.as_ptr().cast::<core::ffi::c_void>();
+    statePtr.stateLog = tableLog;
 }
 
 #[inline]
-pub(crate) unsafe fn FSE_initCState2(
-    statePtr: *mut FSE_CState_t,
-    ct: *const FSE_CTable,
+pub(crate) const fn FSE_symbolTTIndex(tableLog: u32) -> usize {
+    let skip_header = 1;
+
+    match tableLog {
+        0 => skip_header + 1,
+        _ => skip_header + (1 << (tableLog - 1)),
+    }
+}
+
+/// Read the transform of `symbol` out of the symbol transformation table of `ct`.
+#[inline]
+fn FSE_readSymbolTT(
+    ct: &[FSE_CTable],
+    tableLog: u32,
     symbol: u32,
-) {
+) -> FSE_symbolCompressionTransform {
+    let index = FSE_symbolTTIndex(tableLog) + 2 * symbol as usize;
+
+    FSE_symbolCompressionTransform {
+        deltaFindState: ct[index] as core::ffi::c_int,
+        deltaNbBits: ct[index + 1],
+    }
+}
+
+#[inline]
+pub(crate) fn FSE_initCState2(statePtr: &mut FSE_CState_t, ct: &[FSE_CTable], symbol: u32) {
     FSE_initCState(statePtr, ct);
-    let symbolTT =
-        *((*statePtr).symbolTT as *const FSE_symbolCompressionTransform).offset(symbol as isize);
-    let stateTable = (*statePtr).stateTable as *const u16;
-    let nbBitsOut = (symbolTT.deltaNbBits).wrapping_add(((1) << 15) as u32) >> 16;
-    (*statePtr).value = (nbBitsOut << 16).wrapping_sub(symbolTT.deltaNbBits) as ptrdiff_t;
-    (*statePtr).value = *stateTable
-        .offset(((*statePtr).value >> nbBitsOut) + symbolTT.deltaFindState as ptrdiff_t)
-        as ptrdiff_t;
+    let symbolTT = FSE_readSymbolTT(ct, statePtr.stateLog, symbol);
+    let nbBitsOut = (symbolTT.deltaNbBits).wrapping_add((1) << 15) >> 16;
+    let value = (nbBitsOut << 16).wrapping_sub(symbolTT.deltaNbBits) as ptrdiff_t;
+
+    // the state table starts at the third `u16` of `ct`
+    let index = 2 + (value >> nbBitsOut) + symbolTT.deltaFindState as ptrdiff_t;
+    statePtr.value = FSE_readU16(ct, index as usize) as ptrdiff_t;
 }
 
 #[inline]
 pub(crate) unsafe fn FSE_encodeSymbol(
     bitC: &mut BIT_CStream_t,
-    statePtr: *mut FSE_CState_t,
+    statePtr: &mut FSE_CState_t,
     symbol: core::ffi::c_uint,
 ) {
     let symbolTT =
-        *((*statePtr).symbolTT as *const FSE_symbolCompressionTransform).offset(symbol as isize);
-    let stateTable = (*statePtr).stateTable as *const u16;
-    let nbBitsOut = (((*statePtr).value + symbolTT.deltaNbBits as ptrdiff_t) >> 16) as u32;
-    BIT_addBits(bitC, (*statePtr).value as BitContainerType, nbBitsOut);
-    (*statePtr).value = *stateTable
-        .offset(((*statePtr).value >> nbBitsOut) + symbolTT.deltaFindState as ptrdiff_t)
+        *(statePtr.symbolTT as *const FSE_symbolCompressionTransform).offset(symbol as isize);
+    let stateTable = statePtr.stateTable as *const u16;
+    let nbBitsOut = ((statePtr.value + symbolTT.deltaNbBits as ptrdiff_t) >> 16) as u32;
+    BIT_addBits(bitC, statePtr.value as BitContainerType, nbBitsOut);
+    statePtr.value = *stateTable
+        .offset((statePtr.value >> nbBitsOut) + symbolTT.deltaFindState as ptrdiff_t)
         as ptrdiff_t;
 }
 
 #[inline]
-pub(crate) unsafe fn FSE_flushCState(bitC: &mut BIT_CStream_t, statePtr: *const FSE_CState_t) {
-    BIT_addBits(
-        bitC,
-        (*statePtr).value as BitContainerType,
-        (*statePtr).stateLog,
-    );
+pub(crate) unsafe fn FSE_flushCState(bitC: &mut BIT_CStream_t, statePtr: &FSE_CState_t) {
+    BIT_addBits(bitC, statePtr.value as BitContainerType, statePtr.stateLog);
     BIT_flushBits(bitC);
 }
 

lib/compress/fse_compress.rs

@@ -1,52 +1,17 @@
-use libc::{ptrdiff_t, size_t};
+use libc::size_t;
 
 use crate::lib::common::bits::ZSTD_highbit32;
 use crate::lib::common::bitstream::{
-    BIT_CStream_t, BIT_addBits, BIT_closeCStream, BIT_flushBits, BIT_flushBitsFast,
-    BIT_initCStream, BitContainerType,
+    BIT_CStream_t, BIT_closeCStream, BIT_flushBits, BIT_flushBitsFast, BIT_initCStream,
+    BitContainerType,
 };
 use crate::lib::common::error_private::{ERR_isError, Error};
 use crate::lib::common::fse::{
-    FSE_CState_t, FSE_CTable, FSE_encodeSymbol, FSE_symbolCompressionTransform,
-    FSE_DEFAULT_TABLELOG, FSE_MAX_TABLELOG, FSE_MIN_TABLELOG, FSE_NCOUNTBOUND,
+    FSE_CState_t, FSE_CTable, FSE_encodeSymbol, FSE_flushCState, FSE_initCState2,
+    FSE_symbolCompressionTransform, FSE_symbolTTIndex, FSE_writeU16Pair, FSE_DEFAULT_TABLELOG,
+    FSE_MAX_TABLELOG, FSE_MIN_TABLELOG, FSE_NCOUNTBOUND,
 };
-use crate::lib::common::mem::{MEM_read16, MEM_write64};
-
-#[inline]
-unsafe fn FSE_initCState(statePtr: &mut FSE_CState_t, ct: *const FSE_CTable) {
-    let ptr = ct as *const core::ffi::c_void;
-    let u16ptr = ptr as *const u16;
-    let tableLog = MEM_read16(ptr) as u32;
-    statePtr.value = 1 << tableLog;
-    statePtr.stateTable = u16ptr.add(2) as *const core::ffi::c_void;
-    statePtr.symbolTT = ct.add(1).offset(
-        (if tableLog != 0 {
-            1 << tableLog.wrapping_sub(1)
-        } else {
-            1
-        }) as isize,
-    ) as *const core::ffi::c_void;
-    statePtr.stateLog = tableLog;
-}
-
-#[inline]
-unsafe fn FSE_initCState2(statePtr: &mut FSE_CState_t, ct: *const FSE_CTable, symbol: u32) {
-    FSE_initCState(statePtr, ct);
-    let symbolTT =
-        *(statePtr.symbolTT as *const FSE_symbolCompressionTransform).offset(symbol as isize);
-    let stateTable = statePtr.stateTable as *const u16;
-    let nbBitsOut = (symbolTT.deltaNbBits).wrapping_add((1 << 15) as u32) >> 16;
-    statePtr.value = (nbBitsOut << 16).wrapping_sub(symbolTT.deltaNbBits) as ptrdiff_t;
-    statePtr.value = *stateTable
-        .offset((statePtr.value >> nbBitsOut) + symbolTT.deltaFindState as ptrdiff_t)
-        as ptrdiff_t;
-}
-
-#[inline]
-unsafe fn FSE_flushCState(bitC: &mut BIT_CStream_t, statePtr: &FSE_CState_t) {
-    BIT_addBits(bitC, statePtr.value as BitContainerType, statePtr.stateLog);
-    BIT_flushBits(bitC);
-}
+use crate::lib::common::mem::MEM_write64;
 
 /// Same as FSE_buildCTable(), but using an externally allocated scratch buffer (`workSpace`).
 /// wkspSize should be sized to handle worst case situation, which is `1<<max_tableLog * sizeof(FSE_FUNCTION_TYPE)`
@@ -629,23 +594,17 @@ pub(crate) unsafe fn FSE_normalizeCount(
 }
 
 /// Fake FSE_CTable, for rle input (always same symbol).
-pub(crate) unsafe fn FSE_buildCTable_rle(ct: *mut FSE_CTable, symbolValue: u8) -> size_t {
-    let ptr = ct as *mut core::ffi::c_void;
-    let tableU16 = (ptr as *mut u16).add(2);
-    let FSCTptr = (ptr as *mut u32).add(2) as *mut core::ffi::c_void;
-    let symbolTT = FSCTptr as *mut FSE_symbolCompressionTransform;
+pub(crate) fn FSE_buildCTable_rle(ct: &mut [FSE_CTable], symbolValue: u8) -> size_t {
+    // header: a tableLog of zero, and `symbolValue` as the maximum symbol value
+    ct[0] = FSE_writeU16Pair(0, symbolValue as u16);
 
-    // header
-    *tableU16.sub(2) = 0;
-    *tableU16.sub(1) = symbolValue as u16;
-
-    // Build table
-    *tableU16 = 0;
-    *tableU16.add(1) = 0; // just in case
+    // the (two-entry) state table, zeroed just in case
+    ct[1] = FSE_writeU16Pair(0, 0);
 
     // Build Symbol Transformation Table
-    (*symbolTT.offset(symbolValue as isize)).deltaNbBits = 0;
-    (*symbolTT.offset(symbolValue as isize)).deltaFindState = 0;
+    let index = FSE_symbolTTIndex(0) + 2 * symbolValue as usize;
+    ct[index] = 0; // deltaFindState
+    ct[index + 1] = 0; // deltaNbBits
 
     0
 }
@@ -655,7 +614,7 @@ unsafe fn FSE_compress_usingCTable_generic(
     dstSize: size_t,
     src: *const core::ffi::c_void,
     mut srcSize: size_t,
-    ct: *const FSE_CTable,
+    ct: &[FSE_CTable],
     fast: bool,
 ) -> size_t {
     let istart = src as *const u8;
@@ -773,7 +732,7 @@ pub(crate) unsafe fn FSE_compress_usingCTable(
     dstSize: size_t,
     src: *const core::ffi::c_void,
     srcSize: size_t,
-    ct: *const FSE_CTable,
+    ct: &[FSE_CTable],
 ) -> size_t {
     let fast = dstSize
         >= srcSize

lib/compress/hist.rs

@@ -31,7 +31,7 @@ pub unsafe fn HIST_add(
 
 pub unsafe fn HIST_count_simple(
     count: *mut core::ffi::c_uint,
-    maxSymbolValuePtr: *mut core::ffi::c_uint,
+    maxSymbolValuePtr: &mut core::ffi::c_uint,
     src: *const core::ffi::c_void,
     srcSize: size_t,
 ) -> core::ffi::c_uint {
@@ -83,7 +83,7 @@ pub unsafe fn HIST_count_simple(
 /// histogram's alphabet is larger than *maxSymbolValuePtr)
 unsafe fn HIST_count_parallel_wksp(
     count: *mut core::ffi::c_uint,
-    maxSymbolValuePtr: *mut core::ffi::c_uint,
+    maxSymbolValuePtr: &mut core::ffi::c_uint,
     source: *const core::ffi::c_void,
     sourceSize: size_t,
     check: HIST_checkInput_e,
@@ -181,7 +181,7 @@ unsafe fn HIST_count_parallel_wksp(
 /// `workSpaceSize` must be >= HIST_WKSP_SIZE
 pub unsafe fn HIST_countFast_wksp(
     count: *mut core::ffi::c_uint,
-    maxSymbolValuePtr: *mut core::ffi::c_uint,
+    maxSymbolValuePtr: &mut core::ffi::c_uint,
     source: *const core::ffi::c_void,
     sourceSize: size_t,
     workSpace: *mut core::ffi::c_void,
@@ -211,7 +211,7 @@ pub unsafe fn HIST_countFast_wksp(
 /// `workSpace` size must be table of >= HIST_WKSP_SIZE_U32 unsigned
 pub unsafe fn HIST_count_wksp(
     count: *mut core::ffi::c_uint,
-    maxSymbolValuePtr: *mut core::ffi::c_uint,
+    maxSymbolValuePtr: &mut core::ffi::c_uint,
     source: *const core::ffi::c_void,
     sourceSize: size_t,
     workSpace: *mut core::ffi::c_void,
@@ -248,7 +248,7 @@ pub unsafe fn HIST_count_wksp(
 /// fast variant (unsafe : won't check if src contains values beyond count[] limit)
 pub unsafe fn HIST_countFast(
     count: *mut core::ffi::c_uint,
-    maxSymbolValuePtr: *mut core::ffi::c_uint,
+    maxSymbolValuePtr: &mut core::ffi::c_uint,
     source: *const core::ffi::c_void,
     sourceSize: size_t,
 ) -> size_t {
@@ -265,7 +265,7 @@ pub unsafe fn HIST_countFast(
 
 pub unsafe fn HIST_count(
     count: *mut core::ffi::c_uint,
-    maxSymbolValuePtr: *mut core::ffi::c_uint,
+    maxSymbolValuePtr: &mut core::ffi::c_uint,
     src: *const core::ffi::c_void,
     srcSize: size_t,
 ) -> size_t {

lib/compress/huf_compress.rs

@@ -172,7 +172,7 @@ unsafe fn HUF_compressWeights(
             oend.offset_from_unsigned(op),
             weightTable.as_ptr().cast::<c_void>(),
             wtSize,
-            ((*wksp).CTable).as_mut_ptr(),
+            &(*wksp).CTable,
         );
         if ERR_isError(cSize) {
             return cSize;