aboutsummaryrefslogtreecommitdiff
path: root/bv1/codec/src/huff.rs
diff options
context:
space:
mode:
authormetamuffin <metamuffin@disroot.org>2023-03-09 22:48:33 +0100
committermetamuffin <metamuffin@disroot.org>2023-03-09 22:48:33 +0100
commit680b08e6b9d64284b7992fb52a23e5f891291406 (patch)
treeabe30a9f18be09ef931b4b6357d216f6ba982095 /bv1/codec/src/huff.rs
parentc45f80a14ecd00914eb1d4e8f628b74a713667ba (diff)
downloadvideo-codec-experiments-680b08e6b9d64284b7992fb52a23e5f891291406.tar
video-codec-experiments-680b08e6b9d64284b7992fb52a23e5f891291406.tar.bz2
video-codec-experiments-680b08e6b9d64284b7992fb52a23e5f891291406.tar.zst
rename + readme
Diffstat (limited to 'bv1/codec/src/huff.rs')
-rw-r--r--bv1/codec/src/huff.rs190
1 files changed, 190 insertions, 0 deletions
diff --git a/bv1/codec/src/huff.rs b/bv1/codec/src/huff.rs
new file mode 100644
index 0000000..6d74c42
--- /dev/null
+++ b/bv1/codec/src/huff.rs
@@ -0,0 +1,190 @@
+use std::io::{Read, Result, Write};
+
+#[derive(Debug, Clone)]
+enum HT {
+ Branch(Box<HT>, Box<HT>),
+ Terminal(u8),
+}
+
+pub fn write_huff(buf: &[u8], w: &mut impl Write) -> Result<usize> {
+ let mut w = BitIO::new(w);
+ assert!(buf.len() <= 0xffffff, "huff frame too big");
+ w.wbyte((buf.len() & 0xff) as u8)?;
+ w.wbyte(((buf.len() >> 8) & 0xff) as u8)?;
+ w.wbyte(((buf.len() >> 16) & 0xff) as u8)?;
+
+ let mut probs = [0usize; 256];
+ for b in buf {
+ probs[*b as usize] += 1;
+ }
+ let tree = HT::from_probabilities(probs);
+ let mut table = [0u32; 256];
+ tree.create_lut(&mut table, 1);
+ tree.write(&mut w)?;
+
+
+ for b in buf {
+ let mut k = table[*b as usize];
+ while k != 1 {
+ w.wbit((k & 1) == 1)?;
+ k >>= 1;
+ }
+ }
+
+ w.flush()?;
+ Ok(w.position)
+}
+
+pub fn read_huff(r: &mut impl Read, buf: &mut Vec<u8>) -> Result<usize> {
+ let mut r = BitIO::new(r);
+
+ let mut len = 0usize;
+ len |= r.rbyte()? as usize;
+ len |= (r.rbyte()? as usize) << 8;
+ len |= (r.rbyte()? as usize) << 16;
+
+ let root = HT::read(&mut r)?;
+ let root = match &root {
+ HT::Branch(a, b) => [a, b],
+ _ => panic!("no!"),
+ };
+
+ let mut cursor = root;
+ while buf.len() != len {
+ let v = r.rbit()?;
+ match cursor[v as usize].as_ref() {
+ HT::Branch(a, b) => {
+ cursor = [a, b];
+ }
+ HT::Terminal(n) => {
+ buf.push(*n);
+ cursor = root;
+ }
+ }
+ }
+ Ok(r.position)
+}
+
+impl HT {
+ pub fn from_probabilities(ps: [usize; 256]) -> Self {
+ let mut parts = ps
+ .into_iter()
+ .enumerate()
+ .map(|(n, p)| (p, HT::Terminal(n as u8)))
+ .collect::<Vec<_>>();
+
+ while parts.len() != 1 {
+ parts.sort_by_key(|e| -(e.0 as isize));
+ let ((ap, at), (bp, bt)) = (parts.pop().unwrap(), parts.pop().unwrap());
+ parts.push((ap + bp + 1, HT::Branch(box at, box bt)))
+ }
+ parts[0].1.clone()
+ }
+ pub fn create_lut(&self, table: &mut [u32; 256], mut prefix: u32) {
+ assert!(self.depth() < 30, "too deep! doesnt fit {}", self.depth());
+ match self {
+ HT::Branch(a, b) => {
+ let pz = 32 - prefix.leading_zeros();
+ prefix ^= 1 << pz;
+ prefix ^= 1 << (pz - 1);
+ a.create_lut(table, prefix);
+ prefix ^= 1 << (pz - 1);
+ b.create_lut(table, prefix);
+ }
+ HT::Terminal(n) => {
+ assert_eq!(table[*n as usize], 0);
+ table[*n as usize] = prefix
+ }
+ }
+ }
+ pub fn depth(&self) -> usize {
+ match self {
+ HT::Branch(a, b) => a.depth().max(b.depth()) + 1,
+ HT::Terminal(_) => 0,
+ }
+ }
+ pub fn write(&self, w: &mut BitIO<impl Write>) -> Result<()> {
+ match self {
+ HT::Branch(a, b) => {
+ w.wbit(false)?;
+ a.write(w)?;
+ b.write(w)?;
+ }
+ HT::Terminal(n) => {
+ w.wbit(true)?;
+ w.wbyte(*n)?;
+ }
+ }
+ Ok(())
+ }
+ pub fn read(r: &mut BitIO<impl Read>) -> Result<Self> {
+ match r.rbit()? {
+ false => Ok(Self::Branch(box Self::read(r)?, box Self::read(r)?)),
+ true => Ok(Self::Terminal(r.rbyte()?)),
+ }
+ }
+}
+
+pub struct BitIO<T> {
+ inner: T,
+ byte: u8,
+ position: usize,
+}
+impl<T> BitIO<T> {
+ pub fn new(inner: T) -> Self {
+ Self {
+ inner,
+ byte: 0,
+ position: 0,
+ }
+ }
+}
+impl<T: Write> BitIO<T> {
+ #[inline]
+ pub fn wbit(&mut self, b: bool) -> Result<()> {
+ self.byte <<= 1;
+ self.byte |= b as u8;
+ self.position += 1;
+ if self.position & 0b111 == 0 {
+ self.inner.write_all(&[self.byte])?;
+ }
+ Ok(())
+ }
+ #[inline]
+ pub fn wbyte(&mut self, v: u8) -> Result<()> {
+ for i in 0..8 {
+ self.wbit((v & (1 << i)) != 0)?;
+ }
+ Ok(())
+ }
+ pub fn flush(&mut self) -> Result<()> {
+ while self.position & 0b111 != 0 {
+ self.wbit(false)?;
+ }
+ Ok(())
+ }
+}
+
+impl<T: Read> BitIO<T> {
+ #[inline]
+ pub fn rbit(&mut self) -> Result<bool> {
+ if self.position & 0b111 == 0 {
+ let mut buf = [0];
+ self.inner.read_exact(&mut buf)?;
+ self.byte = buf[0];
+ }
+
+ let v = (self.byte & 0b10000000) != 0;
+ self.byte <<= 1;
+ self.position += 1;
+ Ok(v)
+ }
+ #[inline]
+ pub fn rbyte(&mut self) -> Result<u8> {
+ let mut v = 0u8;
+ for i in 0..8 {
+ v |= (self.rbit()? as u8) << i;
+ }
+ Ok(v)
+ }
+}