fix mtree, switch to nv12 pixfmtHEADmaster

1 files changed, 48 insertions, 56 deletions

diff --git a/mtree-test/src/encode.rs b/mtree-test/src/encode.rs
index d746e7f..0aae8ee 100644
--- a/mtree-test/src/encode.rs
+++ b/mtree-test/src/encode.rs
@@ -1,6 +1,6 @@
 use crate::{AbsRef, BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
 use framework::BitstreamFilter;
-use glam::{I16Vec2, i16vec2};
+use glam::{IVec2, ivec2};
 use rayon::{
     iter::{IntoParallelRefIterator, ParallelIterator},
     join,
@@ -8,7 +8,7 @@ use rayon::{
 use std::{collections::VecDeque, time::Instant};
 
 pub struct Enc {
-    res: I16Vec2,
+    res: IVec2,
     last: LastFrames,
     tree: MTree,
     frame_num: u64,
@@ -21,13 +21,13 @@ impl BitstreamFilter for Enc {
 
     fn new(width: u32, height: u32) -> Self {
         Self {
-            res: i16vec2(width as i16, height as i16),
+            res: ivec2(width as i32, height as i32),
             last: LastFrames {
                 frame_offset: 0,
                 frames: VecDeque::new(),
             },
             tree: MTree::Leaf(AbsRef {
-                off: i16vec2(0, 0),
+                off: ivec2(0, 0),
                 frame: u64::MAX,
             }),
             frame_num: 0,
@@ -35,21 +35,23 @@ impl BitstreamFilter for Enc {
         }
     }
     fn process_block(&mut self, frame: Vec<u8>) -> Vec<u8> {
-        let mut frame = Frame(frame);
+        let mut frame = Frame {
+            data: frame,
+            res: self.res,
+        };
         let mut out = Vec::new();
         let mut num_refs = 0;
 
         let t = Instant::now();
         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 boff = ivec2(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,
@@ -59,20 +61,19 @@ impl BitstreamFilter for Enc {
                 if rdist < BLOCK_SIZE as u32 * BLOCK_SIZE as u32 * 50 {
                     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());
+                    out.extend((r.off.x as i16).to_le_bytes());
+                    out.extend((r.off.y as i16).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,
+                        &mut frame,
+                        IVec2::splat(BLOCK_SIZE),
                         r.off,
+                        boff,
                     );
                     num_refs += 1;
                 } else {
                     out.push(0);
-                    frame.export_rect(self.res, boff..boff + I16Vec2::splat(BLOCK_SIZE), &mut out);
+                    frame.export_rect(boff, IVec2::splat(BLOCK_SIZE), &mut out);
                 }
             }
         }
@@ -89,7 +90,7 @@ impl BitstreamFilter for Enc {
                     for yo in 0..self.res.y - BLOCK_SIZE {
                         rs.push(AbsRef {
                             frame: frame_index,
-                            off: i16vec2(xo, yo),
+                            off: ivec2(xo, yo),
                         });
                     }
                 }
@@ -102,7 +103,7 @@ impl BitstreamFilter for Enc {
                             &self.last,
                             AbsRef {
                                 frame: frame_index,
-                                off: i16vec2(xo, yo),
+                                off: ivec2(xo, yo),
                             },
                         )
                     }
@@ -135,14 +136,14 @@ impl MTree {
             MTree::Leaf(r) => *r,
         }
     }
-    pub fn bulk_insert(&mut self, res: I16Vec2, last: &LastFrames, mut rs: Vec<AbsRef>) {
+    pub fn bulk_insert(&mut self, res: IVec2, last: &LastFrames, mut rs: Vec<AbsRef>) {
         match self {
             MTree::Branch(b) => {
                 let dirs = rs
                     .par_iter()
                     .map(|r| {
-                        distance_absref(res, last, b[0].center(), *r)
-                            < distance_absref(res, last, b[1].center(), *r)
+                        distance_absref(last, b[0].center(), *r)
+                            < distance_absref(last, b[1].center(), *r)
                     })
                     .collect::<Vec<bool>>();
 
@@ -178,11 +179,10 @@ impl MTree {
             }
         }
     }
-    pub fn insert(&mut self, res: I16Vec2, last: &LastFrames, r: AbsRef) {
+    pub fn insert(&mut self, res: IVec2, 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)
+                if distance_absref(last, b[0].center(), r) < distance_absref(last, b[1].center(), r)
                 {
                     b[0].insert(res, last, r)
                 } else {
@@ -198,26 +198,18 @@ impl MTree {
             }
         }
     }
-    pub fn search(
-        &mut self,
-        res: I16Vec2,
-        last: &LastFrames,
-        frame: &Frame,
-        off: I16Vec2,
-    ) -> AbsRef {
+    pub fn search(&mut self, res: IVec2, last: &LastFrames, frame: &Frame, off: IVec2) -> 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,
@@ -234,39 +226,39 @@ impl MTree {
     }
 }
 
-fn distance_absref(res: I16Vec2, last: &LastFrames, a: AbsRef, b: AbsRef) -> u32 {
+fn distance_absref(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)
+    distance(aframe, bframe, a.off, b.off)
 }
-fn distance(res: I16Vec2, aframe: &Frame, bframe: &Frame, aoff: I16Vec2, boff: I16Vec2) -> u32 {
+fn distance(aframe: &Frame, bframe: &Frame, aoff: IVec2, boff: IVec2) -> u32 {
+    assert_eq!(aframe.res, bframe.res);
+    let res = aframe.res;
     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;
-        // }
 
-        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);
+    // Luma
+    for y in 0..BLOCK_SIZE {
+        let ay_off = aoff.x + (y + aoff.y) * res.x;
+        let by_off = boff.x + (y + boff.y) * res.x;
 
-        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 x in 0..BLOCK_SIZE {
+            diff += aframe.data[(ay_off + x) as usize].abs_diff(bframe.data[(by_off + x) as usize])
+                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;
+    }
+
+    // Chroma
+    let uvplane_off = res.x * res.y;
+    for y in 0..BLOCK_SIZE / 2 {
+        let ay_off = uvplane_off + (aoff.x & !1) + (y + aoff.y / 2) * res.x;
+        let by_off = uvplane_off + (boff.x & !1) + (y + boff.y / 2) * res.x;
+
+        for x in 0..BLOCK_SIZE {
+            diff += aframe.data[(ay_off + x) as usize].abs_diff(bframe.data[(by_off + x) as usize])
+                as u32
+                * 2;
         }
     }
+
     diff
 }