1 files changed, 123 insertions, 0 deletions
diff --git a/test2/src/main.rs b/test2/src/main.rs
index 1f10f76..bbbf8d3 100644
--- a/test2/src/main.rs
+++ b/test2/src/main.rs
@@ -4,6 +4,7 @@ pub mod encode;
 use decode::Dec;
 use encode::Enc;
 use framework::bitstream_filter_main;
+use glam::IVec2;
 use std::env::args;
 
 pub const BLOCK_SIZE: i32 = 8;
@@ -15,3 +16,125 @@ fn main() -> Result<(), std::io::Error> {
         _ => panic!("unknown mode"),
     }
 }
+
+#[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 export_block(&self, off: IVec2, buf: &mut Vec<u8>) {
+        // Luma
+        for y in 0..BLOCK_SIZE {
+            let y_off = off.x + (y + off.y) * self.res.x;
+            buf.extend(&self.data[y_off as usize..(y_off + BLOCK_SIZE) as usize]);
+        }
+        // Chroma
+        let uvplane_off = self.res.x * self.res.y;
+        for y in 0..BLOCK_SIZE / 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 + BLOCK_SIZE) as usize]);
+        }
+    }
+    fn import_block(&mut self, off: IVec2, buf: &[u8]) -> usize {
+        let mut p = 0;
+        // Luma
+        for y in 0..BLOCK_SIZE {
+            let y_off = off.x + (y + off.y) * self.res.x;
+            self.data[y_off as usize..(y_off + BLOCK_SIZE) as usize]
+                .copy_from_slice(&buf[p..p + BLOCK_SIZE as usize]);
+            p += BLOCK_SIZE as usize;
+        }
+        // Chroma
+        let uvplane_off = self.res.x * self.res.y;
+        for y in 0..BLOCK_SIZE / 2 {
+            let y_off = uvplane_off + (off.x & !1) + (y + off.y / 2) * self.res.x;
+            self.data[y_off as usize..(y_off + BLOCK_SIZE) as usize]
+                .copy_from_slice(&buf[p..p + BLOCK_SIZE as usize]);
+            p += BLOCK_SIZE as usize;
+        }
+        p
+    }
+    fn import_block_diff(&mut self, off: IVec2, buf: &[u8]) -> usize {
+        // Luma
+        for y in 0..BLOCK_SIZE {
+            let y_off = off.x + (y + off.y) * self.res.x;
+            let i = y * BLOCK_SIZE;
+            for x in 0..BLOCK_SIZE {
+                self.data[(y_off + x) as usize] = self.data[(y_off + x) as usize]
+                    .saturating_add_signed(buf[(i + x) as usize] as i8 - 127);
+            }
+        }
+        // Chroma
+        let uvplane_off = self.res.x * self.res.y;
+        for y in 0..BLOCK_SIZE / 2 {
+            let y_off = uvplane_off + (off.x & !1) + (y + off.y / 2) * self.res.x;
+            let i = BLOCK_SIZE * BLOCK_SIZE + y * BLOCK_SIZE;
+            for x in 0..BLOCK_SIZE {
+                self.data[(y_off + x) as usize] = self.data[(y_off + x) as usize]
+                    .saturating_add_signed(buf[(i + x) as usize] as i8 - 127);
+            }
+        }
+        (BLOCK_SIZE * BLOCK_SIZE + BLOCK_SIZE * BLOCK_SIZE / 4) 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 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;
+
+            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]);
+        }
+    }
+    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 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;
+            }
+        }
+
+        diff
+    }
+}