infra for compressed literal

9 files changed, 283 insertions, 21 deletions

diff --git a/evc/Cargo.lock b/evc/Cargo.lock
index 22ef8be..9e3ca03 100644
--- a/evc/Cargo.lock
+++ b/evc/Cargo.lock
@@ -18,6 +18,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "216261ddc8289130e551ddcd5ce8a064710c0d064a4d2895c67151c92b5443f6"
 
 [[package]]
+name = "autocfg"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
+
+[[package]]
 name = "bitflags"
 version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -114,6 +120,7 @@ dependencies = [
  "clap",
  "env_logger",
  "log",
+ "rustdct",
 ]
 
 [[package]]
@@ -187,6 +194,34 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d"
 
 [[package]]
+name = "num-complex"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ae39348c8bc5fbd7f40c727a9925f03517afd2ab27d46702108b6a7e5414c19"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-integer"
+version = "0.1.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9"
+dependencies = [
+ "autocfg",
+ "num-traits",
+]
+
+[[package]]
+name = "num-traits"
+version = "0.2.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "578ede34cf02f8924ab9447f50c28075b4d3e5b269972345e7e0372b38c6cdcd"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
 name = "once_cell"
 version = "1.16.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -199,6 +234,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9b7820b9daea5457c9f21c69448905d723fbd21136ccf521748f23fd49e723ee"
 
 [[package]]
+name = "primal-check"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9df7f93fd637f083201473dab4fee2db4c429d32e55e3299980ab3957ab916a0"
+dependencies = [
+ "num-integer",
+]
+
+[[package]]
 name = "proc-macro-error"
 version = "1.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -258,6 +302,30 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "456c603be3e8d448b072f410900c09faf164fbce2d480456f50eea6e25f9c848"
 
 [[package]]
+name = "rustdct"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8b61555105d6a9bf98797c063c362a1d24ed8ab0431655e38f1cf51e52089551"
+dependencies = [
+ "rustfft",
+]
+
+[[package]]
+name = "rustfft"
+version = "6.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e17d4f6cbdb180c9f4b2a26bbf01c4e647f1e1dea22fe8eb9db54198b32f9434"
+dependencies = [
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "primal-check",
+ "strength_reduce",
+ "transpose",
+ "version_check",
+]
+
+[[package]]
 name = "rustix"
 version = "0.36.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -272,6 +340,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "strength_reduce"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fe895eb47f22e2ddd4dabc02bce419d2e643c8e3b585c78158b349195bc24d82"
+
+[[package]]
 name = "strsim"
 version = "0.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -298,6 +372,16 @@ dependencies = [
 ]
 
 [[package]]
+name = "transpose"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6522d49d03727ffb138ae4cbc1283d3774f0d10aa7f9bf52e6784c45daf9b23"
+dependencies = [
+ "num-integer",
+ "strength_reduce",
+]
+
+[[package]]
 name = "unicode-ident"
 version = "1.0.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/evc/Cargo.toml b/evc/Cargo.toml
index f88e947..7a1178b 100644
--- a/evc/Cargo.toml
+++ b/evc/Cargo.toml
@@ -8,3 +8,4 @@ clap = { version = "*", features = ["derive"] }
 anyhow = "1.0.66"
 log = "0.4.17"
 env_logger = "0.10.0"
