fix mtree, switch to nv12 pixfmtHEADmaster

3 files changed, 121 insertions, 139 deletions

diff --git a/mtree-test/src/decode.rs b/mtree-test/src/decode.rs
index 9c29717..97557da 100644
--- a/mtree-test/src/decode.rs
+++ b/mtree-test/src/decode.rs
@@ -1,10 +1,10 @@
 use crate::{BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
 use framework::BitstreamFilter;
-use glam::{I16Vec2, i16vec2};
+use glam::{IVec2, ivec2};
 use std::collections::VecDeque;
 
 pub struct Dec {
-    res: I16Vec2,
+    res: IVec2,
     last: LastFrames,
 }
 impl BitstreamFilter for Dec {
@@ -13,7 +13,7 @@ impl BitstreamFilter for Dec {
 
     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(),
@@ -26,12 +26,11 @@ impl BitstreamFilter for Dec {
 
         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 brect = boff..boff + I16Vec2::splat(BLOCK_SIZE);
+                let boff = ivec2(bx * BLOCK_SIZE, by * BLOCK_SIZE);
                 let kind = a[0];
                 a = &a[1..];
                 if kind == 0 {
-                    let size = frame.import_rect(self.res, brect, &a);
+                    let size = frame.import_rect(boff, IVec2::splat(BLOCK_SIZE), &a);
                     a = &a[size..];
                 } else {
                     let rframeindex =
@@ -44,12 +43,11 @@ impl BitstreamFilter for Dec {
 
                     let rframe = &self.last.frames[(rframeindex - self.last.frame_offset) as usize];
                     frame_to_frame_rect_copy(
-                        self.res,
-                        &mut frame,
                         rframe,
-                        I16Vec2::splat(BLOCK_SIZE),
+                        &mut frame,
+                        IVec2::splat(BLOCK_SIZE),
+                        ivec2(offx as i32, offy as i32),
                         boff,
-                        i16vec2(offx, offy),
                     );
                 }
             }
@@ -57,6 +55,6 @@ impl BitstreamFilter for Dec {
 
         self.last.frames.push_back(frame.clone());
 
-        frame.0
+        frame.data
     }
 }
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
 }
diff --git a/mtree-test/src/main.rs b/mtree-test/src/main.rs
index 3952722..927ecac 100644
--- a/mtree-test/src/main.rs
+++ b/mtree-test/src/main.rs
@@ -4,8 +4,8 @@ pub mod encode;
 use decode::Dec;
 use encode::Enc;
 use framework::bitstream_filter_main;
-use glam::{I16Vec2, i16vec2};
-use std::{collections::VecDeque, env::args, ops::Range};
+use glam::IVec2;
+use std::{collections::VecDeque, env::args};
 
 fn main() -> Result<(), std::io::Error> {
     match args().nth(1).unwrap().as_str() {
@@ -15,96 +15,88 @@ fn main() -> Result<(), std::io::Error> {
     }
 }
 
-pub const BLOCK_SIZE: i16 = 8;
+pub const BLOCK_SIZE: i32 = 8;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct AbsRef {
-    pub off: I16Vec2,
+    pub off: IVec2,
     pub frame: u64,
 }
 
 #[derive(Clone)]
-pub struct Frame(pub Vec<u8>);
+pub struct Frame {
+    pub res: IVec2,
+    pub data: Vec<u8>,
+}
 
 impl Frame {
-    pub fn new(res: I16Vec2) -> Self {
-        let res = res.as_usizevec2();
-        Self(vec![0; res.x * res.y + 2 * (res.x / 2 * res.y / 2)])
-    }
-
-    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 slice(&self, res: I16Vec2, y: i16, x: Range<i16>) -> [&[u8]; 3] {
-        let start = self.index(res, i16vec2(x.start, y));
-        let end = self.index(res, i16vec2(x.end, y));
-        [
-            &self.0[start[0]..end[0]],
-            &self.0[start[1]..end[1]],
-            &self.0[start[2]..end[2]],
-        ]
-    }
-    pub fn get(&self, res: I16Vec2, p: I16Vec2) -> [u8; 3] {
-        self.index(res, p).map(|i| self.0[i])
+    pub fn new(res: IVec2) -> Self {
+        let ures = res.as_usizevec2();
+        Self {
+            res,
+            data: vec![0; ures.x * ures.y + 2 * (ures.x / 2 * ures.y / 2)],
+        }
     }
-
-    pub fn export_rect(&self, res: I16Vec2, r: Range<I16Vec2>, out: &mut Vec<u8>) {
-        for y in r.start.y..r.end.y {
-            let slices = self.slice(res, y, r.start.x..r.end.x);
-            out.extend(slices[0]);
-            out.extend(slices[1]);
-            out.extend(slices[2]);
+    pub fn export_rect(&self, off: IVec2, size: IVec2, buf: &mut Vec<u8>) {
+        // Luma
+        for y in 0..size.y {
+            let y_off = off.x + (y + off.y) * self.res.x;
+            buf.extend(&self.data[y_off as usize..(y_off + size.x) as usize]);
+        }
+        // Chroma
+        let uvplane_off = self.res.x * self.res.y;
+        for y in 0..size.y / 2 {
+            let y_off = uvplane_off + (off.x & !1) + (y + off.y / 2) * self.res.x;
+            buf.extend(&self.data[y_off as usize..(y_off + size.x) as usize]);
         }
     }
-    pub fn import_rect(&mut self, res: I16Vec2, r: Range<I16Vec2>, mut data: &[u8]) -> usize {
-        let mut off = 0;
-        for y in r.start.y..r.end.y {
-            let start = self.index(res, i16vec2(r.start.x, y));
-            let end = self.index(res, i16vec2(r.end.x, y));
-
-            let mut do_slice = |s: &mut [u8]| {
-                s.copy_from_slice(&data[..s.len()]);
-                data = &data[s.len()..];
-                off += s.len();
-            };
-
-            do_slice(&mut self.0[start[0]..end[0]]);
-            do_slice(&mut self.0[start[1]..end[1]]);
-            do_slice(&mut self.0[start[2]..end[2]]);
+    pub fn import_rect(&mut self, off: IVec2, size: IVec2, buf: &[u8]) -> usize {
+        let mut p = 0;
+        // Luma
+        for y in 0..size.y {
+            let y_off = off.x + (y + off.y) * self.res.x;
+            self.data[y_off as usize..(y_off + size.x) as usize]
+                .copy_from_slice(&buf[p..p + size.x as usize]);
+            p += size.x as usize;
+        }
+        // Chroma
+        let uvplane_off = self.res.x * self.res.y;
+        for y in 0..size.y / 2 {
+            let y_off = uvplane_off + (off.x & !1) + (y + off.y / 2) * self.res.x;
+            self.data[y_off as usize..(y_off + size.x) as usize]
+                .copy_from_slice(&buf[p..p + size.x as usize]);
+            p += size.x as usize;
         }
-        off
+        p
     }
 }
 
 pub fn frame_to_frame_rect_copy(
-    res: I16Vec2,
-    aframe: &mut Frame,
-    bframe: &Frame,
-    size: I16Vec2,
-    aoff: I16Vec2,
-    boff: I16Vec2,
+    aframe: &Frame,
+    bframe: &mut Frame,
+    size: IVec2,
+    aoff: IVec2,
+    boff: IVec2,
 ) {
+    assert_eq!(aframe.res, bframe.res);
+    let res = aframe.res;
+    // Luma
     for y in 0..size.y {
-        let astart = aframe.index(res, i16vec2(aoff.x, aoff.y + y));
-        let aend = aframe.index(res, i16vec2(aoff.x + size.x, aoff.y + y));
-        let bstart = bframe.index(res, i16vec2(boff.x, boff.y + y));
-        let bend = bframe.index(res, i16vec2(boff.x + size.x, boff.y + y));
+        let ay_off = aoff.x + (y + aoff.y) * res.x;
+        let by_off = boff.x + (y + boff.y) * res.x;
+
+        bframe.data[by_off as usize..(by_off + size.x) as usize]
+            .copy_from_slice(&aframe.data[ay_off as usize..(ay_off + size.x) as usize]);
+    }
+
+    // Chroma
+    let uvplane_off = res.x * res.y;
+    for y in 0..size.y / 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;
 
-        aframe.0[astart[0]..aend[0]].copy_from_slice(&bframe.0[bstart[0]..bend[0]]);
-        aframe.0[astart[1]..aend[1]].copy_from_slice(&bframe.0[bstart[1]..bend[1]]);
-        aframe.0[astart[2]..aend[2]].copy_from_slice(&bframe.0[bstart[2]..bend[2]]);
+        bframe.data[by_off as usize..(by_off + size.x) as usize]
+            .copy_from_slice(&aframe.data[ay_off as usize..(ay_off + size.x) as usize]);
     }
 }