#![cfg_attr(docsrs, doc(cfg(feature = "arc_squeeze")))]
extern crate alloc;
use alloc::vec;
use alloc::vec::Vec;
use crate::error::Error;
use crate::traits::{Algorithm, RawDecoder, RawEncoder, RawProgress};
#[derive(Debug, Clone, Copy, Default)]
pub struct ArcSqueeze;
impl Algorithm for ArcSqueeze {
const NAME: &'static str = "squeeze";
type Encoder = Encoder;
type Decoder = Decoder;
type EncoderConfig = ();
type DecoderConfig = ();
fn encoder_with(_: ()) -> Encoder {
Encoder::new()
}
fn decoder_with(_: ()) -> Decoder {
Decoder::new()
}
}
const RLE_FLAG: u8 = 0x90;
const SPEOF: u16 = 256;
const NUM_SYMBOLS: usize = 257;
const MAX_NODES: usize = 2 * NUM_SYMBOLS;
type Node = (i16, i16);
type Code = (u32, u8);
fn rle_encode(input: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(input.len());
let mut i = 0usize;
while i < input.len() {
let b = input[i];
out.push(b);
if b == RLE_FLAG {
out.push(0);
i += 1;
continue;
}
let mut run = 1usize;
while i + run < input.len() && input[i + run] == b && run < 255 {
run += 1;
}
if run >= 3 {
out.push(RLE_FLAG);
out.push(run as u8);
i += run;
} else {
i += 1;
}
}
out
}
#[derive(Debug, Default)]
struct RleDecoder {
last: u8,
have_last: bool,
awaiting_count: bool,
}
impl RleDecoder {
fn reset(&mut self) {
self.last = 0;
self.have_last = false;
self.awaiting_count = false;
}
fn feed(&mut self, b: u8, out: &mut Vec<u8>) -> Result<(), Error> {
if self.awaiting_count {
self.awaiting_count = false;
if b == 0 {
out.push(RLE_FLAG);
self.last = RLE_FLAG;
self.have_last = true;
} else {
if !self.have_last {
return Err(Error::Corrupt);
}
for _ in 1..b {
out.push(self.last);
}
}
return Ok(());
}
if b == RLE_FLAG {
self.awaiting_count = true;
} else {
out.push(b);
self.last = b;
self.have_last = true;
}
Ok(())
}
fn pending(&self) -> bool {
self.awaiting_count
}
}
fn build_tree(freq: &[u32; NUM_SYMBOLS]) -> (Vec<Node>, Vec<Code>) {
#[derive(Clone, Copy)]
struct Item {
weight: u64,
node: i32,
order: u32,
}
let mut heap: Vec<Item> = Vec::new();
let mut order = 0u32;
for (sym, &f) in freq.iter().enumerate() {
if f > 0 {
heap.push(Item {
weight: f as u64,
node: -(sym as i32) - 1,
order,
});
order += 1;
}
}
let mut nodes: Vec<(i32, i32)> = Vec::new();
if heap.len() == 1 {
let only = heap[0].node;
nodes.push((only, only));
return finish_tree(nodes);
}
fn pop_min(heap: &mut Vec<Item>) -> Item {
let mut best = 0usize;
for i in 1..heap.len() {
if heap[i].weight < heap[best].weight
|| (heap[i].weight == heap[best].weight && heap[i].order < heap[best].order)
{
best = i;
}
}
heap.swap_remove(best)
}
while heap.len() > 1 {
let a = pop_min(&mut heap);
let b = pop_min(&mut heap);
let idx = nodes.len() as i32;
nodes.push((a.node, b.node));
heap.push(Item {
weight: a.weight.saturating_add(b.weight),
node: idx,
order,
});
order += 1;
}
finish_tree(nodes)
}
fn finish_tree(nodes: Vec<(i32, i32)>) -> (Vec<Node>, Vec<Code>) {
let root = nodes.len() as i32 - 1;
let mut new_table: Vec<(i16, i16)> = Vec::new();
let mut codes: Vec<(u32, u8)> = vec![(0u32, 0u8); NUM_SYMBOLS];
let mut remap: Vec<i32> = vec![-1; nodes.len().max(1)];
let mut stack: Vec<i32> = Vec::new();
if root >= 0 {
remap[root as usize] = 0;
new_table.push((0, 0)); stack.push(root);
}
while let Some(old) = stack.pop() {
let (l, r) = nodes[old as usize];
let mut resolve = |child: i32, table: &mut Vec<(i16, i16)>, stack: &mut Vec<i32>| -> i16 {
if child < 0 {
child as i16 } else {
if remap[child as usize] < 0 {
let nid = table.len() as i32;
remap[child as usize] = nid;
table.push((0, 0));
stack.push(child);
}
remap[child as usize] as i16
}
};
let nl = resolve(l, &mut new_table, &mut stack);
let nr = resolve(r, &mut new_table, &mut stack);
let nid = remap[old as usize] as usize;
new_table[nid] = (nl, nr);
}
fn assign(table: &[(i16, i16)], node: i16, bits: u32, len: u8, codes: &mut [(u32, u8)]) {
if node < 0 {
let sym = (-(node as i32) - 1) as usize;
if sym < codes.len() {
codes[sym] = (bits, len);
}
return;
}
let (l, r) = table[node as usize];
if len < 32 {
assign(table, l, bits, len + 1, codes);
assign(table, r, bits | (1u32 << len), len + 1, codes);
}
}
if !new_table.is_empty() {
assign(&new_table, 0, 0, 0, &mut codes);
}
(new_table, codes)
}
#[derive(Debug)]
pub struct Encoder {
input: Vec<u8>,
out: Vec<u8>,
out_head: usize,
built: bool,
completed: bool,
}
impl Encoder {
pub fn new() -> Self {
Self {
input: Vec::new(),
out: Vec::new(),
out_head: 0,
built: false,
completed: false,
}
}
fn build(&mut self) {
let rle = rle_encode(&self.input);
let mut freq = [0u32; NUM_SYMBOLS];
for &b in &rle {
freq[b as usize] += 1;
}
freq[SPEOF as usize] = 1;
let (table, codes) = build_tree(&freq);
let mut out = Vec::new();
let n = table.len() as u16;
out.push((n & 0xFF) as u8);
out.push((n >> 8) as u8);
for &(l, r) in &table {
let lu = l as u16;
let ru = r as u16;
out.push((lu & 0xFF) as u8);
out.push((lu >> 8) as u8);
out.push((ru & 0xFF) as u8);
out.push((ru >> 8) as u8);
}
let mut bit_acc: u32 = 0;
let mut bit_count: u8 = 0;
let mut emit = |bits: u32, len: u8, out: &mut Vec<u8>| {
bit_acc |= bits << bit_count;
bit_count += len;
while bit_count >= 8 {
out.push(bit_acc as u8);
bit_acc >>= 8;
bit_count -= 8;
}
};
for &b in &rle {
let (bits, len) = codes[b as usize];
emit(bits, len, &mut out);
}
let (eb, el) = codes[SPEOF as usize];
emit(eb, el, &mut out);
if bit_count > 0 {
out.push(bit_acc as u8);
}
self.out = out;
self.built = true;
}
fn drain(&mut self, output: &mut [u8]) -> usize {
let available = self.out.len() - self.out_head;
let n = available.min(output.len());
output[..n].copy_from_slice(&self.out[self.out_head..self.out_head + n]);
self.out_head += n;
n
}
}
impl Default for Encoder {
fn default() -> Self {
Self::new()
}
}
impl RawEncoder for Encoder {
fn raw_encode(&mut self, input: &[u8], _output: &mut [u8]) -> Result<RawProgress, Error> {
self.input.extend_from_slice(input);
Ok(RawProgress {
consumed: input.len(),
written: 0,
done: false,
})
}
fn raw_finish(&mut self, output: &mut [u8]) -> Result<RawProgress, Error> {
if self.completed {
return Ok(RawProgress {
consumed: 0,
written: 0,
done: true,
});
}
if !self.built {
self.build();
}
let written = self.drain(output);
let done = self.out_head >= self.out.len();
if done {
self.completed = true;
}
Ok(RawProgress {
consumed: 0,
written,
done,
})
}
fn raw_reset(&mut self) {
self.input.clear();
self.out.clear();
self.out_head = 0;
self.built = false;
self.completed = false;
}
}
#[derive(Debug)]
pub struct Decoder {
in_buf: Vec<u8>,
in_pos: usize,
num_nodes: Option<usize>,
table: Vec<(i16, i16)>,
header_done: bool,
empty_stream: bool,
bit_acc: u32,
bit_count: u8,
cur_node: i16,
rle: RleDecoder,
emit_buf: Vec<u8>,
emit_head: usize,
eos: bool,
completed: bool,
}
impl Decoder {
pub fn new() -> Self {
Self {
in_buf: Vec::new(),
in_pos: 0,
num_nodes: None,
table: Vec::new(),
header_done: false,
empty_stream: false,
bit_acc: 0,
bit_count: 0,
cur_node: 0,
rle: RleDecoder::default(),
emit_buf: Vec::new(),
emit_head: 0,
eos: false,
completed: false,
}
}
fn parse_header(&mut self) -> Result<bool, Error> {
if self.header_done {
return Ok(true);
}
if self.num_nodes.is_none() {
if self.in_buf.len() - self.in_pos < 2 {
return Ok(false);
}
let n = (self.in_buf[self.in_pos] as usize)
| ((self.in_buf[self.in_pos + 1] as usize) << 8);
self.in_pos += 2;
if n > MAX_NODES {
return Err(Error::InvalidHuffmanTree);
}
self.num_nodes = Some(n);
if n == 0 {
self.empty_stream = true;
self.header_done = true;
return Ok(true);
}
}
let n = self.num_nodes.unwrap();
if self.in_buf.len() - self.in_pos < 4 * n {
return Ok(false);
}
self.table.clear();
self.table.reserve(n);
for _ in 0..n {
let l =
(self.in_buf[self.in_pos] as i16) | ((self.in_buf[self.in_pos + 1] as i16) << 8);
let r = (self.in_buf[self.in_pos + 2] as i16)
| ((self.in_buf[self.in_pos + 3] as i16) << 8);
self.in_pos += 4;
self.table.push((l, r));
}
for &(l, r) in &self.table {
for child in [l, r] {
if child >= 0 {
if child as usize >= n {
return Err(Error::InvalidHuffmanTree);
}
} else {
let val = (-(child as i32) - 1) as i64;
if val > SPEOF as i64 {
return Err(Error::InvalidHuffmanTree);
}
}
}
}
self.header_done = true;
Ok(true)
}
fn pump(&mut self) -> Result<(), Error> {
if self.empty_stream {
self.eos = true;
return Ok(());
}
loop {
if self.eos {
return Ok(());
}
if self.emit_buf.len() - self.emit_head > 16 * 1024 {
return Ok(());
}
let leaf = loop {
let node = self.cur_node;
if node < 0 || node as usize >= self.table.len() {
return Err(Error::InvalidHuffmanTree);
}
if self.bit_count == 0 {
if self.in_pos >= self.in_buf.len() {
return Ok(()); }
self.bit_acc = self.in_buf[self.in_pos] as u32;
self.bit_count = 8;
self.in_pos += 1;
}
let bit = self.bit_acc & 1;
self.bit_acc >>= 1;
self.bit_count -= 1;
let (l, r) = self.table[node as usize];
let next = if bit == 0 { l } else { r };
if next < 0 {
self.cur_node = 0;
break (-(next as i32) - 1) as u16;
} else {
self.cur_node = next;
}
};
if leaf == SPEOF {
self.eos = true;
return Ok(());
}
if leaf > 255 {
return Err(Error::InvalidHuffmanTree);
}
self.rle.feed(leaf as u8, &mut self.emit_buf)?;
}
}
fn drain_emit(&mut self, out: &mut [u8]) -> usize {
let available = self.emit_buf.len() - self.emit_head;
let n = available.min(out.len());
out[..n].copy_from_slice(&self.emit_buf[self.emit_head..self.emit_head + n]);
self.emit_head += n;
if self.emit_head == self.emit_buf.len() {
self.emit_buf.clear();
self.emit_head = 0;
}
n
}
fn compact_in_buf(&mut self) {
if self.in_pos >= 64 * 1024 {
self.in_buf.drain(0..self.in_pos);
self.in_pos = 0;
}
}
}
impl Default for Decoder {
fn default() -> Self {
Self::new()
}
}
impl RawDecoder for Decoder {
fn raw_decode(&mut self, input: &[u8], output: &mut [u8]) -> Result<RawProgress, Error> {
let consumed = input.len();
self.in_buf.extend_from_slice(input);
let mut written = 0usize;
if !self.header_done && !self.parse_header()? {
return Ok(RawProgress {
consumed,
written,
done: false,
});
}
if self.emit_head < self.emit_buf.len() {
written += self.drain_emit(&mut output[written..]);
}
while written < output.len() {
if self.emit_head >= self.emit_buf.len() {
self.pump()?;
if self.emit_head >= self.emit_buf.len() {
break;
}
}
written += self.drain_emit(&mut output[written..]);
}
self.compact_in_buf();
let done = self.eos && self.emit_head >= self.emit_buf.len();
if done {
self.completed = true;
}
Ok(RawProgress {
consumed,
written,
done,
})
}
fn raw_finish(&mut self, output: &mut [u8]) -> Result<RawProgress, Error> {
if self.completed {
return Ok(RawProgress {
consumed: 0,
written: 0,
done: true,
});
}
if !self.header_done && !self.parse_header()? {
return Err(Error::UnexpectedEnd);
}
let mut written = 0usize;
if self.emit_head < self.emit_buf.len() {
written += self.drain_emit(&mut output[written..]);
}
while written < output.len() {
if self.emit_head >= self.emit_buf.len() {
self.pump()?;
if self.emit_head >= self.emit_buf.len() {
break;
}
}
written += self.drain_emit(&mut output[written..]);
}
if self.emit_head < self.emit_buf.len() {
return Ok(RawProgress {
consumed: 0,
written,
done: false,
});
}
if !self.eos {
return Err(Error::UnexpectedEnd);
}
if self.rle.pending() {
return Err(Error::UnexpectedEnd);
}
self.completed = true;
Ok(RawProgress {
consumed: 0,
written,
done: true,
})
}
fn raw_reset(&mut self) {
self.in_buf.clear();
self.in_pos = 0;
self.num_nodes = None;
self.table.clear();
self.header_done = false;
self.empty_stream = false;
self.bit_acc = 0;
self.bit_count = 0;
self.cur_node = 0;
self.rle.reset();
self.emit_buf.clear();
self.emit_head = 0;
self.eos = false;
self.completed = false;
}
}