1 files changed, 75 insertions, 20 deletions

diff --git a/mtree-test/src/bin/encode.rs b/mtree-test/src/bin/encode.rs
index b37fdf6..3cdb1c8 100644
--- a/mtree-test/src/bin/encode.rs
+++ b/mtree-test/src/bin/encode.rs
@@ -1,14 +1,12 @@
 use framework::{BitstreamFilter, bitstream_filter_main};
 use glam::{I16Vec2, i16vec2};
-use mtree_test::{AbsRef, Frame};
+use mtree_test::{AbsRef, BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
 use std::{collections::VecDeque, io::Result};
 
 fn main() -> Result<()> {
     bitstream_filter_main::<Enc>()
 }
 
-const BLOCK_SIZE: i16 = 16;
-
 struct Enc {
     res: I16Vec2,
     last: LastFrames,
@@ -27,18 +25,57 @@ impl BitstreamFilter for Enc {
             },
             tree: MTree::Leaf(AbsRef {
                 off: i16vec2(0, 0),
-                frame: 0,
+                frame: u64::MAX,
             }),
         }
     }
     fn process_block(&mut self, frame: Vec<u8>) -> Vec<u8> {
-        let frame = Frame(frame);
+        let mut frame = Frame(frame);
+
+        let mut out = Vec::new();
+        let mut num_refs = 0;
+
+        for bx in 0..self.res.x / BLOCK_SIZE {
+            for by in 0..self.res.y / BLOCK_SIZE {
+                let boff = i16vec2(bx * BLOCK_SIZE, by * BLOCK_SIZE);
+                let r = self.tree.search(self.res, &self.last, &frame, boff);
+
+                let rdist = if r.frame == u64::MAX {
+                    u32::MAX
+                } else {
+                    distance(
+                        self.res,
+                        &self.last.frames[(r.frame - self.last.frame_offset) as usize],
+                        &frame,
+                        r.off,
+                        boff,
+                    )
+                };
+                if rdist < 1000 {
+                    out.push(1);
+                    out.extend(r.frame.to_le_bytes());
+                    out.extend(r.off.x.to_le_bytes());
+                    out.extend(r.off.y.to_le_bytes());
+                    frame_to_frame_rect_copy(
+                        self.res,
+                        &mut frame,
+                        &self.last.frames[(r.frame - self.last.frame_offset) as usize],
+                        I16Vec2::splat(BLOCK_SIZE),
+                        boff,
+                        r.off,
+                    );
+                    num_refs += 1;
+                } else {
+                    out.push(0);
+                    frame.export_rect(self.res, boff..boff + I16Vec2::splat(BLOCK_SIZE), &mut out);
+                }
+            }
+        }
 
         self.last.frames.push_back(frame);
         let frame_index = self.last.frame_offset + self.last.frames.len() as u64 - 1;
 
         for xo in 0..self.res.x - BLOCK_SIZE {
-            eprintln!("aaaaaaa {}  aaaa", self.tree.depth());
             for yo in 0..self.res.y - BLOCK_SIZE {
                 self.tree.insert(
                     self.res,
@@ -50,16 +87,12 @@ impl BitstreamFilter for Enc {
                 )
             }
         }
+        eprintln!("aaaaaaa {} {num_refs} aaaa", self.tree.depth());
 
-        Vec::new()
+        out
     }
 }
 
-struct LastFrames {
-    frame_offset: u64,
-    frames: VecDeque<Frame>,
-}
-
 enum MTree {
     Branch(Box<[MTree; 2]>),
     Leaf(AbsRef),
@@ -88,7 +121,13 @@ impl MTree {
                     b[1].insert(res, last, r)
                 }
             }
-            MTree::Leaf(l) => *self = MTree::Branch(Box::new([MTree::Leaf(*l), MTree::Leaf(r)])),
+            MTree::Leaf(l) => {
+                if l.frame == u64::MAX {
+                    *l = r
+                } else {
+                    *self = MTree::Branch(Box::new([MTree::Leaf(*l), MTree::Leaf(r)]))
+                }
+            }
         }
     }
     pub fn search(
@@ -135,14 +174,30 @@ fn distance_absref(res: I16Vec2, last: &LastFrames, a: AbsRef, b: AbsRef) -> u32
 fn distance(res: I16Vec2, aframe: &Frame, bframe: &Frame, aoff: I16Vec2, boff: I16Vec2) -> u32 {
     let mut diff = 0;
     for yo in 0..BLOCK_SIZE {
-        for xo in 0..BLOCK_SIZE {
-            let off = i16vec2(xo, yo);
-            let asample = aframe.get(res, aoff + off);
-            let bsample = bframe.get(res, boff + off);
+        // for xo in 0..BLOCK_SIZE {
+        //     let off = i16vec2(xo, yo);
+        //     let asample = aframe.get(res, aoff + off);
+        //     let bsample = bframe.get(res, boff + off);
 
-            diff += asample[0].abs_diff(bsample[0]) as u32;
-            diff += asample[1].abs_diff(bsample[1]) as u32;
-            diff += asample[2].abs_diff(bsample[2]) as u32;
+        //     diff += asample[0].abs_diff(bsample[0]) as u32;
+        //     diff += asample[1].abs_diff(bsample[1]) as u32;
+        //     diff += asample[2].abs_diff(bsample[2]) as u32;
+        // }
+
+        let aslice = aframe.slice(res, yo + aoff.y, aoff.x..aoff.x + BLOCK_SIZE);
+        let bslice = bframe.slice(res, yo + boff.y, boff.x..boff.x + BLOCK_SIZE);
+
+        for i in 0..BLOCK_SIZE as usize {
+            diff += aslice[0][i].abs_diff(bslice[0][i]) as u32;
+            diff += aslice[0][i].abs_diff(bslice[0][i]) as u32;
+        }
+        for i in 0..BLOCK_SIZE as usize / 2 {
+            diff += aslice[1][i].abs_diff(bslice[1][i]) as u32;
+            diff += aslice[1][i].abs_diff(bslice[1][i]) as u32;
+        }
+        for i in 0..BLOCK_SIZE as usize / 2 {
+            diff += aslice[2][i].abs_diff(bslice[2][i]) as u32;
+            diff += aslice[2][i].abs_diff(bslice[2][i]) as u32;
         }
     }
     diff