pub use decoding::Decoder;
pub use encoding::quantize_weights;
pub use encoding::Encoder;
pub use spec::Spec;
mod decoding;
mod encoding;
mod spec;
pub(crate) type AnsState = u32;
pub(crate) type Symbol = u32;
#[cfg(test)]
mod tests {
use crate::ans::spec::Spec;
use crate::ans::{AnsState, Decoder, Encoder, Symbol};
use crate::bit_reader::BitReader;
use crate::bit_writer::BitWriter;
use crate::bits;
use crate::constants::Bitlen;
use crate::errors::PcoResult;
fn assert_recovers(spec: &Spec, symbols: Vec<Symbol>, expected_byte_len: usize) -> PcoResult<()> {
let encoder = Encoder::new(spec);
let mut state = encoder.default_state();
let mut to_write = Vec::new();
for &symbol in symbols.iter().rev() {
let (new_state, bitlen) = encoder.encode(state, symbol);
to_write.push((state, bitlen));
state = new_state;
}
let mut compressed = Vec::new();
let mut writer = BitWriter::new(&mut compressed, 10);
for (val, bitlen) in to_write.into_iter().rev() {
unsafe { writer.write_uint(bits::lowest_bits_fast(val, bitlen), bitlen) };
writer.flush()?;
}
writer.finish_byte();
writer.flush()?;
drop(writer);
assert_eq!(compressed.len(), expected_byte_len);
let final_state = state;
let table_size = 1 << encoder.size_log();
compressed.extend(&vec![0; 100]);
let mut reader = BitReader::new(&compressed, expected_byte_len, 0);
let decoder = Decoder::new(spec, &vec![]);
let mut decoded = Vec::new();
let mut state_idx = final_state - table_size;
for &symbol in symbols.iter() {
let node = &decoder.nodes[state_idx as usize];
decoded.push(symbol as Symbol);
state_idx = node.next_state_idx_base as AnsState
+ unsafe { reader.read_uint::<AnsState>(node.bits_to_read as Bitlen) };
}
assert_eq!(decoded, symbols);
Ok(())
}
#[test]
fn ans_encoder_decoder_dense() -> PcoResult<()> {
let spec = Spec {
size_log: 3,
state_symbols: vec![0, 1, 2, 0, 1, 2, 0, 1],
symbol_weights: vec![3, 3, 2],
};
let symbols = vec![2, 0, 1, 1, 1, 0, 0, 1, 2];
assert_recovers(&spec, symbols, 2)?;
let mut symbols = Vec::new();
for _ in 0..200 {
symbols.push(0);
symbols.push(1);
symbols.push(2);
}
assert_recovers(&spec, symbols, 125)?;
Ok(())
}
#[test]
fn ans_encoder_decoder_sparse() -> PcoResult<()> {
let spec = Spec {
size_log: 3,
state_symbols: vec![0, 0, 0, 0, 0, 0, 0, 1],
symbol_weights: vec![7, 1],
};
let mut symbols = Vec::new();
for _ in 0..100 {
for _ in 0..7 {
symbols.push(0);
}
symbols.push(1);
}
assert_recovers(&spec, symbols, 50)?;
Ok(())
}
}