473 lines
15 KiB
Rust
473 lines
15 KiB
Rust
//! Range encoder (RFC 6716 §5.1), the exact inverse of the §4.1 range decoder.
|
|
//!
|
|
//! Range-coded symbols are written from the front of the buffer and raw bits from the back; the
|
|
//! bit accounting ([`RangeEncoder::tell`], [`RangeEncoder::tell_frac`]) matches the decoder's, so
|
|
//! every budget decision the decoder takes can be mirrored exactly.
|
|
|
|
const SYM_BITS: u32 = 8;
|
|
const CODE_BITS: u32 = 32;
|
|
const SYM_MAX: u32 = (1 << SYM_BITS) - 1;
|
|
const CODE_SHIFT: u32 = CODE_BITS - SYM_BITS - 1;
|
|
const CODE_TOP: u32 = 1 << (CODE_BITS - 1);
|
|
const CODE_BOT: u32 = CODE_TOP >> SYM_BITS;
|
|
/// Bit resolution of [`RangeEncoder::tell_frac`] (1/8 bit).
|
|
pub const BITRES: u32 = 3;
|
|
|
|
/// Number of bits needed to represent `x` (0 for 0).
|
|
#[inline]
|
|
pub fn ilog(x: u32) -> i32 {
|
|
(32 - x.leading_zeros()) as i32
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct RangeEncoder {
|
|
out: Vec<u8>,
|
|
end: Vec<u8>,
|
|
end_window: u32,
|
|
nend_bits: i32,
|
|
nbits_total: i32,
|
|
rng: u32,
|
|
low: u32,
|
|
rem: i32,
|
|
ext: u32,
|
|
}
|
|
|
|
impl Default for RangeEncoder {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl RangeEncoder {
|
|
pub fn new() -> Self {
|
|
RangeEncoder {
|
|
out: Vec::new(),
|
|
end: Vec::new(),
|
|
end_window: 0,
|
|
nend_bits: 0,
|
|
nbits_total: (CODE_BITS + 1) as i32,
|
|
rng: CODE_TOP,
|
|
low: 0,
|
|
rem: -1,
|
|
ext: 0,
|
|
}
|
|
}
|
|
|
|
/// Current range (equals the decoder's range after the same symbols).
|
|
#[cfg(test)]
|
|
pub fn rng(&self) -> u32 {
|
|
self.rng
|
|
}
|
|
|
|
fn carry_out(&mut self, c: u32) {
|
|
if c != SYM_MAX {
|
|
let carry = c >> SYM_BITS;
|
|
if self.rem >= 0 {
|
|
self.out.push((self.rem as u32 + carry) as u8);
|
|
}
|
|
if self.ext > 0 {
|
|
let sym = (SYM_MAX + carry) & SYM_MAX;
|
|
for _ in 0..self.ext {
|
|
self.out.push(sym as u8);
|
|
}
|
|
self.ext = 0;
|
|
}
|
|
self.rem = (c & SYM_MAX) as i32;
|
|
} else {
|
|
self.ext += 1;
|
|
}
|
|
}
|
|
|
|
#[inline]
|
|
fn normalize(&mut self) {
|
|
while self.rng <= CODE_BOT {
|
|
self.carry_out(self.low >> CODE_SHIFT);
|
|
self.low = (self.low << SYM_BITS) & (CODE_TOP - 1);
|
|
self.rng <<= SYM_BITS;
|
|
self.nbits_total += SYM_BITS as i32;
|
|
}
|
|
}
|
|
|
|
/// Encodes the symbol with cumulative range `[fl, fh)` of total `ft`.
|
|
pub fn encode(&mut self, fl: u32, fh: u32, ft: u32) {
|
|
let r = self.rng / ft;
|
|
if fl > 0 {
|
|
self.low += self.rng - r * (ft - fl);
|
|
self.rng = r * (fh - fl);
|
|
} else {
|
|
self.rng -= r * (ft - fh);
|
|
}
|
|
self.normalize();
|
|
}
|
|
|
|
/// Like [`Self::encode`] with `ft = 1 << bits`.
|
|
pub fn encode_bin(&mut self, fl: u32, fh: u32, bits: u32) {
|
|
let r = self.rng >> bits;
|
|
let ft = 1u32 << bits;
|
|
if fl > 0 {
|
|
self.low += self.rng - r * (ft - fl);
|
|
self.rng = r * (fh - fl);
|
|
} else {
|
|
self.rng -= r * (ft - fh);
|
|
}
|
|
self.normalize();
|
|
}
|
|
|
|
/// Encodes a binary symbol whose probability of being 1 is `1/2^logp`.
|
|
pub fn bit_logp(&mut self, val: bool, logp: u32) {
|
|
let r = self.rng;
|
|
let s = r >> logp;
|
|
let r2 = r - s;
|
|
if val {
|
|
self.low += r2;
|
|
self.rng = s;
|
|
} else {
|
|
self.rng = r2;
|
|
}
|
|
self.normalize();
|
|
}
|
|
|
|
/// Encodes symbol `s` with an "inverse CDF" table (`ft = 1 << ftb`).
|
|
pub fn icdf(&mut self, s: usize, icdf: &[u8], ftb: u32) {
|
|
let r = self.rng >> ftb;
|
|
if s > 0 {
|
|
self.low += self.rng - r * icdf[s - 1] as u32;
|
|
self.rng = r * (icdf[s - 1] as u32 - icdf[s] as u32);
|
|
} else {
|
|
self.rng -= r * icdf[s] as u32;
|
|
}
|
|
self.normalize();
|
|
}
|
|
|
|
/// Encodes a uniformly distributed integer `fl` in `[0, ft)` (`ft >= 2`).
|
|
pub fn uint(&mut self, fl: u32, ft: u32) {
|
|
debug_assert!(ft > 1 && fl < ft);
|
|
let ft1 = ft - 1;
|
|
let ftb = ilog(ft1);
|
|
if ftb > 8 {
|
|
let ftb = (ftb - 8) as u32;
|
|
let top = (ft1 >> ftb) + 1;
|
|
let v = fl >> ftb;
|
|
self.encode(v, v + 1, top);
|
|
self.bits(fl & ((1 << ftb) - 1), ftb);
|
|
} else {
|
|
self.encode(fl, fl + 1, ft);
|
|
}
|
|
}
|
|
|
|
/// Writes `bits` raw bits (0..=25) to the end of the buffer.
|
|
pub fn bits(&mut self, val: u32, bits: u32) {
|
|
if bits == 0 {
|
|
return;
|
|
}
|
|
let mut window = self.end_window;
|
|
let mut used = self.nend_bits;
|
|
if used as u32 + bits > CODE_BITS {
|
|
while used >= SYM_BITS as i32 {
|
|
self.end.push((window & SYM_MAX) as u8);
|
|
window >>= SYM_BITS;
|
|
used -= SYM_BITS as i32;
|
|
}
|
|
}
|
|
window |= val << used;
|
|
used += bits as i32;
|
|
self.end_window = window;
|
|
self.nend_bits = used;
|
|
self.nbits_total += bits as i32;
|
|
}
|
|
|
|
/// Number of bits written so far (rounded up), identical to the decoder's `tell`.
|
|
#[inline]
|
|
pub fn tell(&self) -> i32 {
|
|
self.nbits_total - ilog(self.rng)
|
|
}
|
|
|
|
/// Number of 1/8 bits written so far (rounded up), identical to the decoder's `tell_frac`.
|
|
pub fn tell_frac(&self) -> i32 {
|
|
let nbits = (self.nbits_total as u32) << BITRES;
|
|
let mut l = ilog(self.rng) as u32;
|
|
let mut r = self.rng >> (l - 16);
|
|
for _ in 0..BITRES {
|
|
r = (r * r) >> 15;
|
|
let b = r >> 16;
|
|
l = (l << 1) | b;
|
|
r >>= b;
|
|
}
|
|
nbits.wrapping_sub(l) as i32
|
|
}
|
|
|
|
/// Finalises the stream into a buffer of exactly `size` bytes (range data at the front, raw
|
|
/// bits at the back, zero padding in between; a partial raw-bit byte is merged into the byte
|
|
/// just before the full raw-bit bytes).
|
|
pub fn done(mut self, size: usize) -> Vec<u8> {
|
|
let mut l = CODE_BITS as i32 - ilog(self.rng);
|
|
let mut msk = (CODE_TOP - 1) >> l;
|
|
let mut end = (self.low.wrapping_add(msk)) & !msk;
|
|
if (end | msk) >= self.low.wrapping_add(self.rng) {
|
|
l += 1;
|
|
msk >>= 1;
|
|
end = (self.low.wrapping_add(msk)) & !msk;
|
|
}
|
|
while l > 0 {
|
|
self.carry_out(end >> CODE_SHIFT);
|
|
end = (end << SYM_BITS) & (CODE_TOP - 1);
|
|
l -= SYM_BITS as i32;
|
|
}
|
|
if self.rem >= 0 || self.ext > 0 {
|
|
self.carry_out(0);
|
|
}
|
|
let mut window = self.end_window;
|
|
let mut used = self.nend_bits;
|
|
while used >= SYM_BITS as i32 {
|
|
self.end.push((window & SYM_MAX) as u8);
|
|
window >>= SYM_BITS;
|
|
used -= SYM_BITS as i32;
|
|
}
|
|
let mut buf = vec![0u8; size];
|
|
let full = self.end.len().min(size);
|
|
let room = size - full;
|
|
let n = self.out.len().min(room);
|
|
buf[..n].copy_from_slice(&self.out[..n]);
|
|
for (k, &b) in self.end.iter().take(full).enumerate() {
|
|
buf[size - 1 - k] = b;
|
|
}
|
|
if used > 0 && room > 0 {
|
|
buf[room - 1] |= window as u8;
|
|
}
|
|
buf
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(crate) mod dec {
|
|
//! Minimal range decoder (RFC 6716 §4.1) used to verify the encoder.
|
|
use super::*;
|
|
const CODE_EXTRA: u32 = (CODE_BITS - 2) % SYM_BITS + 1;
|
|
|
|
pub struct RangeDecoder<'a> {
|
|
buf: &'a [u8],
|
|
offs: usize,
|
|
end_offs: usize,
|
|
end_window: u32,
|
|
nend_bits: i32,
|
|
nbits_total: i32,
|
|
rng: u32,
|
|
val: u32,
|
|
rem: u32,
|
|
}
|
|
|
|
impl<'a> RangeDecoder<'a> {
|
|
pub fn new(buf: &'a [u8]) -> Self {
|
|
let mut d = RangeDecoder {
|
|
buf,
|
|
offs: 0,
|
|
end_offs: 0,
|
|
end_window: 0,
|
|
nend_bits: 0,
|
|
nbits_total: (CODE_BITS + 1 - ((CODE_BITS - CODE_EXTRA) / SYM_BITS) * SYM_BITS) as i32,
|
|
rng: 1 << CODE_EXTRA,
|
|
val: 0,
|
|
rem: 0,
|
|
};
|
|
d.rem = d.read_byte();
|
|
d.val = d.rng - 1 - (d.rem >> (SYM_BITS - CODE_EXTRA));
|
|
d.normalize();
|
|
d
|
|
}
|
|
fn read_byte(&mut self) -> u32 {
|
|
if self.offs < self.buf.len() {
|
|
self.offs += 1;
|
|
self.buf[self.offs - 1] as u32
|
|
} else {
|
|
0
|
|
}
|
|
}
|
|
fn read_byte_from_end(&mut self) -> u32 {
|
|
if self.end_offs < self.buf.len() {
|
|
self.end_offs += 1;
|
|
self.buf[self.buf.len() - self.end_offs] as u32
|
|
} else {
|
|
0
|
|
}
|
|
}
|
|
fn normalize(&mut self) {
|
|
while self.rng <= CODE_BOT {
|
|
self.nbits_total += SYM_BITS as i32;
|
|
self.rng <<= SYM_BITS;
|
|
let sym = self.rem;
|
|
self.rem = self.read_byte();
|
|
let sym = ((sym << SYM_BITS) | self.rem) >> (SYM_BITS - CODE_EXTRA);
|
|
self.val = ((self.val << SYM_BITS).wrapping_add(SYM_MAX & !sym)) & (CODE_TOP - 1);
|
|
}
|
|
}
|
|
pub fn rng(&self) -> u32 {
|
|
self.rng
|
|
}
|
|
pub fn decode(&mut self, ft: u32) -> u32 {
|
|
let ext = self.rng / ft;
|
|
let s = self.val / ext;
|
|
ft - (s + 1).min(ft)
|
|
}
|
|
pub fn update(&mut self, fl: u32, fh: u32, ft: u32) {
|
|
let ext = self.rng / ft;
|
|
let s = ext * (ft - fh);
|
|
self.val = self.val.wrapping_sub(s);
|
|
self.rng = if fl > 0 { ext * (fh - fl) } else { self.rng - s };
|
|
self.normalize();
|
|
}
|
|
pub fn bit_logp(&mut self, logp: u32) -> bool {
|
|
let r = self.rng;
|
|
let d = self.val;
|
|
let s = r >> logp;
|
|
let ret = d < s;
|
|
if !ret {
|
|
self.val = d - s;
|
|
}
|
|
self.rng = if ret { s } else { r - s };
|
|
self.normalize();
|
|
ret
|
|
}
|
|
pub fn icdf(&mut self, icdf: &[u8], ftb: u32) -> usize {
|
|
let mut s = self.rng;
|
|
let d = self.val;
|
|
let r = s >> ftb;
|
|
let mut ret = 0usize;
|
|
let mut t;
|
|
loop {
|
|
t = s;
|
|
s = r * icdf[ret] as u32;
|
|
if d >= s {
|
|
break;
|
|
}
|
|
ret += 1;
|
|
}
|
|
self.val = d - s;
|
|
self.rng = t - s;
|
|
self.normalize();
|
|
ret
|
|
}
|
|
pub fn uint(&mut self, ft: u32) -> u32 {
|
|
let ft1 = ft - 1;
|
|
let ftb = ilog(ft1);
|
|
if ftb > 8 {
|
|
let ftb = (ftb - 8) as u32;
|
|
let top = (ft1 >> ftb) + 1;
|
|
let s = self.decode(top);
|
|
self.update(s, s + 1, top);
|
|
(s << ftb) | self.bits(ftb)
|
|
} else {
|
|
let s = self.decode(ft);
|
|
self.update(s, s + 1, ft);
|
|
s
|
|
}
|
|
}
|
|
pub fn bits(&mut self, bits: u32) -> u32 {
|
|
let mut window = self.end_window;
|
|
let mut available = self.nend_bits;
|
|
if (available as u32) < bits {
|
|
loop {
|
|
window |= self.read_byte_from_end() << available;
|
|
available += SYM_BITS as i32;
|
|
if available > (CODE_BITS - SYM_BITS) as i32 {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
let ret = window & ((1u32 << bits) - 1);
|
|
self.end_window = window >> bits;
|
|
self.nend_bits = available - bits as i32;
|
|
self.nbits_total += bits as i32;
|
|
ret
|
|
}
|
|
pub fn tell(&self) -> i32 {
|
|
self.nbits_total - ilog(self.rng)
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::dec::RangeDecoder;
|
|
use super::*;
|
|
|
|
struct Lcg(u32);
|
|
impl Lcg {
|
|
fn next(&mut self) -> u32 {
|
|
self.0 = self.0.wrapping_mul(1664525).wrapping_add(1013904223);
|
|
self.0 >> 8
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn starts_at_one_bit() {
|
|
let e = RangeEncoder::new();
|
|
assert_eq!(e.tell(), 1);
|
|
assert_eq!(e.tell_frac(), 8);
|
|
}
|
|
|
|
#[test]
|
|
fn roundtrip_mixed_symbols_and_tell_matches() {
|
|
let mut rng = Lcg(7);
|
|
let icdf: [u8; 4] = [200, 120, 30, 0];
|
|
for trial in 0..300 {
|
|
let n = 1 + (rng.next() % 300) as usize;
|
|
let mut ops = Vec::new();
|
|
for _ in 0..n {
|
|
match rng.next() % 5 {
|
|
0 => {
|
|
let ft = 2 + rng.next() % 1000;
|
|
ops.push((0u8, rng.next() % ft, ft));
|
|
}
|
|
1 => {
|
|
let logp = 1 + rng.next() % 15;
|
|
ops.push((1, rng.next().is_multiple_of(5) as u32, logp));
|
|
}
|
|
2 => {
|
|
let bits = 1 + rng.next() % 24;
|
|
ops.push((2, rng.next() & ((1 << bits) - 1), bits));
|
|
}
|
|
3 => {
|
|
let ft = 2 + rng.next() % 100_000_000;
|
|
ops.push((3, rng.next() % ft, ft));
|
|
}
|
|
_ => ops.push((4, rng.next() % 4, 0)),
|
|
}
|
|
}
|
|
let mut e = RangeEncoder::new();
|
|
let mut tells = Vec::new();
|
|
for &(k, v, p) in &ops {
|
|
match k {
|
|
0 => e.encode(v, v + 1, p),
|
|
1 => e.bit_logp(v != 0, p),
|
|
2 => e.bits(v, p),
|
|
3 => e.uint(v, p),
|
|
_ => e.icdf(v as usize, &icdf, 8),
|
|
}
|
|
tells.push(e.tell());
|
|
}
|
|
// Exactly the number of bytes `tell` says were used must suffice.
|
|
let size = (e.tell() as usize).div_ceil(8);
|
|
let final_rng = e.rng();
|
|
let buf = e.done(size);
|
|
assert_eq!(buf.len(), size);
|
|
let mut d = RangeDecoder::new(&buf);
|
|
for (i, &(k, v, p)) in ops.iter().enumerate() {
|
|
let got = match k {
|
|
0 => {
|
|
let s = d.decode(p);
|
|
d.update(s, s + 1, p);
|
|
s
|
|
}
|
|
1 => d.bit_logp(p) as u32,
|
|
2 => d.bits(p),
|
|
3 => d.uint(p),
|
|
_ => d.icdf(&icdf, 8) as u32,
|
|
};
|
|
assert_eq!(got, v, "trial {trial} op {i} kind {k}");
|
|
assert_eq!(d.tell(), tells[i], "tell mismatch trial {trial} op {i}");
|
|
}
|
|
assert_eq!(d.rng(), final_rng);
|
|
}
|
|
}
|
|
}
|