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)
+    }
+}