huff doesnt work

3 files changed, 197 insertions, 16 deletions

diff --git a/lvc/src/bin/main.rs b/lvc/src/bin/main.rs
index d5c9bbe..e579f72 100644
--- a/lvc/src/bin/main.rs
+++ b/lvc/src/bin/main.rs
@@ -3,6 +3,7 @@ use lvc::{
     debug::draw_debug,
     decode::decode,
     encode::{encode, EncodeConfig},
+    huff::{read_huffman, write_huffman},
     Block, Frame, Pixel, PixelValue, View, P2,
 };
 use std::{
@@ -58,7 +59,15 @@ fn main() {
             let mut stdout = BufWriter::new(stdout());
 
             loop {
-                let b = Block::read(&mut stdin, View::all(size)).unwrap();
+                let mut buf = vec![];
+                read_huffman(&mut stdin, &mut buf).unwrap();
+
+                eprintln!("{:?}", &buf[..20]);
+                let b = {
+                    let mut buf = std::io::Cursor::new(&mut buf);
+                    Block::read(&mut buf, View::all(size)).unwrap()
+                };
+
                 decode(&last_frame, &mut frame, View::all(size), &b);
 
                 if let Some(debug_frame) = &mut debug_frame {
@@ -94,12 +103,34 @@ fn main() {
 
                 let t = Instant::now();
                 let b: Block = encode(&last_frame, &frame, View::all(size), &config);
-                eprintln!("frame {frame_number} took {:?}", t.elapsed());
-
-                b.write(&mut stdout);
+                let time_encode = t.elapsed();
 
+                let t = Instant::now();
                 decode(&last_frame, &mut frame, View::all(size), &b);
                 last_frame = frame;
+                let time_decode = t.elapsed();
+
+                let mut buf = vec![];
+                let mut bufw = std::io::Cursor::new(&mut buf);
+                b.write(&mut bufw).unwrap();
+                drop(bufw);
+                let t = Instant::now();
+                let bits_raw = buf.len() * 8;
+                // eprintln!("{:?}", &buf[..20]);
+                let bits_huff = write_huffman(&buf, &mut stdout).unwrap();
+                let time_huff = t.elapsed();
+                drop(buf);
+
+                eprintln!(
+                    "frame {frame_number}: {:?}",
+                    time_decode + time_huff + time_encode
+                );
+                eprintln!("\tencode {time_encode:?}");
+                eprintln!("\tdecode {time_decode:?}");
+                eprintln!(
+                    "\thuff {time_huff:?} ({}%)",
+                    ((bits_huff as f32 / bits_raw as f32) * 100.0).round()
+                );
             }
         }
     }
diff --git a/lvc/src/bin/test.rs b/lvc/src/bin/test.rs
new file mode 100644
index 0000000..0f630d3
--- /dev/null
+++ b/lvc/src/bin/test.rs
@@ -0,0 +1,46 @@
+#[cfg(test)]
+mod test {
+
+    use lvc::huff::{read_huffman, write_huffman, BitIO};
+    use std::io::Cursor;
+
+    #[test]
+    fn test_bitio() {
+        let mut buf = Vec::<u8>::new();
+
+        {
+            let mut b = BitIO::new(Cursor::new(&mut buf));
+            b.wbit(true).unwrap();
+            b.wbit(true).unwrap();
+            b.wbit(true).unwrap();
+            b.wbit(true).unwrap();
+            b.wbit(false).unwrap();
+            b.wbit(true).unwrap();
+            b.wbit(true).unwrap();
+            b.wbit(true).unwrap();
+            b.wbit(true).unwrap();
+            b.flush().unwrap();
+        }
+        {
+            let mut b = BitIO::new(Cursor::new(&mut buf));
+            for _ in 0..9 {
+                eprintln!("{:?}", b.rbit().unwrap())
+            }
+        }
+    }
+
+    #[test]
+    fn test_huff() {
+        let a = vec![1; 10000];
+        let mut b = vec![];
+
+        let mut buf = Vec::<u8>::new();
+        write_huffman(&a, &mut Cursor::new(&mut buf)).unwrap();
+        eprintln!("out {buf:x?}");
+        read_huffman(&mut Cursor::new(&mut buf), &mut b).unwrap();
+
+        assert_eq!(a, b)
+    }
+}
+
+fn main() {}
diff --git a/lvc/src/huff.rs b/lvc/src/huff.rs
index 5943207..b21d38d 100644
--- a/lvc/src/huff.rs
+++ b/lvc/src/huff.rs
@@ -1,4 +1,4 @@
-use std::io::Write;
+use std::io::{Read, Result, Write};
 
 #[derive(Debug, Clone)]
 enum HT {
@@ -6,27 +6,66 @@ enum HT {
     Terminal(u8),
 }
 
-fn write_huffman(buf: &[u8], w: &mut impl Write) -> usize {
+pub fn write_huffman(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 = [0u16; 256];
+    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.wbit((k & 1) != 0)?;
+            k >>= 1;
+        }
+    }
+
+    w.flush()?;
+    Ok(w.position)
+}
+
+pub fn read_huffman(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;
+
+    eprintln!("{len:?}");
+    let root = HT::read(&mut r)?;
+    eprintln!("{root:?}");
+    let root = match &root {
+        HT::Branch(a, b) => [a, b],
+        _ => panic!("no!"),
+    };
+
+    let mut cn = root;
+    for _ in 0..len * 8 {
+        let v = r.rbit()?;
+        match cn[v as usize].as_ref() {
+            HT::Branch(a, b) => {
+                cn = [a, b];
+            }
+            HT::Terminal(n) => {
+                buf.push(*n);
+                cn = root;
+            }
         }
     }
 
-    w.position
+    Ok(len)
 }
 
 impl HT {
@@ -40,11 +79,12 @@ impl HT {
         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, HT::Branch(box at, box bt)))
+            parts.push((ap + bp + 1, HT::Branch(box at, box bt)))
         }
         parts[0].1.clone()
     }
-    pub fn create_lut(&self, table: &mut [u16; 256], prefix: u16) {
+    pub fn create_lut(&self, table: &mut [u32; 256], prefix: u32) {
+        assert!(self.depth() < 30, "too deep! doesnt fit {}", self.depth());
         match self {
             HT::Branch(a, b) => {
                 a.create_lut(table, (prefix << 1) | 0);
@@ -56,9 +96,35 @@ impl HT {
             }
         }
     }
+    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()?)),
+        }
+    }
 }
 
-struct BitIO<T> {
+pub struct BitIO<T> {
     inner: T,
     byte: u8,
     position: usize,
@@ -74,12 +140,50 @@ impl<T> BitIO<T> {
 }
 impl<T: Write> BitIO<T> {
     #[inline]
-    pub fn wbit(&mut self, b: bool) {
+    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]).unwrap();
+            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)
     }
 }