add huffman to test2

4 files changed, 188 insertions, 5 deletions

diff --git a/test2/src/decode.rs b/test2/src/decode.rs
index da0538b..016a9d1 100644
--- a/test2/src/decode.rs
+++ b/test2/src/decode.rs
@@ -1,4 +1,4 @@
-use crate::{BLOCK_SIZE, Frame};
+use crate::{BLOCK_SIZE, Frame, huffman::decode_huff};
 use framework::BitstreamFilter;
 use glam::{IVec2, ivec2};
 
@@ -19,7 +19,8 @@ impl BitstreamFilter for Dec {
     }
 
     fn process_block(&mut self, a: Vec<u8>) -> Vec<u8> {
-        let mut buf = a.as_slice();
+        let buf = decode_huff(a);
+        let mut buf = buf.as_slice();
         let mut frame = Frame::new(self.res);
         for by in 0..frame.res.y / BLOCK_SIZE {
             for bx in 0..frame.res.x / BLOCK_SIZE {
diff --git a/test2/src/encode.rs b/test2/src/encode.rs
index 7eb82e0..53ca8a9 100644
--- a/test2/src/encode.rs
+++ b/test2/src/encode.rs
@@ -1,4 +1,4 @@
-use crate::{BLOCK_SIZE, Frame};
+use crate::{BLOCK_SIZE, Frame, huffman::encode_huff};
 use framework::BitstreamFilter;
 use glam::{IVec2, ivec2};
 
@@ -79,7 +79,7 @@ impl BitstreamFilter for Enc {
                     if neg {
                         d = (-(d as i8)) as u8;
                     }
-                    d &= 0xff << (6 - d.leading_zeros()); // only keep 3 msb
+                    d &= 0xff << (7 - d.leading_zeros()); // only keep 3 msb
                     if neg {
                         d = (-(d as i8)) as u8;
                     }
@@ -95,6 +95,6 @@ impl BitstreamFilter for Enc {
         }
 
         self.last = frame.clone();
-        out
+        encode_huff(&out)
     }
 }
diff --git a/test2/src/huffman.rs b/test2/src/huffman.rs
new file mode 100644
index 0000000..40a6730
--- /dev/null
+++ b/test2/src/huffman.rs
@@ -0,0 +1,181 @@
+#[derive(Debug, Clone)]
+enum HT {
+    Branch(Box<HT>, Box<HT>),
+    Terminal(u8),
+}
+
+pub fn encode_huff(buf: &[u8]) -> Vec<u8> {
+    let mut w = BitIO::new(Vec::new());
+    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();
+    w.buffer
+}
+
+pub fn decode_huff(r: Vec<u8>) -> Vec<u8> {
+    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;
+    let mut buf = Vec::new();
+    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;
+            }
+        }
+    }
+    buf
+}
+
+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::new(at), Box::new(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) {
+        match self {
+            HT::Branch(a, b) => {
+                w.wbit(false);
+                a.write(w);
+                b.write(w);
+            }
+            HT::Terminal(n) => {
+                w.wbit(true);
+                w.wbyte(*n);
+            }
+        }
+    }
+    pub fn read(r: &mut BitIO) -> Self {
+        match r.rbit() {
+            false => Self::Branch(Box::new(Self::read(r)), Box::new(Self::read(r))),
+            true => Self::Terminal(r.rbyte()),
+        }
+    }
+}
+
+pub struct BitIO {
+    buffer: Vec<u8>,
+    byte: u8,
+    position: usize,
+}
+impl BitIO {
+    pub fn new(inner: Vec<u8>) -> Self {
+        Self {
+            buffer: inner,
+            byte: 0,
+            position: 0,
+        }
+    }
+    #[inline]
+    pub fn wbit(&mut self, b: bool) {
+        self.byte <<= 1;
+        self.byte |= b as u8;
+        self.position += 1;
+        if self.position & 0b111 == 0 {
+            self.buffer.push(self.byte)
+        }
+    }
+    #[inline]
+    pub fn wbyte(&mut self, v: u8) {
+        for i in 0..8 {
+            self.wbit((v & (1 << i)) != 0);
+        }
+    }
+
+    #[inline]
+    pub fn flush(&mut self) {
+        while self.position & 0b111 != 0 {
+            self.wbit(false);
+        }
+    }
+
+    #[inline]
+    pub fn rbit(&mut self) -> bool {
+        if self.position & 0b111 == 0 {
+            self.byte = self.buffer[self.position >> 3];
+        }
+
+        let v = (self.byte & 0b10000000) != 0;
+        self.byte <<= 1;
+        self.position += 1;
+        v
+    }
+
+    #[inline]
+    pub fn rbyte(&mut self) -> u8 {
+        let mut v = 0u8;
+        for i in 0..8 {
+            v |= (self.rbit() as u8) << i;
+        }
+        v
+    }
+}
diff --git a/test2/src/main.rs b/test2/src/main.rs
index aca978e..f145202 100644
--- a/test2/src/main.rs
+++ b/test2/src/main.rs
@@ -1,5 +1,6 @@
 pub mod decode;
 pub mod encode;
+pub mod huffman;
 
 use decode::Dec;
 use encode::Enc;