+rustdct = "0.7.1"
diff --git a/evc/src/bin/encode.rs b/evc/src/bin/encode.rs
index 32af242..4b7e26c 100644
--- a/evc/src/bin/encode.rs
+++ b/evc/src/bin/encode.rs
@@ -2,6 +2,7 @@ use anyhow::Context;
 use clap::Parser;
 use evc::{
     codec::{
+        compress::compress_block,
         decode::decode_block,
         encode::{encode_block, EncodeConfig, EncodeMode},
     },
@@ -26,7 +27,7 @@ pub struct EncodeArgs {
     #[arg(short, long, default_value = "8")]
     jobs: usize,
 
-    #[arg(short, long, default_value = "4")]
+    #[arg(short, long, default_value = "8")]
     min_block_size: isize,
 
     #[arg(short = 't', long, default_value = "200")]
@@ -70,10 +71,24 @@ fn main() -> anyhow::Result<()> {
         let v1 = frame.view();
         let v2 = prev_frame.view();
 
-        let root = encode_block(v1, v2, &config);
+        let (error, mut root) = encode_block(v1, v2, &config);
+
+        compress_block(
+            &mut root,
+            Vec2 {
+                x: size.x as usize,
+                y: size.y as usize,
+            },
+        );
+
         root.write(&mut output, size)
             .context("writing encoded frame")?;
 
+        info!(
+            "cumulative error: {error} ({} per pixel)",
+            error / frame.view().area() as f64
+        );
+
         decode_block(
             &root,
             frame.view_mut(),
diff --git a/evc/src/block.rs b/evc/src/block.rs
index cbf69bf..69cc460 100644
--- a/evc/src/block.rs
+++ b/evc/src/block.rs
@@ -1,5 +1,5 @@
 use crate::{
-    format::ser::{ConstSizeSerExt, Sink, Small, Source},
+    format::ser::{ConstSizeSerExt, Ser, Sink, Small, Source},
     helpers::vector::Vec2,
     helpers::{matrix::Mat2, pixel::Pixel},
 };
@@ -8,7 +8,7 @@ use anyhow::bail;
 #[derive(Clone, Debug, PartialEq)]
 pub enum Block {
     Literal(Vec<Pixel>),
-    CompressedLiteral(Vec<Pixel>),
+    CompressedLiteral(Vec<u8>),
     Split(Box<[Block; 2]>),
     Reference { translation: Vec2<isize> },
     AdvancedReference(AdvancedReference),
@@ -32,7 +32,10 @@ impl Block {
                 sink.put(0u8)?;
                 pixels.write_const_size(sink, size.area() as usize)?;
             }
-            Block::CompressedLiteral(_) => bail!("compressed literal is not supported"),
+            Block::CompressedLiteral(data) => {
+                sink.put(1u8)?;
+                data.write(sink)?;
+            }
             Block::Split(box [a, b]) => {
                 sink.put(2u8)?;
                 let [asize, bsize] = split_size(size);
@@ -58,7 +61,7 @@ impl Block {
     pub fn read(source: &mut impl std::io::Read, size: Vec2<isize>) -> anyhow::Result<Self> {
         Ok(match source.get::<u8>()? {
             0 => Block::Literal(Vec::read_const_size(source, size.area() as usize)?),
-            1 => bail!("compressed literal is not supported"),
+            1 => Block::CompressedLiteral(Vec::read(source)?),
             2 => Block::Split(Box::new({
                 let [asize, bsize] = split_size(size);
                 let a = Block::read(source, asize)?;
diff --git a/evc/src/codec/compress.rs b/evc/src/codec/compress.rs
new file mode 100644
index 0000000..e011669
--- /dev/null
+++ b/evc/src/codec/compress.rs
@@ -0,0 +1,151 @@
+use crate::{
+    block::Block,
+    frame::Frame,
+    helpers::{pixel::Pixel, vector::Vec2},
+    view::View,
+};
+use rustdct::DctPlanner;
+
+pub fn compress_block(block: &mut Block, size: Vec2<usize>) {
+    match block {
+        Block::Literal(p) => {
+            let comp = lit_compress(size.x, size.y, &p);
+            *block = Block::CompressedLiteral(comp)
+        }
+        Block::Split(box [a, b]) => {
+            let vert = size.x > size.y;
+            compress_block(
+                a,
+                if vert {
+                    Vec2 {
+                        x: size.x / 2,
+                        y: size.y,
+                    }
+                } else {
+                    Vec2 {
+                        x: size.x,
+                        y: size.y / 2,
+                    }
+                },
+            );
+            compress_block(
+                b,
+                if vert {
+                    Vec2 {
+                        x: size.x - size.x / 2,
+                        y: size.y,
+                    }
+                } else {
+                    Vec2 {
+                        x: size.x,
+                        y: size.y - size.y / 2,
+                    }
+                },
+            );
+        }
+        _ => (),
+    }
+}
+
+pub fn lit_compress(w: usize, h: usize, pixels: &[Pixel]) -> Vec<u8> {
+    let mut out = vec![];
+    for ci in 0..3 {
+        let mut ch = vec![0.; w * h];
+        for y in 0..h {
+            for x in 0..w {
+                let p = pixels[x + y * w];
+                ch[x + y * w] = match ci {
+                    0 => p.r,
+                    1 => p.g,
+                    2 => p.b,
+                    _ => unreachable!(),
+                } as f32
+            }
+        }
+
+        norm_dct_channel(w, h, &mut ch);
+        for i in 0..w * h {
+            out.push(ch[i] as u8)
+        }
+    }
+    out
+}
+pub fn lit_decompress(compressed: &[u8], mut target: View<&mut Frame>) {
+    let (w, h) = (target.size.x as usize, target.size.y as usize);
+    for ci in 0..3 {
+        let mut ch = compressed[ci * w * h..(ci + 1) * w * h]
+            .iter()
+            .map(|v| *v as f32)
+            .collect::<Vec<_>>();
+        norm_idct_channel(w, h, &mut ch);
+        for y in 0..h {
+            for x in 0..w {
+                let p = &mut target[Vec2 {
+                    x: x as isize,
+                    y: y as isize,
+                }];
+                *(match ci {
+                    0 => &mut p.r,
+                    1 => &mut p.g,
+                    2 => &mut p.b,
+                    _ => unreachable!(),
+                }) = ch[x + y * w] as u8
+            }
+        }
+    }
+}
+
+fn norm_dct_channel(w: usize, h: usize, values: &mut [f32]) {
+    let d = DctPlanner::new().plan_dct2(w);
+    let mut temp = vec![0.; d.get_scratch_len()];
+    for c in values.chunks_exact_mut(w) {
+        d.process_dct2_with_scratch(c, &mut temp)
+    }
+
+    let mut values_trans = transpose(w, h, values);
+    let d = DctPlanner::new().plan_dct2(h);
+    let mut temp = vec![0.; d.get_scratch_len()];
+    for c in values_trans.chunks_exact_mut(h) {
+        d.process_dct2_with_scratch(c, &mut temp)
+    }
+
+    let scale = 2. / (w as f32 * h as f32).sqrt();
+    for (i, v) in values_trans.iter().enumerate() {
+        values[i] = *v * scale
+    }
+}
+
+fn norm_idct_channel(w: usize, h: usize, values: &mut [f32]) {
+    let d = DctPlanner::new().plan_dct2(w);
+    let mut temp = vec![0.; d.get_scratch_len()];
+    for c in values.chunks_exact_mut(w) {
+        d.process_dct3_with_scratch(c, &mut temp)
+    }
+
+    let mut values_trans = transpose(w, h, values);
+    let d = DctPlanner::new().plan_dct2(h);
+    let mut temp = vec![0.; d.get_scratch_len()];
+    for c in values_trans.chunks_exact_mut(h) {
+        d.process_dct3_with_scratch(c, &mut temp)
+    }
+
+    let scale = 2. / (w as f32 * h as f32).sqrt();
+    for (i, v) in values_trans.iter().enumerate() {
+        values[i] = *v * scale
+    }
+}
+
+fn transpose(w: usize, h: usize, values: &[f32]) -> Vec<f32> {
+    let mut i = 0;
+    let mut it;
+    let mut transposed = vec![0.; values.len()];
+    for row in 0..h {
+        it = row;
+        for _ in 0..w {
+            transposed[it] = values[i];
+            i += 1;
+            it += h;
+        }
+    }
+    transposed
+}
diff --git a/evc/src/codec/decode.rs b/evc/src/codec/decode.rs
index 61234da..087483c 100644
--- a/evc/src/codec/decode.rs
+++ b/evc/src/codec/decode.rs
@@ -2,6 +2,8 @@ use crate::{
     block::Block, frame::Frame, helpers::threading::both_par, refsampler::Sampler, view::View,
 };
 
+use super::compress::lit_decompress;
+
 pub struct DecodeConfig {
     pub max_threads: usize,
 }
@@ -29,7 +31,9 @@ pub fn decode_block(
                 config.max_threads,
             );
         }
-        Block::CompressedLiteral(_) => todo!(),
+        Block::CompressedLiteral(data) => {
+            lit_decompress(&data, target);
+        }
         Block::Reference { translation } => target.copy_from(&prev.offset(*translation)),
         Block::AdvancedReference(r) => target.copy_from_sampler(&Sampler::from_refblock(prev, r)),
     }
diff --git a/evc/src/codec/encode/mod.rs b/evc/src/codec/encode/mod.rs
index 336f298..76fb481 100644
--- a/evc/src/codec/encode/mod.rs
+++ b/evc/src/codec/encode/mod.rs
@@ -30,7 +30,7 @@ pub enum EncodeMode {
     AdvancedPartial,
 }
 
-pub fn encode_block(view: View<&Frame>, prev: View<&Frame>, config: &EncodeConfig) -> Block {
+pub fn encode_block(view: View<&Frame>, prev: View<&Frame>, config: &EncodeConfig) -> (f64, Block) {
     let (diff, refblock) = if view.area() > config.max_diff_area {
         (
             f64::INFINITY,
@@ -57,10 +57,10 @@ pub fn encode_block(view: View<&Frame>, prev: View<&Frame>, config: &EncodeConfi
         (diff - irrelevance, refblock)
     };
     if diff < config.ref_thres {
-        refblock
+        (diff, refblock)
     } else {
         if view.size.x < config.min_block_size || view.size.y < config.min_block_size {
-            Block::Literal(view.pixels())
+            (0.0, Block::Literal(view.pixels()))
         } else {
             let [av, bv] =
                 unsafe { std::mem::transmute::<_, [View<&'static Frame>; 2]>(view.split()) };
@@ -69,7 +69,7 @@ pub fn encode_block(view: View<&Frame>, prev: View<&Frame>, config: &EncodeConfi
             let config = unsafe { std::mem::transmute::<_, &'static EncodeConfig>(config) };
 
             // only bother to do multithreading, when the block is big.
-            let (a, b) = if view.area() > 100 {
+            let ((ad, a), (bd, b)) = if view.area() > 100 {
                 both_par(
                     || encode_block(av, ap, config),
                     || encode_block(bv, bp, config),
@@ -79,15 +79,18 @@ pub fn encode_block(view: View<&Frame>, prev: View<&Frame>, config: &EncodeConfi
                 (encode_block(av, ap, config), encode_block(bv, bp, config))
             };
 
-            if a.is_literal() && b.is_literal() {
-                Block::Literal(view.pixels())
-            } else if Block::identical_ref(&a, &b) {
-                Block::Reference {
-                    translation: Vec2::<isize>::ZERO,
-                }
-            } else {
-                Block::Split(Box::new([a, b]))
-            }
+            (
+                ad + bd,
+                if a.is_literal() && b.is_literal() {
+                    Block::Literal(view.pixels())
+                } else if Block::identical_ref(&a, &b) {
+                    Block::Reference {
+                        translation: Vec2::<isize>::ZERO,
+                    }
+                } else {
+                    Block::Split(Box::new([a, b]))
+                },
+            )
         }
     }
 }
diff --git a/evc/src/codec/mod.rs b/evc/src/codec/mod.rs
index 1c80d24..3203e6e 100644
--- a/evc/src/codec/mod.rs
+++ b/evc/src/codec/mod.rs
@@ -1,2 +1,3 @@
 pub mod decode;
 pub mod encode;
+pub mod compress;
diff --git a/evc/src/view.rs b/evc/src/view.rs
index f153597..6f34965 100644
--- a/evc/src/view.rs
+++ b/evc/src/view.rs
@@ -187,7 +187,7 @@ impl View<&mut Frame> {
     pub fn set_pixels(&mut self, pixels: &Vec<Pixel>) {
         for y in 0..self.size.y {
             for x in 0..self.size.x {
-                self[(x, y)] = pixels[(x * self.size.y + y) as usize]
+                self[(x, y)] = pixels[(y * self.size.x + x) as usize]
             }
         }
     }