fix mtree, switch to nv12 pixfmtHEADmaster

1 files changed, 64 insertions, 72 deletions

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]);
     }
 }