a

2 files changed, 157 insertions, 3 deletions

diff --git a/mtree-test/src/bin/encode.rs b/mtree-test/src/bin/encode.rs
index f85e4c0..b37fdf6 100644
--- a/mtree-test/src/bin/encode.rs
+++ b/mtree-test/src/bin/encode.rs
@@ -1,13 +1,18 @@
 use framework::{BitstreamFilter, bitstream_filter_main};
 use glam::{I16Vec2, i16vec2};
-use std::io::Result;
+use mtree_test::{AbsRef, Frame};
+use std::{collections::VecDeque, io::Result};
 
 fn main() -> Result<()> {
     bitstream_filter_main::<Enc>()
 }
 
+const BLOCK_SIZE: i16 = 16;
+
 struct Enc {
     res: I16Vec2,
+    last: LastFrames,
+    tree: MTree,
 }
 impl BitstreamFilter for Enc {
     const INPUT_CODEC_ID: &str = "V_UNCOMPRESSED";
@@ -16,9 +21,129 @@ impl BitstreamFilter for Enc {
     fn new(width: u32, height: u32) -> Self {
         Self {
             res: i16vec2(width as i16, height as i16),
+            last: LastFrames {
+                frame_offset: 0,
+                frames: VecDeque::new(),
+            },
+            tree: MTree::Leaf(AbsRef {
+                off: i16vec2(0, 0),
+                frame: 0,
+            }),
+        }
+    }
+    fn process_block(&mut self, frame: Vec<u8>) -> Vec<u8> {
+        let frame = Frame(frame);
+
+        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,
+                    &self.last,
+                    AbsRef {
+                        frame: frame_index,
+                        off: i16vec2(xo, yo),
+                    },
+                )
+            }
         }
+
+        Vec::new()
+    }
+}
+
+struct LastFrames {
+    frame_offset: u64,
+    frames: VecDeque<Frame>,
+}
+
+enum MTree {
+    Branch(Box<[MTree; 2]>),
+    Leaf(AbsRef),
+}
+impl MTree {
+    pub fn depth(&self) -> usize {
+        (match self {
+            MTree::Branch(b) => b[0].depth().max(b[1].depth()),
+            MTree::Leaf(_) => 0,
+        }) + 1
     }
-    fn process_block(&mut self, a: Vec<u8>) -> Vec<u8> {
-        a
+    pub fn center(&self) -> AbsRef {
+        match self {
+            MTree::Branch(b) => b[0].center(),
+            MTree::Leaf(r) => *r,
+        }
+    }
+    pub fn insert(&mut self, res: I16Vec2, last: &LastFrames, r: AbsRef) {
+        match self {
+            MTree::Branch(b) => {
+                if distance_absref(res, last, b[0].center(), r)
+                    < distance_absref(res, last, b[1].center(), r)
+                {
+                    b[0].insert(res, last, r)
+                } else {
+                    b[1].insert(res, last, r)
+                }
+            }
+            MTree::Leaf(l) => *self = MTree::Branch(Box::new([MTree::Leaf(*l), MTree::Leaf(r)])),
+        }
+    }
+    pub fn search(
+        &mut self,
+        res: I16Vec2,
+        last: &LastFrames,
+        frame: &Frame,
+        off: I16Vec2,
+    ) -> AbsRef {
+        match self {
+            MTree::Branch(b) => {
+                let acenter = b[0].center();
+                let bcenter = b[1].center();
+                let adist = distance(
+                    res,
+                    &last.frames[(acenter.frame - last.frame_offset) as usize],
+                    frame,
+                    acenter.off,
+                    off,
+                );
+                let bdist = distance(
+                    res,
+                    &last.frames[(bcenter.frame - last.frame_offset) as usize],
+                    frame,
+                    bcenter.off,
+                    off,
+                );
+                if adist < bdist {
+                    b[0].search(res, last, frame, off)
+                } else {
+                    b[1].search(res, last, frame, off)
+                }
+            }
+            MTree::Leaf(l) => *l,
+        }
+    }
+}
+
+fn distance_absref(res: I16Vec2, last: &LastFrames, a: AbsRef, b: AbsRef) -> u32 {
+    let aframe = &last.frames[(a.frame - last.frame_offset) as usize];
+    let bframe = &last.frames[(b.frame - last.frame_offset) as usize];
+    distance(res, aframe, bframe, a.off, b.off)
+}
+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);
+
+            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
 }
diff --git a/mtree-test/src/lib.rs b/mtree-test/src/lib.rs
index e69de29..1214cbf 100644
--- a/mtree-test/src/lib.rs
+++ b/mtree-test/src/lib.rs
@@ -0,0 +1,29 @@
+use glam::I16Vec2;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct AbsRef {
+    pub off: I16Vec2,
+    pub frame: u64,
+}
+
+pub struct Frame(pub Vec<u8>);
+impl Frame {
+    pub fn index(&self, res: I16Vec2, p: I16Vec2) -> [usize; 3] {
+        let res = res.as_usizevec2();
+        let p = p.as_usizevec2();
+
+        let ystride = res.y;
+        let ysize = res.x * ystride;
+        let uvstride = res.x / 2;
+        let usize = uvstride * (res.y / 2);
+        let puv = p / 2;
+        [
+            p.x + p.y * ystride,
+            ysize + puv.x + puv.y * uvstride,
+            ysize + usize + puv.x + puv.y * uvstride,
+        ]
+    }
+    pub fn get(&self, res: I16Vec2, p: I16Vec2) -> [u8; 3] {
+        self.index(res, p).map(|i| self.0[i])
+    }
+}