Revision control
Copy as Markdown
Other Tools
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// Originally written for jxl-oxide.
use crate::bit_reader::BitReader;
use crate::error::{Error, Result};
const LOG_SUM_PROBS: usize = 12;
const SUM_PROBS: u16 = 1 << LOG_SUM_PROBS;
const RLE_MARKER_SYM: u16 = LOG_SUM_PROBS as u16 + 1;
#[derive(Debug)]
struct AnsHistogram {
buckets: Vec<Bucket>,
log_bucket_size: usize,
bucket_mask: u32,
// For optimizing fast-lossless case.
single_symbol: Option<u32>,
}
// log_alphabet_size <= 8 and log_bucket_size <= 7, so u8 is sufficient for symbols and cutoffs.
#[derive(Debug, Copy, Clone)]
#[repr(C)]
struct Bucket {
alias_symbol: u8,
alias_cutoff: u8,
dist: u16,
alias_offset: u16,
alias_dist_xor: u16,
}
impl AnsHistogram {
fn decode_dist_two_symbols(br: &mut BitReader, dist: &mut [u16]) -> Result<usize> {
let table_size = dist.len();
let v0 = Self::read_u8(br)? as usize;
let v1 = Self::read_u8(br)? as usize;
if v0 == v1 {
return Err(Error::InvalidAnsHistogram);
}
let alphabet_size = v0.max(v1) + 1;
if alphabet_size > table_size {
return Err(Error::InvalidAnsHistogram);
}
let prob = br.read(LOG_SUM_PROBS)? as u16;
dist[v0] = prob;
dist[v1] = SUM_PROBS - prob;
Ok(alphabet_size)
}
fn decode_dist_single_symbol(br: &mut BitReader, dist: &mut [u16]) -> Result<usize> {
let table_size = dist.len();
let val = Self::read_u8(br)? as usize;
let alphabet_size = val + 1;
if alphabet_size > table_size {
return Err(Error::InvalidAnsHistogram);
}
dist[val] = SUM_PROBS;
Ok(alphabet_size)
}
fn decode_dist_evenly_distributed(br: &mut BitReader, dist: &mut [u16]) -> Result<usize> {
let table_size = dist.len();
let alphabet_size = Self::read_u8(br)? as usize + 1;
if alphabet_size > table_size {
return Err(Error::InvalidAnsHistogram);
}
let base = SUM_PROBS as usize / alphabet_size;
let remainder = SUM_PROBS as usize % alphabet_size;
dist[0..remainder].fill(base as u16 + 1);
dist[remainder..alphabet_size].fill(base as u16);
Ok(alphabet_size)
}
fn decode_dist_complex(br: &mut BitReader, dist: &mut [u16]) -> Result<usize> {
let table_size = dist.len();
let mut len = 0usize;
while len < 3 {
if br.read(1)? != 0 {
len += 1;
} else {
break;
}
}
let shift = (br.read(len)? + (1 << len) - 1) as i16;
if shift > 13 {
return Err(Error::InvalidAnsHistogram);
}
let alphabet_size = Self::read_u8(br)? as usize + 3;
if alphabet_size > table_size {
return Err(Error::InvalidAnsHistogram);
}
// TODO(tirr-c): This could be an array of length `SUM_PROB / 4` (4 is from the minimum
// value of `repeat_count`). Change if using array is faster.
let mut repeat_ranges = Vec::new();
let mut omit_data = None;
let mut idx = 0;
while idx < alphabet_size {
dist[idx] = Self::read_prefix(br)?;
if dist[idx] == RLE_MARKER_SYM {
let repeat_count = Self::read_u8(br)? as usize + 4;
if idx + repeat_count > alphabet_size {
return Err(Error::InvalidAnsHistogram);
}
repeat_ranges.push(idx..(idx + repeat_count));
idx += repeat_count;
continue;
}
match &mut omit_data {
Some((log, pos)) => {
if dist[idx] > *log {
*log = dist[idx];
*pos = idx;
}
}
data => {
*data = Some((dist[idx], idx));
}
}
idx += 1;
}
let Some((_, omit_pos)) = omit_data else {
return Err(Error::InvalidAnsHistogram);
};
if dist.get(omit_pos + 1) == Some(&RLE_MARKER_SYM) {
return Err(Error::InvalidAnsHistogram);
}
let mut repeat_range_idx = 0usize;
let mut acc = 0;
let mut prev_dist = 0u16;
for (idx, code) in dist.iter_mut().enumerate() {
if repeat_range_idx < repeat_ranges.len()
&& repeat_ranges[repeat_range_idx].start <= idx
{
if repeat_ranges[repeat_range_idx].end == idx {
repeat_range_idx += 1;
} else {
*code = prev_dist;
acc += *code;
// dist[omit_pos] > 0
if acc >= SUM_PROBS {
return Err(Error::InvalidAnsHistogram);
}
continue;
}
}
if *code == 0 {
prev_dist = 0;
continue;
}
if idx == omit_pos {
prev_dist = 0;
continue;
}
if *code > 1 {
let zeros = (*code - 1) as i16;
let bitcount = (shift - ((LOG_SUM_PROBS as i16 - zeros) >> 1)).clamp(0, zeros);
*code = (1 << zeros) + ((br.read(bitcount as usize)? as u16) << (zeros - bitcount));
}
prev_dist = *code;
acc += *code;
// dist[omit_pos] > 0
if acc >= SUM_PROBS {
return Err(Error::InvalidAnsHistogram);
}
}
dist[omit_pos] = SUM_PROBS - acc;
Ok(alphabet_size)
}
fn build_alias_map(alphabet_size: usize, log_bucket_size: usize, dist: &[u16]) -> Vec<Bucket> {
#[derive(Debug)]
struct WorkingBucket {
dist: u16,
alias_symbol: u16,
alias_offset: u16,
alias_cutoff: u16,
}
let bucket_size = 1u16 << log_bucket_size;
let mut buckets: Vec<_> = dist
.iter()
.enumerate()
.map(|(i, &dist)| WorkingBucket {
dist,
alias_symbol: if i < alphabet_size { i as u16 } else { 0 },
alias_offset: 0,
alias_cutoff: dist,
})
.collect();
let mut underfull = Vec::new();
let mut overfull = Vec::new();
for (idx, &WorkingBucket { dist, .. }) in buckets.iter().enumerate() {
match dist.cmp(&bucket_size) {
std::cmp::Ordering::Less => underfull.push(idx),
std::cmp::Ordering::Equal => {}
std::cmp::Ordering::Greater => overfull.push(idx),
}
}
while let (Some(o), Some(u)) = (overfull.pop(), underfull.pop()) {
let by = bucket_size - buckets[u].alias_cutoff;
buckets[o].alias_cutoff -= by;
buckets[u].alias_symbol = o as u16;
buckets[u].alias_offset = buckets[o].alias_cutoff;
match buckets[o].alias_cutoff.cmp(&bucket_size) {
std::cmp::Ordering::Less => underfull.push(o),
std::cmp::Ordering::Equal => {}
std::cmp::Ordering::Greater => overfull.push(o),
}
}
// Assertion failure happens only if `dist` doesn't sum to `SUM_PROB`, which is checked
// before building alias map.
assert!(overfull.is_empty() && underfull.is_empty());
buckets
.iter()
.enumerate()
.map(|(idx, bucket)| {
if bucket.alias_cutoff == bucket_size {
Bucket {
dist: bucket.dist,
alias_symbol: idx as u8,
alias_offset: 0,
alias_cutoff: 0,
alias_dist_xor: 0,
}
} else {
Bucket {
dist: bucket.dist,
alias_symbol: bucket.alias_symbol as u8,
alias_offset: bucket.alias_offset - bucket.alias_cutoff,
alias_cutoff: bucket.alias_cutoff as u8,
alias_dist_xor: bucket.dist ^ buckets[bucket.alias_symbol as usize].dist,
}
}
})
.collect()
}
// log_alphabet_size: 5 + u(2)
pub fn decode(br: &mut BitReader, log_alpha_size: usize) -> Result<Self> {
debug_assert!((5..=8).contains(&log_alpha_size));
let table_size = (1u16 << log_alpha_size) as usize;
// 4 <= log_bucket_size <= 7
let log_bucket_size = LOG_SUM_PROBS - log_alpha_size;
let bucket_size = 1u16 << log_bucket_size;
let bucket_mask = bucket_size as u32 - 1;
let mut dist = vec![0u16; table_size];
let alphabet_size = if br.read(1)? != 0 {
if br.read(1)? != 0 {
Self::decode_dist_two_symbols(br, &mut dist)?
} else {
Self::decode_dist_single_symbol(br, &mut dist)?
}
} else if br.read(1)? != 0 {
Self::decode_dist_evenly_distributed(br, &mut dist)?
} else {
Self::decode_dist_complex(br, &mut dist)?
};
if let Some(single_sym_idx) = dist.iter().position(|&d| d == SUM_PROBS) {
let buckets = dist
.into_iter()
.enumerate()
.map(|(i, dist)| Bucket {
dist,
alias_symbol: single_sym_idx as u8,
alias_offset: bucket_size * i as u16,
alias_cutoff: 0,
alias_dist_xor: dist ^ SUM_PROBS,
})
.collect();
return Ok(Self {
buckets,
log_bucket_size,
bucket_mask,
single_symbol: Some(single_sym_idx as u32),
});
}
Ok(Self {
buckets: Self::build_alias_map(alphabet_size, log_bucket_size, &dist),
log_bucket_size,
bucket_mask,
single_symbol: None,
})
}
fn read_u8(bitstream: &mut BitReader) -> Result<u8> {
Ok(if bitstream.read(1)? != 0 {
let n = bitstream.read(3)?;
((1 << n) + bitstream.read(n as usize)?) as u8
} else {
0
})
}
fn read_prefix(br: &mut BitReader) -> Result<u16> {
// Prefix code lookup table.
#[rustfmt::skip]
const TABLE: [(u8, u8); 128] = [
(10, 3), (12, 7), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4),
(10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4),
(10, 3), ( 0, 5), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4),
(10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4),
(10, 3), (11, 6), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4),
(10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4),
(10, 3), ( 0, 5), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4),
(10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4),
(10, 3), (13, 7), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4),
(10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4),
(10, 3), ( 0, 5), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4),
(10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4),
(10, 3), (11, 6), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4),
(10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4),
(10, 3), ( 0, 5), (7, 3), (3, 4), (6, 3), (8, 3), (9, 3), (5, 4),
(10, 3), ( 4, 4), (7, 3), (1, 4), (6, 3), (8, 3), (9, 3), (2, 4),
];
let index = br.peek(7);
let (sym, bits) = TABLE[index as usize];
br.consume(bits as usize)?;
Ok(sym as u16)
}
}
impl AnsHistogram {
#[inline]
pub fn read(&self, br: &mut BitReader, state: &mut u32) -> u32 {
let idx = *state & 0xfff;
let i = (idx >> self.log_bucket_size) as usize;
let pos = idx & self.bucket_mask;
debug_assert!(self.buckets.len().is_power_of_two());
let bucket = self.buckets[i & (self.buckets.len() - 1)];
let alias_symbol = bucket.alias_symbol as u32;
let alias_cutoff = bucket.alias_cutoff as u32;
let dist = bucket.dist as u32;
let map_to_alias = (pos >= alias_cutoff) as u32;
let offset = (bucket.alias_offset as u32) * map_to_alias;
let dist_xor = (bucket.alias_dist_xor as u32) * map_to_alias;
let dist = dist ^ dist_xor;
let symbol = (alias_symbol * map_to_alias) | (i as u32 * (1 - map_to_alias));
let offset = offset + pos;
let next_state = (*state >> LOG_SUM_PROBS) * dist + offset;
let select_appended = (next_state < (1 << 16)) as u32;
let appended_state = (next_state << 16) | (br.peek(16) as u32);
*state = (appended_state * select_appended) | (next_state * (1 - select_appended));
br.consume_optimistic((16 * select_appended) as usize);
symbol
}
// For optimizing fast-lossless case.
#[inline]
pub fn single_symbol(&self) -> Option<u32> {
self.single_symbol
}
}
#[derive(Debug)]
pub struct AnsCodes {
histograms: Vec<AnsHistogram>,
}
impl AnsCodes {
pub fn decode(num: usize, log_alpha_size: usize, br: &mut BitReader) -> Result<AnsCodes> {
let histograms = (0..num)
.map(|_| AnsHistogram::decode(br, log_alpha_size))
.collect::<Result<_>>()?;
Ok(Self { histograms })
}
pub fn single_symbol(&self, ctx: usize) -> Option<u32> {
self.histograms[ctx].single_symbol()
}
}
#[derive(Debug)]
pub struct AnsReader(u32);
impl AnsReader {
/// Expected final ANS state.
const CHECKSUM: u32 = 0x130000;
pub fn new_unused() -> Self {
Self(Self::CHECKSUM)
}
pub fn init(br: &mut BitReader) -> Result<Self> {
let initial_state = br.read(32)? as u32;
Ok(Self(initial_state))
}
#[inline]
pub fn read(&mut self, codes: &AnsCodes, br: &mut BitReader, ctx: usize) -> u32 {
codes.histograms[ctx].read(br, &mut self.0)
}
pub fn check_final_state(self) -> Result<()> {
if self.0 == Self::CHECKSUM {
Ok(())
} else {
Err(Error::AnsChecksumMismatch)
}
}
pub(super) fn checkpoint(&self) -> Self {
Self(self.0)
}
}
#[cfg(test)]
mod test {
use super::*;
fn validate_buckets(buckets: &[Bucket]) {
let dist_sum: u16 = buckets.iter().map(|bucket| bucket.dist).sum();
assert_eq!(dist_sum, SUM_PROBS);
}
#[test]
fn single_symbol() {
// Single symbol of 20
let mut br = BitReader::new(&[0b00100101, 0b01]);
let histogram = AnsHistogram::decode(&mut br, 5).unwrap();
validate_buckets(&histogram.buckets);
assert_eq!(histogram.buckets[20].dist, SUM_PROBS);
assert_eq!(histogram.single_symbol, Some(20));
// Single symbol of 32 (invalid)
let mut br = BitReader::new(&[0b00101101, 0b000]);
assert!(AnsHistogram::decode(&mut br, 5).is_err());
}
#[test]
fn two_symbols() {
// two symbols of 10 and 20, where the prob of symbol 10 is 256
let mut br = BitReader::new(&[0b10011111, 0b10010010, 0b00000000, 0b00010]);
let histogram = AnsHistogram::decode(&mut br, 5).unwrap();
validate_buckets(&histogram.buckets);
assert_eq!(histogram.buckets[10].dist, 256);
assert_eq!(histogram.buckets[20].dist, SUM_PROBS - 256);
}
#[test]
fn decode_arb() {
arbtest::arbtest(|u| {
let total_len = u.arbitrary_len::<u8>()?;
let mut buf = vec![0u8; total_len];
u.fill_buffer(&mut buf)?;
let mut br = BitReader::new(&buf);
loop {
match AnsHistogram::decode(&mut br, 8) {
Ok(histogram) => validate_buckets(&histogram.buckets),
Err(Error::OutOfBounds(_)) => break,
Err(_) => {}
}
}
Ok(())
});
}
}