use alloc::vec;
use alloc::vec::Vec;
use super::bits::BitWriter;
use super::{EntropyError, highbit32};
pub(crate) const MIN_TABLELOG: u32 = 5;
pub(crate) const MAX_TABLELOG: u32 = 12;
pub(crate) const DEFAULT_TABLELOG: u32 = 11;
pub(crate) const NCOUNT_BOUND: usize = 512;
pub(crate) const MAX_SYMBOL_VALUE: u32 = 255;
pub(crate) const fn block_bound(size: usize) -> usize {
size + (size >> 7) + 4 + core::mem::size_of::<usize>()
}
pub(crate) const fn compress_bound(size: usize) -> usize {
NCOUNT_BOUND + block_bound(size)
}
const fn table_step(table_size: u32) -> u32 {
(table_size >> 1) + (table_size >> 3) + 3
}
fn min_table_log(src_size: usize, max_symbol_value: u32) -> u32 {
let min_bits_src = highbit32(src_size as u32) + 1;
let min_bits_symbols = highbit32(max_symbol_value) + 2;
min_bits_src.min(min_bits_symbols)
}
pub(crate) fn optimal_table_log_internal(
max_table_log: u32,
src_size: usize,
max_symbol_value: u32,
minus: u32,
) -> u32 {
let max_bits_src = highbit32((src_size as u32).wrapping_sub(1)).wrapping_sub(minus);
let mut table_log = max_table_log;
let min_bits = min_table_log(src_size, max_symbol_value);
if table_log == 0 {
table_log = DEFAULT_TABLELOG;
}
if max_bits_src < table_log {
table_log = max_bits_src;
}
if min_bits > table_log {
table_log = min_bits;
}
table_log.clamp(MIN_TABLELOG, MAX_TABLELOG)
}
pub(crate) fn optimal_table_log(max_table_log: u32, src_size: usize, max_symbol_value: u32) -> u32 {
optimal_table_log_internal(max_table_log, src_size, max_symbol_value, 2)
}
fn normalize_m2(
norm: &mut [i16],
table_log: u32,
count: &[u32],
mut total: usize,
max_symbol_value: u32,
low_prob_count: i16,
) -> Result<(), EntropyError> {
const NOT_YET_ASSIGNED: i16 = -2;
let mut distributed: u32 = 0;
let low_threshold = (total >> table_log) as u32;
let mut low_one = ((total * 3) >> (table_log + 1)) as u32;
for (slot, &c) in norm.iter_mut().zip(&count[..=max_symbol_value as usize]) {
if c == 0 {
*slot = 0;
continue;
}
if c <= low_threshold {
*slot = low_prob_count;
distributed += 1;
total -= c as usize;
continue;
}
if c <= low_one {
*slot = 1;
distributed += 1;
total -= c as usize;
continue;
}
*slot = NOT_YET_ASSIGNED;
}
let mut to_distribute = (1u32 << table_log) - distributed;
if to_distribute == 0 {
return Ok(());
}
if total / to_distribute as usize > low_one as usize {
low_one = ((total * 3) / (to_distribute as usize * 2)) as u32;
for s in 0..=max_symbol_value as usize {
if norm[s] == NOT_YET_ASSIGNED && count[s] <= low_one {
norm[s] = 1;
distributed += 1;
total -= count[s] as usize;
}
}
to_distribute = (1u32 << table_log) - distributed;
}
if distributed == max_symbol_value + 1 {
let mut max_v = 0usize;
let mut max_c = 0u32;
for (s, &c) in count[..=max_symbol_value as usize].iter().enumerate() {
if c > max_c {
max_v = s;
max_c = c;
}
}
norm[max_v] = norm[max_v].wrapping_add(to_distribute as i16);
return Ok(());
}
if total == 0 {
let mut s = 0usize;
while to_distribute > 0 {
if norm[s] > 0 {
to_distribute -= 1;
norm[s] += 1;
}
s = (s + 1) % (max_symbol_value as usize + 1);
}
return Ok(());
}
let v_step_log = 62 - u64::from(table_log);
let mid = (1u64 << (v_step_log - 1)) - 1;
let r_step =
(((1u64 << v_step_log) * u64::from(to_distribute)) + mid) / u64::from(total as u32);
let mut tmp_total = mid;
for (slot, &c) in norm.iter_mut().zip(&count[..=max_symbol_value as usize]) {
if *slot == NOT_YET_ASSIGNED {
let end = tmp_total.wrapping_add(u64::from(c).wrapping_mul(r_step));
let s_start = (tmp_total >> v_step_log) as u32;
let s_end = (end >> v_step_log) as u32;
let weight = s_end.wrapping_sub(s_start);
if weight < 1 {
return Err(EntropyError::Generic);
}
*slot = weight as i16;
tmp_total = end;
}
}
Ok(())
}
pub(crate) fn normalize_count(
norm: &mut [i16],
mut table_log: u32,
count: &[u32],
total: usize,
max_symbol_value: u32,
use_low_prob_count: bool,
) -> Result<u32, EntropyError> {
const RTB_TABLE: [u64; 8] = [
0, 473_195, 504_333, 520_860, 550_000, 700_000, 750_000, 830_000,
];
if table_log == 0 {
table_log = DEFAULT_TABLELOG;
}
if table_log < MIN_TABLELOG {
return Err(EntropyError::Generic);
}
if table_log > MAX_TABLELOG {
return Err(EntropyError::TableLogTooLarge);
}
if table_log < min_table_log(total, max_symbol_value) {
return Err(EntropyError::Generic);
}
let low_prob_count: i16 = if use_low_prob_count { -1 } else { 1 };
let scale = 62 - u64::from(table_log);
let step = (1u64 << 62) / u64::from(total as u32);
let v_step = 1u64 << (scale - 20);
let mut still_to_distribute: i32 = 1 << table_log;
let mut largest = 0usize;
let mut largest_p: i16 = 0;
let low_threshold = (total >> table_log) as u32;
for s in 0..=max_symbol_value as usize {
if count[s] as usize == total {
return Ok(0);
}
if count[s] == 0 {
norm[s] = 0;
continue;
}
if count[s] <= low_threshold {
norm[s] = low_prob_count;
still_to_distribute -= 1;
} else {
let scaled = u64::from(count[s]).wrapping_mul(step);
let mut proba = (scaled >> scale) as i16;
if proba < 8 {
let rest_to_beat = v_step * RTB_TABLE[proba as usize];
let remainder = scaled.wrapping_sub(u64::from(proba as u16) << scale);
proba = proba.wrapping_add(i16::from(remainder > rest_to_beat));
}
if proba > largest_p {
largest_p = proba;
largest = s;
}
norm[s] = proba;
still_to_distribute -= i32::from(proba);
}
}
if -still_to_distribute >= i32::from(norm[largest] >> 1) {
normalize_m2(
norm,
table_log,
count,
total,
max_symbol_value,
low_prob_count,
)?;
} else {
norm[largest] = norm[largest].wrapping_add(still_to_distribute as i16);
}
Ok(table_log)
}
#[derive(Clone, Copy, Debug, Default)]
struct SymbolTransform {
delta_find_state: i32,
delta_nb_bits: u32,
}
#[derive(Debug)]
pub(crate) struct CTable {
table_log: u32,
state_table: Vec<u16>,
symbol_tt: Vec<SymbolTransform>,
}
impl CTable {
pub(crate) fn build(
norm: &[i16],
max_symbol_value: u32,
table_log: u32,
) -> Result<Self, EntropyError> {
if table_log > MAX_TABLELOG || max_symbol_value > MAX_SYMBOL_VALUE {
return Err(EntropyError::TableLogTooLarge);
}
let max_sv1 = max_symbol_value as usize + 1;
if norm.len() < max_sv1 {
return Err(EntropyError::Generic);
}
let table_size = 1u32 << table_log;
let table_mask = table_size - 1;
let step = table_step(table_size);
let mut cumul = vec![0u16; max_sv1 + 1];
let mut table_symbol = vec![0u8; table_size as usize];
let mut high_threshold = table_size - 1;
for u in 1..=max_sv1 {
if norm[u - 1] == -1 {
cumul[u] = cumul[u - 1] + 1;
table_symbol[high_threshold as usize] = (u - 1) as u8;
high_threshold = high_threshold.wrapping_sub(1);
} else {
cumul[u] = cumul[u - 1].wrapping_add(norm[u - 1] as u16);
}
}
cumul[max_sv1] = (table_size + 1) as u16;
let mut position = 0u32;
for (symbol, &freq) in norm[..max_sv1].iter().enumerate() {
for _ in 0..freq.max(0) {
table_symbol[position as usize] = symbol as u8;
position = (position + step) & table_mask;
while position > high_threshold {
position = (position + step) & table_mask;
}
}
}
let mut state_table = vec![0u16; table_size as usize];
for (u, &symbol) in table_symbol.iter().enumerate() {
let slot = &mut cumul[usize::from(symbol)];
state_table[usize::from(*slot)] = (table_size as usize + u) as u16;
*slot += 1;
}
let mut symbol_tt = vec![SymbolTransform::default(); max_sv1];
let mut total: u32 = 0;
for (s, &count) in norm[..max_sv1].iter().enumerate() {
match count {
0 => {
symbol_tt[s].delta_nb_bits =
((table_log + 1) << 16).wrapping_sub(1 << table_log);
}
-1 | 1 => {
symbol_tt[s].delta_nb_bits = (table_log << 16).wrapping_sub(1 << table_log);
symbol_tt[s].delta_find_state = total.wrapping_sub(1) as i32;
total += 1;
}
_ => {
let count = count as u32;
let max_bits_out = table_log - highbit32(count - 1);
let min_state_plus = count << max_bits_out;
symbol_tt[s].delta_nb_bits = (max_bits_out << 16).wrapping_sub(min_state_plus);
symbol_tt[s].delta_find_state = total.wrapping_sub(count) as i32;
total += count;
}
}
}
Ok(Self {
table_log,
state_table,
symbol_tt,
})
}
}
#[derive(Clone, Copy, Debug, Default)]
struct CState {
value: i64,
}
impl CState {
fn new(ct: &CTable, symbol: u8) -> Self {
let tt = ct.symbol_tt[usize::from(symbol)];
let nb_bits_out = tt.delta_nb_bits.wrapping_add(1 << 15) >> 16;
let value = i64::from((nb_bits_out << 16).wrapping_sub(tt.delta_nb_bits));
let index = (value >> nb_bits_out) + i64::from(tt.delta_find_state);
Self {
value: i64::from(ct.state_table[index as usize]),
}
}
#[inline]
fn encode(&mut self, bits: &mut BitWriter, ct: &CTable, symbol: u8) {
let tt = ct.symbol_tt[usize::from(symbol)];
let nb_bits_out = ((self.value + i64::from(tt.delta_nb_bits)) >> 16) as u32;
bits.add_bits(self.value as u64, nb_bits_out);
let index = (self.value >> nb_bits_out) + i64::from(tt.delta_find_state);
self.value = i64::from(ct.state_table[index as usize]);
}
fn flush(&self, bits: &mut BitWriter, ct: &CTable) {
bits.add_bits(self.value as u64, ct.table_log);
}
}
pub(crate) fn compress_using_ctable(
src: &[u8],
ct: &CTable,
nb_states: usize,
dst_capacity: usize,
) -> Option<Vec<u8>> {
if nb_states != 2 && nb_states != 4 {
return None;
}
if src.len() <= nb_states {
return None;
}
let mut bits = BitWriter::with_capacity(dst_capacity);
let mut states = [CState::default(); 4];
let mut ip = src.len();
let mut s = nb_states - (src.len() % nb_states);
for _ in 0..nb_states {
ip -= 1;
states[s % nb_states] = CState::new(ct, src[ip]);
s += 1;
}
for _ in 0..src.len() % nb_states {
ip -= 1;
states[s % nb_states].encode(&mut bits, ct, src[ip]);
s += 1;
}
debug_assert_eq!(s % nb_states, 0);
s = 0;
debug_assert_eq!(ip % nb_states, 0);
while ip > 0 {
for _ in 0..nb_states {
ip -= 1;
states[s % nb_states].encode(&mut bits, ct, src[ip]);
s += 1;
}
}
for _ in 0..nb_states {
states[s % nb_states].flush(&mut bits, ct);
s += 1;
}
bits.close(dst_capacity)
}
pub(crate) fn write_ncount(
norm: &[i16],
max_symbol_value: u32,
table_log: u32,
) -> Result<Vec<u8>, EntropyError> {
if table_log > MAX_TABLELOG {
return Err(EntropyError::TableLogTooLarge);
}
if table_log < MIN_TABLELOG {
return Err(EntropyError::Generic);
}
let mut out = Vec::with_capacity(NCOUNT_BOUND);
let table_size = 1i32 << table_log;
let alphabet_size = max_symbol_value + 1;
let mut bit_stream: u32 = 0;
let mut bit_count: i32 = 0;
let mut symbol: u32 = 0;
let mut previous_is_0 = false;
bit_stream = bit_stream.wrapping_add((table_log - MIN_TABLELOG) << bit_count);
bit_count += 4;
let mut remaining = table_size + 1;
let mut threshold = table_size;
let mut nb_bits = table_log as i32 + 1;
while symbol < alphabet_size && remaining > 1 {
if previous_is_0 {
let mut start = symbol;
while symbol < alphabet_size && norm[symbol as usize] == 0 {
symbol += 1;
}
if symbol == alphabet_size {
break;
}
while symbol >= start + 24 {
start += 24;
bit_stream = bit_stream.wrapping_add(0xFFFFu32 << bit_count);
out.push(bit_stream as u8);
out.push((bit_stream >> 8) as u8);
bit_stream >>= 16;
}
while symbol >= start + 3 {
start += 3;
bit_stream = bit_stream.wrapping_add(3u32 << bit_count);
bit_count += 2;
}
bit_stream = bit_stream.wrapping_add((symbol - start) << bit_count);
bit_count += 2;
if bit_count > 16 {
out.push(bit_stream as u8);
out.push((bit_stream >> 8) as u8);
bit_stream >>= 16;
bit_count -= 16;
}
}
{
let mut count = i32::from(norm[symbol as usize]);
symbol += 1;
let max = (2 * threshold - 1) - remaining;
remaining -= count.abs();
count += 1;
if count >= threshold {
count += max;
}
bit_stream = bit_stream.wrapping_add((count as u32) << bit_count);
bit_count += nb_bits;
bit_count -= i32::from(count < max);
previous_is_0 = count == 1;
if remaining < 1 {
return Err(EntropyError::Generic);
}
while remaining < threshold {
nb_bits -= 1;
threshold >>= 1;
}
}
if bit_count > 16 {
out.push(bit_stream as u8);
out.push((bit_stream >> 8) as u8);
bit_stream >>= 16;
bit_count -= 16;
}
}
if remaining != 1 {
return Err(EntropyError::Generic);
}
out.push(bit_stream as u8);
out.push((bit_stream >> 8) as u8);
let len = out.len() - 2 + ((bit_count + 7) / 8) as usize;
out.truncate(len);
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
fn histogram(src: &[u8]) -> (Vec<u32>, u32) {
let mut count = vec![0u32; 256];
for &b in src {
count[usize::from(b)] += 1;
}
let max = count.iter().rposition(|&c| c != 0).unwrap_or(0) as u32;
(count, max)
}
#[test]
fn bounds_follow_the_macros() {
assert_eq!(block_bound(0), 12);
assert_eq!(block_bound(256), 256 + 2 + 12);
assert_eq!(compress_bound(0), 524);
}
#[test]
fn optimal_table_log_matches_the_reference_on_tiny_and_large_sources() {
assert_eq!(optimal_table_log(DEFAULT_TABLELOG, 2, 1), 11);
assert_eq!(optimal_table_log(DEFAULT_TABLELOG, 100, 3), 5);
assert_eq!(optimal_table_log(DEFAULT_TABLELOG, 1 << 20, 255), 11);
assert_eq!(optimal_table_log(0, 1 << 20, 255), 11);
assert_eq!(optimal_table_log(14, 1 << 20, 255), 12);
assert_eq!(optimal_table_log_internal(6, 1000, 12, 1), 6);
}
#[test]
fn normalised_counts_sum_to_the_table_and_round_trip_through_ncount() {
let src: Vec<u8> = (0..4000u32).map(|i| ((i * i) % 23) as u8).collect();
let (count, max) = histogram(&src);
let log = optimal_table_log(DEFAULT_TABLELOG, src.len(), max);
let mut norm = vec![0i16; 256];
assert_eq!(
normalize_count(&mut norm, log, &count, src.len(), max, true),
Ok(log)
);
let sum: i32 = norm[..=max as usize]
.iter()
.map(|&n| i32::from(n.abs()))
.sum();
assert_eq!(sum, 1 << log);
let header = write_ncount(&norm, max, log).unwrap();
assert!(!header.is_empty() && header.len() <= NCOUNT_BOUND);
assert_eq!(header[0] & 15, (log - MIN_TABLELOG) as u8);
}
#[test]
fn normalisation_rejects_bad_table_logs_and_reports_constants() {
let (count, max) = histogram(b"aaaabbbbcc");
let mut norm = vec![0i16; 256];
assert_eq!(
normalize_count(&mut norm, 4, &count, 10, max, false),
Err(EntropyError::Generic)
);
assert_eq!(
normalize_count(&mut norm, 13, &count, 10, max, false),
Err(EntropyError::TableLogTooLarge)
);
let wide = (0..=255u8).collect::<Vec<_>>();
let (count, max) = histogram(&wide);
assert_eq!(
normalize_count(&mut norm, 5, &count, 256, max, false),
Err(EntropyError::Generic)
);
let (count, max) = histogram(b"zzzzzzzz");
assert_eq!(normalize_count(&mut norm, 6, &count, 8, max, true), Ok(0));
let (count, max) = histogram(&[1u8, 2, 3, 1, 2, 3, 1, 1]);
assert_eq!(
normalize_count(&mut norm, 0, &count, 8, max, true).map(|log| log >= 5),
Ok(true)
);
}
#[test]
fn the_secondary_normalisation_handles_flat_distributions() {
for symbols in [60usize, 200, 255] {
let mut src = Vec::new();
for s in 0..symbols {
src.extend(core::iter::repeat_n(s as u8, 3 + s % 4));
}
let (count, max) = histogram(&src);
let log = optimal_table_log(DEFAULT_TABLELOG, src.len(), max);
let mut norm = vec![0i16; 256];
normalize_count(&mut norm, log, &count, src.len(), max, false).unwrap();
let sum: i32 = norm[..=max as usize]
.iter()
.map(|&n| i32::from(n.abs()))
.sum();
assert_eq!(sum, 1 << log, "{symbols} symbols");
}
}
#[test]
fn m2_normalisation_distributes_rounding_to_the_largest_or_by_weight() {
let count = [1000u32, 3, 3, 3, 2, 1, 1, 1, 1];
let total: usize = count.iter().map(|&c| c as usize).sum();
let mut norm = vec![0i16; 9];
normalize_m2(&mut norm, 6, &count, total, 8, -1).unwrap();
let sum: i32 = norm.iter().map(|&n| i32::from(n.abs())).sum();
assert_eq!(sum, 64);
let count = [2u32, 3, 2, 2];
let mut norm = vec![0i16; 4];
normalize_m2(&mut norm, 5, &count, 9, 3, 1).unwrap();
assert_eq!(norm.iter().map(|&n| i32::from(n)).sum::<i32>(), 32);
let count = [1u32; 6];
let mut norm = vec![0i16; 6];
normalize_m2(&mut norm, 5, &count, 6, 5, 1).unwrap();
assert_eq!(norm.iter().map(|&n| i32::from(n)).sum::<i32>(), 32);
}
#[test]
fn ctables_encode_with_two_and_four_states() {
let src: Vec<u8> = (0..3000u32).map(|i| ((i * 7 + i / 3) % 11) as u8).collect();
let (count, max) = histogram(&src);
let log = optimal_table_log(DEFAULT_TABLELOG, src.len(), max);
let mut norm = vec![0i16; 256];
normalize_count(&mut norm, log, &count, src.len(), max, true).unwrap();
let ct = CTable::build(&norm, max, log).unwrap();
assert_eq!(ct.table_log, log);
for states in [2, 4] {
let out = compress_using_ctable(&src, &ct, states, compress_bound(src.len())).unwrap();
assert!(out.len() < src.len());
assert_ne!(*out.last().unwrap(), 0, "the end mark is in the last byte");
}
assert!(compress_using_ctable(&src, &ct, 3, 1 << 16).is_none());
assert!(compress_using_ctable(&src[..2], &ct, 2, 1 << 16).is_none());
assert!(compress_using_ctable(&src[..3], &ct, 2, 1 << 16).is_some());
assert!(compress_using_ctable(&src[..4], &ct, 4, 1 << 16).is_none());
assert!(compress_using_ctable(&src, &ct, 2, 16).is_none());
}
#[test]
fn ctable_construction_validates_its_arguments() {
let norm = [16i16, 16];
assert!(CTable::build(&norm, 1, 5).is_ok());
assert_eq!(
CTable::build(&norm, 1, 13).unwrap_err(),
EntropyError::TableLogTooLarge
);
assert_eq!(
CTable::build(&norm, 300, 5).unwrap_err(),
EntropyError::TableLogTooLarge
);
assert_eq!(
CTable::build(&norm, 2, 5).unwrap_err(),
EntropyError::Generic
);
let norm = [-1i16, 0, 28, 3];
let ct = CTable::build(&norm, 3, 5).unwrap();
assert_eq!(ct.table_log, 5);
}
#[test]
fn ncount_headers_reject_inconsistent_counts() {
let norm = [16i16, 16];
assert_eq!(
write_ncount(&norm, 1, 13),
Err(EntropyError::TableLogTooLarge)
);
assert_eq!(write_ncount(&norm, 1, 4), Err(EntropyError::Generic));
assert_eq!(write_ncount(&[20i16, 20], 1, 5), Err(EntropyError::Generic));
assert_eq!(write_ncount(&[8i16, 8], 1, 5), Err(EntropyError::Generic));
assert_eq!(write_ncount(&norm, 1, 5), Ok(alloc::vec![0x10, 0x3f]));
let mut sparse = vec![0i16; 100];
sparse[0] = 16;
sparse[99] = 16;
let header = write_ncount(&sparse, 99, 5).unwrap();
assert!(header.len() > 2);
}
}