test2

9 files changed, 235 insertions, 19 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 44cff8c..2152237 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -79,3 +79,12 @@ dependencies = [
  "crossbeam-deque",
  "crossbeam-utils",
 ]
+
+[[package]]
+name = "test2"
+version = "0.1.0"
+dependencies = [
+ "framework",
+ "glam",
+ "rayon",
+]
diff --git a/Cargo.toml b/Cargo.toml
index eccbbe0..fc8ad92 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,3 +1,3 @@
 [workspace]
-members = ["framework", "mtree-test"]
+members = ["framework", "mtree-test", "test2"]
 resolver = "3"
diff --git a/mtree-test/src/bin/decode.rs b/mtree-test/src/decode.rs
index 992bebc..9c29717 100644
--- a/mtree-test/src/bin/decode.rs
+++ b/mtree-test/src/decode.rs
@@ -1,13 +1,9 @@
-use framework::{BitstreamFilter, bitstream_filter_main};
+use crate::{BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
+use framework::BitstreamFilter;
 use glam::{I16Vec2, i16vec2};
-use mtree_test::{BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
-use std::{collections::VecDeque, io::Result};
+use std::collections::VecDeque;
 
-fn main() -> Result<()> {
-    bitstream_filter_main::<Dec>()
-}
-
-struct Dec {
+pub struct Dec {
     res: I16Vec2,
     last: LastFrames,
 }
diff --git a/mtree-test/src/bin/encode.rs b/mtree-test/src/encode.rs
index 80f3499..d746e7f 100644
--- a/mtree-test/src/bin/encode.rs
+++ b/mtree-test/src/encode.rs
@@ -1,17 +1,13 @@
-use framework::{BitstreamFilter, bitstream_filter_main};
+use crate::{AbsRef, BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
+use framework::BitstreamFilter;
 use glam::{I16Vec2, i16vec2};
-use mtree_test::{AbsRef, BLOCK_SIZE, Frame, LastFrames, frame_to_frame_rect_copy};
 use rayon::{
-    iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator},
+    iter::{IntoParallelRefIterator, ParallelIterator},
     join,
 };
-use std::{collections::VecDeque, io::Result, time::Instant};
+use std::{collections::VecDeque, time::Instant};
 
-fn main() -> Result<()> {
-    bitstream_filter_main::<Enc>()
-}
-
-struct Enc {
+pub struct Enc {
     res: I16Vec2,
     last: LastFrames,
     tree: MTree,
diff --git a/mtree-test/src/lib.rs b/mtree-test/src/main.rs
index fd282db..3952722 100644
--- a/mtree-test/src/lib.rs
+++ b/mtree-test/src/main.rs
@@ -1,5 +1,19 @@
+pub mod decode;
+pub mod encode;
+
+use decode::Dec;
+use encode::Enc;
+use framework::bitstream_filter_main;
 use glam::{I16Vec2, i16vec2};
-use std::{collections::VecDeque, ops::Range};
+use std::{collections::VecDeque, env::args, ops::Range};
+
+fn main() -> Result<(), std::io::Error> {
+    match args().nth(1).unwrap().as_str() {
+        "enc" => bitstream_filter_main::<Enc>(),
+        "dec" => bitstream_filter_main::<Dec>(),
+        _ => panic!("unknown mode"),
+    }
+}
 
 pub const BLOCK_SIZE: i16 = 8;
 
diff --git a/test2/Cargo.toml b/test2/Cargo.toml
new file mode 100644
index 0000000..4663775
--- /dev/null
+++ b/test2/Cargo.toml
@@ -0,0 +1,9 @@
+[package]
+name = "test2"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+framework = { path = "../framework" }
+glam = "0.30.3"
+rayon = "1.10.0"
diff --git a/test2/src/decode.rs b/test2/src/decode.rs
new file mode 100644
index 0000000..974c6d9
--- /dev/null
+++ b/test2/src/decode.rs
@@ -0,0 +1,20 @@
+use framework::BitstreamFilter;
+use glam::{I16Vec2, i16vec2};
+
+pub struct Dec {
+    res: I16Vec2,
+}
+impl BitstreamFilter for Dec {
+    const INPUT_CODEC_ID: &str = "V_VCETEST2";
+    const OUTPUT_CODEC_ID: &str = "V_UNCOMPRESSED";
+
+    fn new(width: u32, height: u32) -> Self {
+        Self {
+            res: i16vec2(width as i16, height as i16),
+        }
+    }
+
+    fn process_block(&mut self, a: Vec<u8>) -> Vec<u8> {
+        a
+    }
+}
diff --git a/test2/src/encode.rs b/test2/src/encode.rs
new file mode 100644
index 0000000..799d8d4
--- /dev/null
+++ b/test2/src/encode.rs
@@ -0,0 +1,155 @@
+use crate::BLOCK_SIZE;
+use framework::BitstreamFilter;
+use glam::{IVec2, ivec2};
+
+pub struct Enc {
+    res: IVec2,
+    last: Frame,
+}
+impl BitstreamFilter for Enc {
+    const INPUT_CODEC_ID: &str = "V_UNCOMPRESSED";
+    const OUTPUT_CODEC_ID: &str = "V_VCETEST2";
+
+    fn new(width: u32, height: u32) -> Self {
+        let res = ivec2(width as i32, height as i32);
+        Self {
+            res,
+            last: Frame::new(res),
+        }
+    }
+
+    fn process_block(&mut self, a: Vec<u8>) -> Vec<u8> {
+        let mut frame = Frame {
+            data: a,
+            res: self.res,
+        };
+
+        for by in 0..frame.res.y / BLOCK_SIZE {
+            for bx in 0..frame.res.x / BLOCK_SIZE {
+                let boff = ivec2(bx * BLOCK_SIZE, by * BLOCK_SIZE);
+
+                let mut best_d = Frame::compare_block(&frame, &self.last, boff, boff);
+                let mut best_off = ivec2(0, 0);
+
+                for granularity in [8, 4, 2, 1] {
+                    for dir in [
+                        ivec2(1, 0),
+                        ivec2(1, 1),
+                        ivec2(0, 1),
+                        ivec2(-1, 1),
+                        ivec2(-1, 0),
+                        ivec2(-1, -1),
+                        ivec2(0, -1),
+                        ivec2(1, -1),
+                    ] {
+                        let roff = boff + dir * granularity;
+                        if roff.x >= 0
+                            && roff.y >= 0
+                            && roff.x + BLOCK_SIZE < frame.res.x
+                            && roff.y + BLOCK_SIZE < frame.res.y
+                        {
+                            let d = Frame::compare_block(&frame, &self.last, boff, roff);
+                            if d < best_d {
+                                best_d = d;
+                                best_off = roff;
+                            }
+                        }
+                    }
+                }
+
+                if best_d < 8 * 8 * 10 {
+                    Frame::copy_block(&self.last, &mut frame, best_off, boff);
+                }
+            }
+        }
+
+        self.last = frame.clone();
+        frame.data
+    }
+}
+
+#[derive(Clone)]
+struct Frame {
+    res: IVec2,
+    data: Vec<u8>,
+}
+
+impl Frame {
+    fn new(res: IVec2) -> Self {
+        Frame {
+            res,
+            data: vec![127; (res.x * res.y + res.x * res.y / 2) as usize],
+        }
+    }
+    fn copy_block(aframe: &Frame, bframe: &mut Frame, aoff: IVec2, boff: IVec2) {
+        assert_eq!(aframe.res, bframe.res);
+        let res = aframe.res;
+
+        // 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;
+
+            bframe.data[by_off as usize..(by_off + BLOCK_SIZE) as usize]
+                .copy_from_slice(&aframe.data[ay_off as usize..(ay_off + BLOCK_SIZE) as usize]);
+        }
+
+        // Chroma
+        let uplane_off = res.x * res.y;
+        let vplane_off = uplane_off + (res.x * res.y) / 4;
+
+        for y in 0..BLOCK_SIZE / 2 {
+            let ay_off = aoff.x / 2 + (y + aoff.y / 2) * res.x / 2;
+            let by_off = boff.x / 2 + (y + boff.y / 2) * res.x / 2;
+
+            for x in 0..BLOCK_SIZE / 2 {
+                let ay_index = ay_off + x;
+                let by_index = by_off + x;
+
+                bframe.data[(by_index + uplane_off) as usize] =
+                    aframe.data[(ay_index + uplane_off) as usize];
+                bframe.data[(by_index + vplane_off) as usize] =
+                    aframe.data[(ay_index + vplane_off) as usize];
+            }
+        }
+    }
+    fn compare_block(aframe: &Frame, bframe: &Frame, aoff: IVec2, boff: IVec2) -> u32 {
+        assert_eq!(aframe.res, bframe.res);
+        let res = aframe.res;
+        let mut diff = 0;
+
+        // 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 x in 0..BLOCK_SIZE {
+                diff += aframe.data[(ay_off + x) as usize]
+                    .abs_diff(bframe.data[(by_off + x) as usize]) as u32
+            }
+        }
+
+        // Chroma
+        let uplane_off = res.x * res.y;
+        let vplane_off = uplane_off + (res.x * res.y) / 4;
+
+        for y in 0..BLOCK_SIZE / 2 {
+            let ay_off = aoff.x / 2 + (y + aoff.y / 2) * res.x / 2;
+            let by_off = boff.x / 2 + (y + boff.y / 2) * res.x / 2;
+
+            for x in 0..BLOCK_SIZE / 2 {
+                let ay_index = ay_off + x;
+                let by_index = by_off + x;
+
+                diff += aframe.data[(ay_index + uplane_off) as usize]
+                    .abs_diff(bframe.data[(by_index + uplane_off) as usize])
+                    as u32;
+                diff += aframe.data[(ay_index + vplane_off) as usize]
+                    .abs_diff(bframe.data[(by_index + vplane_off) as usize])
+                    as u32;
+            }
+        }
+
+        diff
+    }
+}
diff --git a/test2/src/main.rs b/test2/src/main.rs
new file mode 100644
index 0000000..1f10f76
--- /dev/null
+++ b/test2/src/main.rs
@@ -0,0 +1,17 @@
+pub mod decode;
+pub mod encode;
+
+use decode::Dec;
+use encode::Enc;
+use framework::bitstream_filter_main;
+use std::env::args;
+
+pub const BLOCK_SIZE: i32 = 8;
+
+fn main() -> Result<(), std::io::Error> {
+    match args().nth(1).unwrap().as_str() {
+        "enc" => bitstream_filter_main::<Enc>(),
+        "dec" => bitstream_filter_main::<Dec>(),
+        _ => panic!("unknown mode"),
+    }
+}